diff --git a/infra/stats.ts b/infra/stats.ts index 10d37119f0d..29c9537daf8 100644 --- a/infra/stats.ts +++ b/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, diff --git a/packages/stats/core/package.json b/packages/stats/core/package.json index ffedc71d429..92e8ab0e262 100644 --- a/packages/stats/core/package.json +++ b/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" }, diff --git a/packages/stats/core/src/domain/inference.test.ts b/packages/stats/core/src/domain/inference.test.ts index fa71fb51e60..f58e7deab6a 100644 --- a/packages/stats/core/src/domain/inference.test.ts +++ b/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) { diff --git a/packages/stats/core/src/domain/inference.ts b/packages/stats/core/src/domain/inference.ts index 558832f9953..ad246053054 100644 --- a/packages/stats/core/src/domain/inference.ts +++ b/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}, '') diff --git a/packages/stats/core/src/r2-sql.ts b/packages/stats/core/src/r2-sql.ts new file mode 100644 index 00000000000..91093643c0e --- /dev/null +++ b/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 + +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 + } +} + +export class R2Sql extends Context.Service()("@opencode/stats/R2Sql") { + static readonly layer: Layer.Layer = 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)]])), + ), + ) + }), + }), + ) +} diff --git a/packages/stats/core/src/resource.d.ts b/packages/stats/core/src/resource.d.ts index 8343f7baa63..b8017777971 100644 --- a/packages/stats/core/src/resource.d.ts +++ b/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" diff --git a/packages/stats/core/src/stat-sync.ts b/packages/stats/core/src/stat-sync.ts index 736ca852f3f..ceec6f7e6dc 100644 --- a/packages/stats/core/src/stat-sync.ts +++ b/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 = +}) => Effect.Effect = 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, })}`, ) diff --git a/packages/stats/server/src/stat-sync.ts b/packages/stats/server/src/stat-sync.ts index 613fbec5b7d..79766096326 100644 --- a/packages/stats/server/src/stat-sync.ts +++ b/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)))