diff --git a/packages/stats/core/src/domain/inference.test.ts b/packages/stats/core/src/domain/inference.test.ts index 8f7f4374c2e..7563f18615f 100644 --- a/packages/stats/core/src/domain/inference.test.ts +++ b/packages/stats/core/src/domain/inference.test.ts @@ -221,7 +221,7 @@ describe("inference stat normalization", () => { expect(queries[0]).toContain("OR lower(raw_model) IN ('gpt-5-nano', 'grok-code', 'big-pickle')") expect(queries[0]).toContain("OR lower(raw_model) LIKE '%-free'") expect(queries[0]).toContain("THEN 'Free'") - expect(queries[0]).toContain("LIMIT 10000") + expect(queries[0]).not.toContain("LIMIT") expect(queries[0]).toContain("approx_distinct(session) AS sessions") expect(queries[1]).toContain("'geo_model' ELSE 'geo'") expect(queries[1]).toContain("0 AS sessions") diff --git a/packages/stats/core/src/domain/inference.ts b/packages/stats/core/src/domain/inference.ts index 963516625dd..fb45188aadf 100644 --- a/packages/stats/core/src/domain/inference.ts +++ b/packages/stats/core/src/domain/inference.ts @@ -36,9 +36,9 @@ const WEEK_MS = 7 * DAY_MS // from both sources. const LIVE_SOURCE_START = "2026-08-11T10:57:48.186Z" -// 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. +// R2 SQL results are cursor-paginated after aggregation. Separate usage and geo +// queries per day/week avoid combining 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, @@ -327,7 +327,6 @@ FROM filtered GROUP BY GROUPING SETS ( ${groupingSets} ) -LIMIT 10000 ` } diff --git a/packages/stats/core/src/r2-sql.test.ts b/packages/stats/core/src/r2-sql.test.ts index b02f429b162..bb19ca393cd 100644 --- a/packages/stats/core/src/r2-sql.test.ts +++ b/packages/stats/core/src/r2-sql.test.ts @@ -1,5 +1,7 @@ +import { Database } from "bun:sqlite" +import { Effect } from "effect" import { expect, test } from "bun:test" -import { R2SqlQueryError } from "./r2-sql" +import { queryR2SqlPages, R2SqlQueryError, type R2SqlData } from "./r2-sql" test("includes the transport failure in the message used by sync logging", () => { const cause = new DOMException("The operation timed out.", "TimeoutError") @@ -16,3 +18,62 @@ test("preserves an R2 response error and request details", () => { expect(error.status).toBe(400) expect(error.requestId).toBe("test-request") }) + +// Execute the generated pagination SQL against real aggregate rows, including +// groups that share country/model prefixes across a page boundary. +for (const count of [0, 9999, 10000, 10001, 20005]) { + test(`reads all ${count} aggregate rows without losing or repeating groups`, async () => { + using db = new Database(":memory:") + db.run("CREATE TABLE aggregates (dimension TEXT, model TEXT, country TEXT, requests TEXT)") + db.transaction(() => { + const insert = db.prepare("INSERT INTO aggregates VALUES (?, ?, ?, ?)") + Array.from({ length: count }, (_, index) => { + insert.run( + index < 2 ? "geo" : "geo_model", + index < 2 ? null : `model'${Math.floor(index / 3)}`, + `country${index % 3}`, + String(index + 1), + ) + }) + })() + const pages: number[] = [] + const rows = await Effect.runPromise( + queryR2SqlPages("SELECT * FROM aggregates", ["dimension", "model", "country"], (query) => + Effect.sync(() => { + const rows = db.query(query).all() as R2SqlData[] + pages.push(rows.length) + return rows + }), + ), + ) + expect(rows).toHaveLength(count) + expect(new Set(rows.map((row) => row.requests)).size).toBe(count) + expect(rows.reduce((sum, row) => sum + Number(row.requests), 0)).toBe((count * (count + 1)) / 2) + expect(pages).toHaveLength(Math.floor(count / 10000) + 1) + expect(pages.at(-1)).toBe(count % 10000) + }) +} + +test("still rejects capped results without an aggregate cursor", async () => { + const result = await Effect.runPromise( + queryR2SqlPages("SELECT * FROM aggregates LIMIT 10000", undefined, () => + Effect.succeed(Array.from({ length: 10000 }, () => ({}))), + ).pipe(Effect.flip), + ) + expect(result.message).toContain("10000 row limit") +}) + +test("propagates a later page failure instead of returning partial aggregates", async () => { + const calls: string[] = [] + const error = new R2SqlQueryError({ message: "second page failed" }) + const result = await Effect.runPromise( + queryR2SqlPages("SELECT * FROM aggregates", ["model"], (query) => { + calls.push(query) + return calls.length === 1 + ? Effect.succeed(Array.from({ length: 10000 }, (_, index) => ({ model: String(index) }))) + : Effect.fail(error) + }).pipe(Effect.flip), + ) + expect(result).toBe(error) + expect(calls).toHaveLength(2) +}) diff --git a/packages/stats/core/src/r2-sql.ts b/packages/stats/core/src/r2-sql.ts index 1941d439ec4..b9f3af44976 100644 --- a/packages/stats/core/src/r2-sql.ts +++ b/packages/stats/core/src/r2-sql.ts @@ -37,7 +37,7 @@ export class R2SqlQueryError extends Error { export declare namespace R2Sql { export interface Service { - readonly query: (query: string) => Effect.Effect + readonly query: (query: string, columns?: readonly string[]) => Effect.Effect } } @@ -45,75 +45,99 @@ export class R2Sql extends Context.Service()("@opencode/st static readonly layer: Layer.Layer = Layer.succeed( R2Sql, R2Sql.of({ - query: Effect.fn("R2Sql.query")(function* (query: string) { - const startedAt = Date.now() - const response = yield* Effect.tryPromise({ - try: (signal) => { - const options = { - method: "POST", - // Analytical queries can exceed Bun's default five-minute idle timer. - // Bound the whole request and cancel it when the sync is interrupted. - timeout: false, - signal: AbortSignal.any([signal, AbortSignal.timeout(R2_SQL_TIMEOUT_MS)]), - headers: { - Authorization: `Bearer ${Resource.R2SqlAuthToken.value}`, - "Content-Type": "application/json", - }, - body: JSON.stringify({ query }), - } - return Bun.fetch( - `https://api.sql.cloudflarestorage.com/api/v1/accounts/${Resource.R2Sql.accountId}/r2-sql/query/${Resource.R2Sql.bucket}`, - options, - ) - }, - catch: (cause) => - new R2SqlQueryError({ - message: `Failed to run R2 SQL stats query after ${Date.now() - startedAt}ms`, - 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)]])), - ), - ) - }), + query: (query, columns) => queryR2SqlPages(query, columns, fetchRows), }), ) } + +// Cursor columns must uniquely identify an aggregate row. Page after aggregation +// so distinct counts, averages, and country totals retain their original scope. +export const queryR2SqlPages = Effect.fn("R2Sql.query")(function* ( + query: string, + columns: readonly string[] | undefined, + fetchRows: (query: string) => Effect.Effect, +) { + const rows: R2SqlData[] = [] + let cursor: R2SqlData | undefined + while (true) { + const page = yield* fetchRows(columns?.length ? pageQuery(query, columns, cursor) : query) + if (page.length >= R2_SQL_MAX_ROWS && !columns?.length) + return yield* Effect.fail( + new R2SqlQueryError({ message: `R2 SQL stats query reached the ${R2_SQL_MAX_ROWS} row limit` }), + ) + rows.push(...page) + if (page.length < R2_SQL_MAX_ROWS) return rows + cursor = page.at(-1)! + } +}) + +function pageQuery(query: string, columns: readonly string[], cursor?: R2SqlData) { + const keys = columns.map((column) => `COALESCE("${column.replace(/"/g, '""')}", '')`) + const values = columns.map((column) => `'${(cursor?.[column] ?? "").replace(/'/g, "''")}'`) + const after = keys.map((key, index) => + [ + ...keys.slice(0, index).map((prefix, before) => `${prefix} = ${values[before]}`), + `${key} > ${values[index]}`, + ].join(" AND "), + ) + return `SELECT * FROM (${query}) AS stats_page +${cursor ? `WHERE ${after.map((condition) => `(${condition})`).join(" OR ")}` : ""} +ORDER BY ${keys.join(", ")} +LIMIT ${R2_SQL_MAX_ROWS}` +} + +const fetchRows = Effect.fn("R2Sql.fetchRows")(function* (query: string) { + const startedAt = Date.now() + const response = yield* Effect.tryPromise({ + try: (signal) => { + const options = { + method: "POST", + // Analytical queries can exceed Bun's default five-minute idle timer. + // Bound the whole request and cancel it when the sync is interrupted. + timeout: false, + signal: AbortSignal.any([signal, AbortSignal.timeout(R2_SQL_TIMEOUT_MS)]), + headers: { + Authorization: `Bearer ${Resource.R2SqlAuthToken.value}`, + "Content-Type": "application/json", + }, + body: JSON.stringify({ query }), + } + return Bun.fetch( + `https://api.sql.cloudflarestorage.com/api/v1/accounts/${Resource.R2Sql.accountId}/r2-sql/query/${Resource.R2Sql.bucket}`, + options, + ) + }, + catch: (cause) => + new R2SqlQueryError({ + message: `Failed to run R2 SQL stats query after ${Date.now() - startedAt}ms`, + 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, + }), + ) + + 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/stat-sync.ts b/packages/stats/core/src/stat-sync.ts index ed71f2fee63..7cf53ae9368 100644 --- a/packages/stats/core/src/stat-sync.ts +++ b/packages/stats/core/src/stat-sync.ts @@ -54,7 +54,7 @@ export const syncStats: (options?: { const rows = yield* Effect.forEach( queries, (query, index) => - r2Sql.query(query).pipe( + r2Sql.query(query, ["dimension", "tier", "provider", "model", "country"]).pipe( Effect.tap((rows) => Effect.logInfo( `stats query complete ${JSON.stringify({ index, total: queries.length, rows: rows.length })}`,