Parcourir la source

feat(stats): query r2 data catalog

Adam il y a 4 jours
Parent
commit
46a14e685a

+ 13 - 1
infra/stats.ts

@@ -181,6 +181,16 @@ const statsSyncConfig = new sst.Linkable("StatsSyncConfig", {
   },
 })
 
+const r2SqlAuthToken = new sst.Secret("R2SqlAuthToken")
+const r2Sql = new sst.Linkable("R2Sql", {
+  properties: {
+    accountId: "15d29c8639fd3733b1b5486a2acfd968",
+    bucket: `platform-${$app.stage}-lake`,
+    namespace: "inference",
+    table: "generation",
+  },
+})
+
 export const statSync = new sst.aws.Service("StatsSyncService", {
   cluster: lakeCluster,
   architecture: "arm64",
@@ -193,7 +203,9 @@ export const statSync = new sst.aws.Service("StatsSyncService", {
     dockerfile: "packages/stats/server/Dockerfile",
   },
   command: ["bun", "src/stat-sync.ts"],
-  link: [database, inferenceEvent, statsSyncConfig],
+  // Keep the legacy Athena link and IAM permissions during the first R2-backed
+  // release so reverting the application code remains a one-deploy rollback.
+  link: [database, inferenceEvent, r2Sql, r2SqlAuthToken, statsSyncConfig],
   permissions: lakeQueryPermissions,
   scaling: {
     min: 1,

+ 1 - 0
packages/stats/core/package.json

@@ -12,6 +12,7 @@
     "./database": "./src/database.ts",
     "./database/*": "./src/database/*.ts",
     "./domain/*": "./src/domain/*.ts",
+    "./r2-sql": "./src/r2-sql.ts",
     "./runtime": "./src/runtime.ts",
     "./stat-sync": "./src/stat-sync.ts"
   },

+ 22 - 1
packages/stats/core/src/domain/inference.test.ts

@@ -1,5 +1,5 @@
 import { describe, expect, test } from "bun:test"
-import { toGeoAggregate, toModelAggregate, toProviderAggregate } from "./inference"
+import { buildStatsQueries, toGeoAggregate, toModelAggregate, toProviderAggregate } from "./inference"
 import { modelAuthor, normalizeInferenceModel, statModel, statProvider } from "./model-normalization"
 
 describe("inference stat normalization", () => {
@@ -82,6 +82,27 @@ describe("inference stat normalization", () => {
       }),
     ).toMatchObject([{ period_key: "2026-W20" }])
   })
+
+  test("builds bounded R2 SQL queries for each day and week", () => {
+    const queries = buildStatsQueries(new Date("2026-08-10T00:00:00.000Z"), new Date("2026-08-12T12:00:00.000Z"), {
+      namespace: "inference",
+      table: "generation",
+      dataset: "zen",
+    })
+
+    expect(queries).toHaveLength(8)
+    expect(queries[0]).toContain("'week' AS grain")
+    expect(queries[0]).toContain("'2026-W33' AS period_key")
+    expect(queries[2]).toContain("'2026-08-10' AS period_key")
+    expect(queries[6]).toContain("'2026-08-12' AS period_key")
+    expect(queries[0]).toContain('FROM "inference"."generation"')
+    expect(queries[0]).toContain("event_type = 'generation.completed'")
+    expect(queries[0]).toContain("product = 'go'")
+    expect(queries[0]).toContain("LIMIT 10000")
+    expect(queries[0]).toContain("approx_distinct(session) AS sessions")
+    expect(queries[1]).toContain("'geo_model' ELSE 'geo'")
+    expect(queries[1]).toContain("0 AS sessions")
+  })
 })
 
 function aggregate(model: string, provider: string) {

+ 139 - 104
packages/stats/core/src/domain/inference.ts

@@ -1,5 +1,5 @@
 import { Resource } from "sst/resource"
-import type { AthenaData } from "../athena"
+import type { R2SqlData } from "../r2-sql"
 import type { GeoStatAggregate } from "./geo"
 import type { ModelStatAggregate } from "./model"
 import {
@@ -13,22 +13,66 @@ import type { ProviderStatAggregate } from "./provider"
 import { normalizeCountry, normalizeTier, type StatBaseAggregate } from "./stat"
 
 export type StatDimension = "model" | "provider" | "geo" | "geo_model"
+export type StatsQuerySource = { namespace: string; table: string; dataset: string }
+type StatsQueryFamily = "usage" | "geo"
 
-// All stat dimensions and both grains are computed in one query via GROUPING SETS so
-// the source table is scanned once per sync pass; separate queries per dimension (and
-// the previous weekly/daily UNION ALL) each re-scanned the same events.
-export function buildStatsQuery(periodStart: Date, periodEnd: Date) {
-  const periodStartValue = sqlString(periodStart.toISOString())
-  const periodEndValue = sqlString(periodEnd.toISOString())
-  const periodStartDateValue = sqlString(periodStart.toISOString().slice(0, 10))
-  const periodEndDateValue = sqlString(periodEnd.toISOString().slice(0, 10))
-  const sourceTable = [Resource.InferenceEvent.catalog, Resource.InferenceEvent.database, Resource.InferenceEvent.table]
-    .map(sqlIdentifier)
-    .join(".")
+const DAY_MS = 86_400_000
+const WEEK_MS = 7 * DAY_MS
+
+// R2 SQL limits result sets to 10,000 rows and does not support OFFSET. Two
+// queries per day/week keep each result bounded and avoid combining the costly
+// distinct user/session aggregates with the high-cardinality geo dimensions.
+export function buildStatsQueries(periodStart: Date, periodEnd: Date, input?: StatsQuerySource) {
+  const source = input ?? {
+    namespace: Resource.R2Sql.namespace,
+    table: Resource.R2Sql.table,
+    dataset: Resource.StatsSyncConfig.dataset,
+  }
+  return [...statPeriods("week", periodStart, periodEnd), ...statPeriods("day", periodStart, periodEnd)].flatMap(
+    (period) => [buildStatsQuery(period, source, "usage"), buildStatsQuery(period, source, "geo")],
+  )
+}
+
+function buildStatsQuery(
+  period: { grain: "day" | "week"; key: string; start: Date; end: Date },
+  source: StatsQuerySource,
+  family: StatsQueryFamily,
+) {
+  const periodStartValue = sqlString(period.start.toISOString())
+  const periodEndValue = sqlString(period.end.toISOString())
+  const ingestEndValue = sqlString(new Date(period.end.getTime() + DAY_MS).toISOString())
+  const sourceTable = [source.namespace, source.table].map(sqlIdentifier).join(".")
+  const dimensions =
+    family === "usage"
+      ? `CASE WHEN grouping(model) = 0 THEN 'model' ELSE 'provider' END AS dimension,
+  tier,
+  provider,
+  CASE WHEN grouping(model) = 0 THEN model END AS model,
+  CASE WHEN grouping(model) = 0 THEN COALESCE(MAX(NULLIF(provider_model, '')), '') END AS provider_model,
+  null AS country,
+  null AS continent`
+      : `CASE WHEN grouping(model) = 0 THEN 'geo_model' ELSE 'geo' END AS dimension,
+  tier,
+  CASE WHEN grouping(model) = 0 THEN provider ELSE 'all' END AS provider,
+  CASE WHEN grouping(model) = 0 THEN model ELSE 'all' END AS model,
+  null AS provider_model,
+  country,
+  COALESCE(MAX(NULLIF(continent, '')), '') AS continent`
+  const distinctColumns =
+    family === "usage"
+      ? `approx_distinct(session) AS sessions,
+    approx_distinct(user_key) AS unique_users`
+      : `0 AS sessions,
+    0 AS unique_users`
+  const groupingSets =
+    family === "usage"
+      ? `(tier, provider, model),
+  (tier, provider)`
+      : `(tier, country),
+  (tier, provider, model, country)`
   const aggregateColumns = `
-    COUNT(DISTINCT session) AS sessions,
+    ${distinctColumns},
     COUNT(*) AS requests,
-    COUNT(DISTINCT user_key) AS unique_users,
     COALESCE(SUM(tokens_input), 0) AS input_tokens,
     COALESCE(SUM(tokens_output), 0) AS output_tokens,
     COALESCE(SUM(tokens_reasoning), 0) AS reasoning_tokens,
@@ -38,65 +82,57 @@ export function buildStatsQuery(periodStart: Date, periodEnd: Date) {
     COALESCE(SUM(cost_output_microcents), 0) AS output_cost_microcents,
     COALESCE(SUM(cost_total_microcents), 0) AS total_cost_microcents,
     AVG(duration_ms) AS avg_duration_ms,
-    approx_percentile(CAST(duration_ms AS double), 0.5) AS p50_duration_ms,
-    approx_percentile(CAST(duration_ms AS double), 0.95) AS p95_duration_ms,
+    null AS p50_duration_ms,
+    null AS p95_duration_ms,
     AVG(ttfb_ms) AS avg_ttfb_ms,
-    approx_percentile(CAST(ttfb_ms AS double), 0.5) AS p50_ttfb_ms,
-    approx_percentile(CAST(ttfb_ms AS double), 0.95) AS p95_ttfb_ms,
+    null AS p50_ttfb_ms,
+    null AS p95_ttfb_ms,
     AVG(output_tps) AS avg_output_tps,
-    SUM(CASE WHEN status >= 200 AND status < 400 THEN 1 ELSE 0 END) AS success_count,
-    SUM(CASE WHEN status >= 400 THEN 1 ELSE 0 END) AS error_count,
+    SUM(CASE WHEN outcome = 'succeeded' THEN 1 ELSE 0 END) AS success_count,
+    SUM(CASE WHEN outcome = 'failed' THEN 1 ELSE 0 END) AS error_count,
     COUNT(*) AS sample_count`
 
   return `
 WITH normalized AS (
   SELECT
-    from_iso8601_timestamp(event_timestamp) AS event_time,
-    model AS raw_model,
-    ${statModelSql("model", "provider_model")} AS model,
-    COALESCE(NULLIF(provider_model, ''), '') AS provider_model,
-    COALESCE(NULLIF(provider, ''), '') AS raw_provider,
-    UPPER(COALESCE(NULLIF(cf_country, ''), 'ZZ')) AS country,
-    COALESCE(NULLIF(cf_continent, ''), '') AS continent,
-    session,
-    COALESCE(NULLIF(workspace, ''), '') AS workspace,
-    COALESCE(NULLIF(api_key, ''), '') AS api_key,
+    model_requested AS raw_model,
+    ${statModelSql("model_requested", "route_model")} AS model,
+    COALESCE(NULLIF(route_model, ''), '') AS provider_model,
+    COALESCE(NULLIF(provider_id, ''), '') AS raw_provider,
+    UPPER(COALESCE(NULLIF(country, ''), 'ZZ')) AS country,
+    COALESCE(NULLIF(continent, ''), '') AS continent,
+    session_id AS session,
+    COALESCE(NULLIF(workspace_id, ''), '') AS workspace,
+    COALESCE(NULLIF(service_api_key_id, ''), '') AS api_key,
     COALESCE(NULLIF(user_id, ''), '') AS user_id,
-    status,
-    duration AS duration_ms,
-    time_to_first_byte AS ttfb_ms,
-    timestamp_first_byte,
-    timestamp_last_byte,
+    outcome,
+    duration_ms,
+    time_to_first_token_ms AS ttfb_ms,
+    CASE
+      WHEN first_token_at IS NULL OR last_token_at IS NULL THEN null
+      ELSE date_part('epoch', last_token_at) - date_part('epoch', first_token_at)
+    END AS output_seconds,
     tokens_input,
     tokens_output,
     tokens_reasoning,
     tokens_cache_read,
-    tokens_cache_write_5m,
-    tokens_cache_write_1h,
-    cost_input_microcents,
-    cost_output_microcents,
-    cost_total_microcents,
-    cost_input,
-    cost_output,
-    cost_total,
-    source
+    tokens_cache_write,
+    cost_input AS cost_input_microcents,
+    cost_output AS cost_output_microcents,
+    cost_total AS cost_total_microcents
   FROM ${sourceTable}
-  WHERE event_type = 'completions'
-    AND model IS NOT NULL
-    AND model <> ''
-    AND source = 'lite'
-    AND event_date >= ${periodStartDateValue}
-    AND event_date <= ${periodEndDateValue}
-    AND event_timestamp >= ${periodStartValue}
-    AND event_timestamp < ${periodEndValue}
+  WHERE event_type = 'generation.completed'
+    AND source IN ('inference', 'inference-legacy')
+    AND product = 'go'
+    AND model_requested IS NOT NULL
+    AND model_requested <> ''
+    AND __ingest_ts >= ${periodStartValue}
+    AND __ingest_ts < ${ingestEndValue}
+    AND started_at >= ${periodStartValue}
+    AND started_at < ${periodEndValue}
 ), filtered AS (
   SELECT
-    event_time,
-    CASE
-      WHEN source = 'lite' THEN 'Go'
-      WHEN raw_model IN ('gpt-5-nano', 'grok-code', 'big-pickle') OR regexp_like(raw_model, '-free(:global)?$') THEN 'Free'
-      ELSE 'Paid'
-    END AS tier,
+    'Go' AS tier,
     ${statProviderSql("model", "provider_model", "raw_provider")} AS provider,
     provider_model,
     model,
@@ -104,63 +140,39 @@ WITH normalized AS (
     continent,
     session,
     COALESCE(NULLIF(user_id, ''), NULLIF(workspace, ''), NULLIF(api_key, '')) AS user_key,
-    status,
+    outcome,
     duration_ms,
     ttfb_ms,
     CASE
-      WHEN timestamp_last_byte - timestamp_first_byte < 100 THEN null
-      ELSE CAST(tokens_output AS double) / (timestamp_last_byte - timestamp_first_byte) * 1000
+      WHEN output_seconds < 0.1 THEN null
+      ELSE CAST(tokens_output AS double) / output_seconds
     END AS output_tps,
     tokens_input,
     tokens_output,
     tokens_reasoning,
     tokens_cache_read,
-    COALESCE(tokens_cache_read, 0) + COALESCE(tokens_cache_write_5m, 0) + COALESCE(tokens_cache_write_1h, 0) + COALESCE(tokens_input, 0) + COALESCE(tokens_output, 0) AS tokens_total,
-    COALESCE(cost_input_microcents, cost_input * 1000000) AS cost_input_microcents,
-    COALESCE(cost_output_microcents, cost_output * 1000000) AS cost_output_microcents,
-    COALESCE(cost_total_microcents, cost_total * 1000000) AS cost_total_microcents
+    COALESCE(tokens_cache_read, 0) + COALESCE(tokens_cache_write, 0) + COALESCE(tokens_input, 0) + COALESCE(tokens_output, 0) AS tokens_total,
+    cost_input_microcents,
+    cost_output_microcents,
+    cost_total_microcents
   FROM normalized
   WHERE lower(model) NOT IN (${[...EXCLUDED_MODELS].map(sqlString).join(", ")})
-), periods AS (
-  SELECT
-    concat(CAST(year_of_week(event_time) AS varchar), '-W', lpad(CAST(week(event_time) AS varchar), 2, '0')) AS week_key,
-    substr(to_iso8601(date_trunc('day', event_time)), 1, 10) AS day_key,
-    *
-  FROM filtered
 )
 SELECT
-  CASE WHEN grouping(week_key) = 0 THEN 'week' ELSE 'day' END AS grain,
-  COALESCE(week_key, day_key) AS period_key,
-  ${sqlString(Resource.StatsSyncConfig.dataset)} AS dataset,
-  CASE
-    WHEN grouping(country) = 0 AND grouping(model) = 0 THEN 'geo_model'
-    WHEN grouping(country) = 0 THEN 'geo'
-    WHEN grouping(model) = 0 THEN 'model'
-    ELSE 'provider'
-  END AS dimension,
-  tier,
-  CASE WHEN grouping(provider) = 0 THEN provider ELSE 'all' END AS provider,
-  CASE WHEN grouping(model) = 0 THEN model WHEN grouping(country) = 0 THEN 'all' END AS model,
-  CASE WHEN grouping(model) = 0 AND grouping(country) = 1 THEN COALESCE(MAX(NULLIF(provider_model, '')), '') END AS provider_model,
-  CASE WHEN grouping(country) = 0 THEN country END AS country,
-  CASE WHEN grouping(country) = 0 THEN COALESCE(MAX(NULLIF(continent, '')), '') END AS continent,
+  ${sqlString(period.grain)} AS grain,
+  ${sqlString(period.key)} AS period_key,
+  ${sqlString(source.dataset)} AS dataset,
+  ${dimensions},
   ${aggregateColumns}
-FROM periods
+FROM filtered
 GROUP BY GROUPING SETS (
-  (week_key, tier, provider, model),
-  (week_key, tier, provider),
-  (week_key, tier, country),
-  (week_key, tier, provider, model, country),
-  (day_key, tier, provider, model),
-  (day_key, tier, provider),
-  (day_key, tier, country),
-  (day_key, tier, provider, model, country)
+  ${groupingSets}
 )
-ORDER BY grain, period_key, total_tokens DESC
+LIMIT 10000
 `
 }
 
-export function toModelAggregate(data: AthenaData): ModelStatAggregate[] {
+export function toModelAggregate(data: R2SqlData): ModelStatAggregate[] {
   const model = statModel(data.model, data.provider_model)
   const provider = statProvider(model, data.provider_model, data.provider)
   if (!provider) return []
@@ -170,13 +182,13 @@ export function toModelAggregate(data: AthenaData): ModelStatAggregate[] {
   ])
 }
 
-export function toProviderAggregate(data: AthenaData): ProviderStatAggregate[] {
+export function toProviderAggregate(data: R2SqlData): ProviderStatAggregate[] {
   return toStatBaseAggregate(data).flatMap((base) => [
     { ...base, provider: statProvider(data.model, data.provider_model, data.provider) || "unknown" },
   ])
 }
 
-export function toGeoAggregate(data: AthenaData): GeoStatAggregate[] {
+export function toGeoAggregate(data: R2SqlData): GeoStatAggregate[] {
   return toStatBaseAggregate(data).flatMap((base) => [
     {
       ...base,
@@ -188,7 +200,7 @@ export function toGeoAggregate(data: AthenaData): GeoStatAggregate[] {
   ])
 }
 
-function toStatBaseAggregate(data: AthenaData): StatBaseAggregate[] {
+function toStatBaseAggregate(data: R2SqlData): StatBaseAggregate[] {
   const grain = data.grain === "day" || data.grain === "week" ? data.grain : undefined
   if (!grain || !data.period_key) return []
 
@@ -223,21 +235,21 @@ function toStatBaseAggregate(data: AthenaData): StatBaseAggregate[] {
   ]
 }
 
-function integer(data: AthenaData, key: string) {
+function integer(data: R2SqlData, key: string) {
   return Math.round(number(data, key))
 }
 
-function nullableNumber(data: AthenaData, key: string) {
+function nullableNumber(data: R2SqlData, key: string) {
   if (data[key] === undefined || data[key] === "") return null
   return Number(number(data, key).toFixed(2))
 }
 
-function nullableInteger(data: AthenaData, key: string) {
+function nullableInteger(data: R2SqlData, key: string) {
   if (data[key] === undefined || data[key] === "") return null
   return Math.round(number(data, key))
 }
 
-function number(data: AthenaData, key: string) {
+function number(data: R2SqlData, key: string) {
   const value = Number(data[key])
   return Number.isFinite(value) ? value : 0
 }
@@ -250,6 +262,29 @@ function sqlString(value: string) {
   return `'${value.replace(/'/g, "''")}'`
 }
 
+function statPeriods(grain: "day" | "week", periodStart: Date, periodEnd: Date) {
+  const interval = grain === "day" ? DAY_MS : WEEK_MS
+  const count = Math.max(0, Math.ceil((periodEnd.getTime() - periodStart.getTime()) / interval))
+  return Array.from({ length: count }, (_, index) => {
+    const start = new Date(periodStart.getTime() + index * interval)
+    return {
+      grain,
+      key: grain === "day" ? start.toISOString().slice(0, 10) : isoWeekKey(start),
+      start,
+      end: new Date(Math.min(start.getTime() + interval, periodEnd.getTime())),
+    }
+  })
+}
+
+function isoWeekKey(date: Date) {
+  const thursday = new Date(Date.UTC(date.getUTCFullYear(), date.getUTCMonth(), date.getUTCDate()))
+  const day = thursday.getUTCDay() || 7
+  thursday.setUTCDate(thursday.getUTCDate() + 4 - day)
+  const year = thursday.getUTCFullYear()
+  const week = Math.ceil((thursday.getTime() - Date.UTC(year, 0, 1) + DAY_MS) / WEEK_MS)
+  return `${year}-W${String(week).padStart(2, "0")}`
+}
+
 function statModelSql(model: string, providerModel: string) {
   return `COALESCE(NULLIF(regexp_replace(CASE
       WHEN lower(${model}) = 'big-pickle' THEN NULLIF(${providerModel}, '')

+ 105 - 0
packages/stats/core/src/r2-sql.ts

@@ -0,0 +1,105 @@
+import { Context, Effect, Layer, Schema } from "effect"
+import { Resource } from "sst/resource"
+
+const R2_SQL_MAX_ROWS = 10_000
+const R2SqlValue = Schema.Union([Schema.String, Schema.Number, Schema.Boolean, Schema.Null])
+const R2SqlResponse = Schema.Struct({
+  success: Schema.Boolean,
+  result: Schema.optional(
+    Schema.NullOr(
+      Schema.Struct({
+        request_id: Schema.String,
+        rows: Schema.Array(Schema.Record(Schema.String, R2SqlValue)),
+      }),
+    ),
+  ),
+  errors: Schema.Array(Schema.Unknown),
+})
+const decodeResponse = Schema.decodeUnknownEffect(Schema.fromJsonString(R2SqlResponse))
+
+export type R2SqlData = Record<string, string>
+
+export class R2SqlQueryError extends Error {
+  readonly _tag = "R2SqlQueryError"
+  readonly requestId?: string
+  readonly status?: number
+
+  constructor(input: { message: string; requestId?: string; status?: number; cause?: unknown }) {
+    super(input.message, { cause: input.cause })
+    this.name = "R2SqlQueryError"
+    this.requestId = input.requestId
+    this.status = input.status
+  }
+}
+
+export declare namespace R2Sql {
+  export interface Service {
+    readonly query: (query: string) => Effect.Effect<R2SqlData[], R2SqlQueryError>
+  }
+}
+
+export class R2Sql extends Context.Service<R2Sql, R2Sql.Service>()("@opencode/stats/R2Sql") {
+  static readonly layer: Layer.Layer<R2Sql> = Layer.succeed(
+    R2Sql,
+    R2Sql.of({
+      query: Effect.fn("R2Sql.query")(function* (query: string) {
+        const response = yield* Effect.tryPromise({
+          try: () =>
+            Bun.fetch(
+              `https://api.sql.cloudflarestorage.com/api/v1/accounts/${Resource.R2Sql.accountId}/r2-sql/query/${Resource.R2Sql.bucket}`,
+              {
+                method: "POST",
+                headers: {
+                  Authorization: `Bearer ${Resource.R2SqlAuthToken.value}`,
+                  "Content-Type": "application/json",
+                },
+                body: JSON.stringify({ query }),
+              },
+            ),
+          catch: (cause) => new R2SqlQueryError({ message: "Failed to run R2 SQL stats query", cause }),
+        })
+        const body = yield* Effect.tryPromise({
+          try: () => response.text(),
+          catch: (cause) =>
+            new R2SqlQueryError({ message: "Failed to read R2 SQL stats response", status: response.status, cause }),
+        })
+        const decoded = yield* decodeResponse(body).pipe(
+          Effect.mapError(
+            (cause) =>
+              new R2SqlQueryError({
+                message: "R2 SQL returned an invalid stats response",
+                status: response.status,
+                cause,
+              }),
+          ),
+        )
+        if (!response.ok || !decoded.success || !decoded.result)
+          return yield* Effect.fail(
+            new R2SqlQueryError({
+              message: `R2 SQL stats query failed: ${JSON.stringify(decoded.errors)}`,
+              requestId: decoded.result?.request_id,
+              status: response.status,
+            }),
+          )
+
+        // R2 SQL has no OFFSET support and caps LIMIT at 10,000. Each stats
+        // query is scoped to one day or week, and reaching the cap is treated as
+        // an error so a newly high-cardinality period can never be truncated.
+        if (decoded.result.rows.length >= R2_SQL_MAX_ROWS)
+          return yield* Effect.fail(
+            new R2SqlQueryError({
+              message: `R2 SQL stats query reached the ${R2_SQL_MAX_ROWS} row limit`,
+              requestId: decoded.result.request_id,
+              status: response.status,
+            }),
+          )
+
+        return decoded.result.rows.map((row) =>
+          Object.fromEntries(
+            Object.entries(row).flatMap(([key, value]) => (value === null ? [] : [[key, String(value)]])),
+          ),
+        )
+      }),
+    }),
+  )
+}

+ 11 - 0
packages/stats/core/src/resource.d.ts

@@ -11,6 +11,17 @@ declare module "sst/resource" {
       type: "sst.sst.Linkable"
       workgroup: string
     }
+    R2Sql: {
+      accountId: string
+      bucket: string
+      namespace: string
+      table: string
+      type: "sst.sst.Linkable"
+    }
+    R2SqlAuthToken: {
+      type: "sst.sst.Secret"
+      value: string
+    }
     StatsSyncConfig: {
       dataset: string
       type: "sst.sst.Linkable"

+ 14 - 13
packages/stats/core/src/stat-sync.ts

@@ -1,12 +1,12 @@
 import { DateTime, Effect } from "effect"
 import { Resource } from "sst/resource"
-import { Athena, AthenaQueryError, AthenaQueryTimeoutError } from "./athena"
 import { DatabaseError } from "./database"
 import { GeoStatRepo, rowsFromAggregates as geoRowsFromAggregates } from "./domain/geo"
-import { buildStatsQuery, toGeoAggregate, toModelAggregate, toProviderAggregate } from "./domain/inference"
+import { buildStatsQueries, toGeoAggregate, toModelAggregate, toProviderAggregate } from "./domain/inference"
 import { ModelStatRepo, rowsFromAggregates as modelRowsFromAggregates } from "./domain/model"
 import { ProviderStatRepo, rowsFromAggregates as providerRowsFromAggregates } from "./domain/provider"
 import { startOfIsoWeek } from "./domain/stat"
+import { R2Sql, R2SqlQueryError } from "./r2-sql"
 
 const DATALAKE_INGESTION_LAG_MS = 5 * 60_000
 const STATS_DATA_START_MS = new Date("2026-05-28T00:00:00.000Z").getTime()
@@ -18,23 +18,25 @@ const DISPLAY_WINDOW_MS = 56 * 86_400_000
 const INCREMENTAL_LOOKBACK_MS = 2 * 3_600_000
 
 export type SyncStatsResult = { ok: true; rows: number; startedAt: string; periodStart: string; periodEnd: string }
-export type SyncStatsError = AthenaQueryError | AthenaQueryTimeoutError | DatabaseError
+export type SyncStatsError = R2SqlQueryError | DatabaseError
 
 export const syncStats: (options?: {
   full?: boolean
-}) => Effect.Effect<SyncStatsResult, SyncStatsError, Athena | ModelStatRepo | ProviderStatRepo | GeoStatRepo> =
+}) => Effect.Effect<SyncStatsResult, SyncStatsError, R2Sql | ModelStatRepo | ProviderStatRepo | GeoStatRepo> =
   Effect.fn("StatSync.sync")(function* (options?: { full?: boolean }) {
     const startedAt = yield* DateTime.nowAsDate
     const periodEnd = new Date(Math.floor((startedAt.getTime() - DATALAKE_INGESTION_LAG_MS) / 60_000) * 60_000)
     const periodStart = options?.full ? fullPeriodStart(periodEnd) : incrementalPeriodStart(periodEnd)
-    const athena = yield* Athena
+    const r2Sql = yield* R2Sql
     const modelStats = yield* ModelStatRepo
     const providerStats = yield* ProviderStatRepo
     const geoStats = yield* GeoStatRepo
 
     yield* logRuntimeCheck()
 
-    const rows = yield* athena.query(buildStatsQuery(periodStart, periodEnd))
+    const rows = yield* Effect.forEach(buildStatsQueries(periodStart, periodEnd), r2Sql.query, {
+      concurrency: 4,
+    }).pipe(Effect.map((batches) => batches.flat()))
     const modelRows = modelRowsFromAggregates(rows.filter((row) => row.dimension === "model").flatMap(toModelAggregate))
     const providerRows = providerRowsFromAggregates(
       rows.filter((row) => row.dimension === "provider").flatMap(toProviderAggregate),
@@ -77,7 +79,7 @@ export const syncStats: (options?: {
     }
   })
 
-// May 27 was partial, so keep Athena stats anchored at the first complete day.
+// May 27 was partial, so keep stats anchored at the first complete day.
 function fullPeriodStart(periodEnd: Date) {
   return new Date(
     Math.max(
@@ -99,13 +101,12 @@ function incrementalPeriodStart(periodEnd: Date) {
 
 function logRuntimeCheck() {
   return Effect.logInfo(
-    `athena stats runtime check ${JSON.stringify({
-      catalog: Resource.InferenceEvent.catalog,
-      database: Resource.InferenceEvent.database,
+    `r2 sql stats runtime check ${JSON.stringify({
+      accountId: Resource.R2Sql.accountId,
+      bucket: Resource.R2Sql.bucket,
       dataset: Resource.StatsSyncConfig.dataset,
-      table: Resource.InferenceEvent.table,
-      workgroup: Resource.InferenceEvent.workgroup,
-      region: Resource.InferenceEvent.region,
+      namespace: Resource.R2Sql.namespace,
+      table: Resource.R2Sql.table,
       stage: Resource.App.stage,
     })}`,
   )

+ 5 - 5
packages/stats/server/src/stat-sync.ts

@@ -1,6 +1,6 @@
 import * as NodeRuntime from "@effect/platform-node/NodeRuntime"
-import { Athena } from "@opencode-ai/stats-core/athena"
 import { ModelStatRepo } from "@opencode-ai/stats-core/domain/model"
+import { R2Sql } from "@opencode-ai/stats-core/r2-sql"
 import { layer as statsLayer } from "@opencode-ai/stats-core/runtime"
 import { syncStats } from "@opencode-ai/stats-core/stat-sync"
 import { Cause, Duration, Effect, Layer, Schedule } from "effect"
@@ -8,7 +8,7 @@ import { Cause, Duration, Effect, Layer, Schedule } from "effect"
 const SYNC_INTERVAL = "1 hour"
 const SYNC_INTERVAL_MS = 3_600_000
 
-const runtimeLayer = Layer.mergeAll(statsLayer, Athena.layer)
+const runtimeLayer = Layer.mergeAll(statsLayer, R2Sql.layer)
 
 const daemon = Effect.gen(function* () {
   yield* Effect.logInfo("stats sync daemon started")
@@ -40,9 +40,9 @@ const daemon = Effect.gen(function* () {
   yield* pass.pipe(Effect.repeat(Schedule.fixed(SYNC_INTERVAL)))
 }).pipe(Effect.forkScoped)
 
-// A restarted daemon must not immediately re-run the expensive Athena pass; resume
-// the hourly cadence from the last completed sync instead. This caps the Athena
-// spend of a crash loop at one pass per interval.
+// A restarted daemon must not immediately re-run the R2 SQL pass; resume the
+// hourly cadence from the last completed sync instead. This caps the query spend
+// of a crash loop at one pass per interval.
 const initialDelay = Effect.fnUntraced(function* () {
   const modelStats = yield* ModelStatRepo
   const lastSynced = yield* modelStats.lastSyncedAt().pipe(Effect.catchCause(() => Effect.succeed(null)))