fix(stats): paginate sync aggregates (#50198)

This commit is contained in:
Adam
2026-09-20 12:00:07 -05:00
committed by GitHub
parent 8b2de7317b
commit 45ad8dc38a
5 changed files with 161 additions and 77 deletions
@@ -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")
+3 -4
View File
@@ -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
`
}
+62 -1
View File
@@ -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)
})
+94 -70
View File
@@ -37,7 +37,7 @@ export class R2SqlQueryError extends Error {
export declare namespace R2Sql {
export interface Service {
readonly query: (query: string) => Effect.Effect<R2SqlData[], R2SqlQueryError>
readonly query: (query: string, columns?: readonly string[]) => Effect.Effect<R2SqlData[], R2SqlQueryError>
}
}
@@ -45,75 +45,99 @@ export class R2Sql extends Context.Service<R2Sql, R2Sql.Service>()("@opencode/st
static readonly layer: Layer.Layer<R2Sql> = 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<R2SqlData[], R2SqlQueryError>,
) {
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)]]))),
)
})
+1 -1
View File
@@ -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 })}`,