fix(stats): unblock stats sync

This commit is contained in:
Adam
2026-06-21 14:53:46 -05:00
parent 69f1ec22e3
commit 4c6750d464
4 changed files with 83 additions and 44 deletions
+31 -19
View File
@@ -5,28 +5,36 @@ import {
StartQueryExecutionCommand, StartQueryExecutionCommand,
type Row, type Row,
} from "@aws-sdk/client-athena" } from "@aws-sdk/client-athena"
import { Effect, Layer, Schema } from "effect" import { Effect, Layer } from "effect"
import * as Context from "effect/Context" import * as Context from "effect/Context"
import { Resource } from "sst/resource" import { Resource } from "sst/resource"
const ATHENA_MAX_POLL_ATTEMPTS = 60 const ATHENA_MAX_POLL_ATTEMPTS = 300
const ATHENA_PAGE_SIZE = 1000 const ATHENA_PAGE_SIZE = 1000
export type AthenaData = Record<string, string> export type AthenaData = Record<string, string>
export class AthenaQueryError extends Schema.TaggedErrorClass<AthenaQueryError>()("AthenaQueryError", { export class AthenaQueryError extends Error {
message: Schema.String, readonly _tag = "AthenaQueryError"
queryExecutionId: Schema.optional(Schema.String), readonly queryExecutionId?: string
cause: Schema.optional(Schema.Defect()),
}) {}
export class AthenaQueryTimeoutError extends Schema.TaggedErrorClass<AthenaQueryTimeoutError>()( constructor(input: { message: string; queryExecutionId?: string; cause?: unknown }) {
"AthenaQueryTimeoutError", super(input.message, { cause: input.cause })
{ this.name = "AthenaQueryError"
message: Schema.String, this.queryExecutionId = input.queryExecutionId
queryExecutionId: Schema.String, }
}, }
) {}
export class AthenaQueryTimeoutError extends Error {
readonly _tag = "AthenaQueryTimeoutError"
readonly queryExecutionId: string
constructor(input: { message: string; queryExecutionId: string }) {
super(input.message)
this.name = "AthenaQueryTimeoutError"
this.queryExecutionId = input.queryExecutionId
}
}
export declare namespace Athena { export declare namespace Athena {
export interface Service { export interface Service {
@@ -57,7 +65,7 @@ export class Athena extends Context.Service<Athena, Athena.Service>()("@opencode
}) })
const queryExecutionId = started.QueryExecutionId const queryExecutionId = started.QueryExecutionId
if (!queryExecutionId) if (!queryExecutionId)
return yield* new AthenaQueryError({ message: "Athena did not return a query execution id" }) return yield* Effect.fail(new AthenaQueryError({ message: "Athena did not return a query execution id" }))
yield* poll(client, queryExecutionId) yield* poll(client, queryExecutionId)
return yield* results(client, queryExecutionId) return yield* results(client, queryExecutionId)
@@ -87,16 +95,20 @@ const poll: (
if (status?.State === "SUCCEEDED") return if (status?.State === "SUCCEEDED") return
if (status?.State === "FAILED" || status?.State === "CANCELLED") if (status?.State === "FAILED" || status?.State === "CANCELLED")
return yield* new AthenaQueryError({ return yield* Effect.fail(
new AthenaQueryError({
message: `Athena stats query ${status.State.toLowerCase()}: ${status.StateChangeReason ?? "unknown reason"}`, message: `Athena stats query ${status.State.toLowerCase()}: ${status.StateChangeReason ?? "unknown reason"}`,
queryExecutionId, queryExecutionId,
}) }),
)
if (attempt >= ATHENA_MAX_POLL_ATTEMPTS - 1) if (attempt >= ATHENA_MAX_POLL_ATTEMPTS - 1)
return yield* new AthenaQueryTimeoutError({ return yield* Effect.fail(
new AthenaQueryTimeoutError({
message: `Athena stats query ${queryExecutionId} did not complete`, message: `Athena stats query ${queryExecutionId} did not complete`,
queryExecutionId, queryExecutionId,
}) }),
)
return yield* poll(client, queryExecutionId, attempt + 1) return yield* poll(client, queryExecutionId, attempt + 1)
}) })
+24 -9
View File
@@ -44,16 +44,29 @@ export class DrizzleClient extends Context.Service<DrizzleClient, Drizzle>()("@o
) )
} }
export class DatabaseError extends Schema.TaggedErrorClass<DatabaseError>()("DatabaseError", { export class DatabaseError extends Error {
cause: Schema.Defect(), readonly _tag = "DatabaseError"
}) {}
constructor(input: { cause: unknown }) {
super("Database operation failed", { cause: input.cause })
this.name = "DatabaseError"
}
static make(input: { cause: unknown }) {
return new DatabaseError(input)
}
}
export const catchDbError = Effect.mapError((cause) => DatabaseError.make({ cause })) export const catchDbError = Effect.mapError((cause) => DatabaseError.make({ cause }))
export class MigrationError extends Schema.TaggedErrorClass<MigrationError>()("MigrationError", { export class MigrationError extends Error {
message: Schema.String, readonly _tag = "MigrationError"
cause: Schema.optional(Schema.Defect()),
}) {} constructor(input: { message: string; cause?: unknown }) {
super(input.message, { cause: input.cause })
this.name = "MigrationError"
}
}
export const migrate = Effect.fn("Database.migrate")(function* () { export const migrate = Effect.fn("Database.migrate")(function* () {
const settings = yield* DatabaseConfig const settings = yield* DatabaseConfig
@@ -68,9 +81,11 @@ export const migrate = Effect.fn("Database.migrate")(function* () {
catch: (cause) => new MigrationError({ message: "Failed to apply database migrations", cause }), catch: (cause) => new MigrationError({ message: "Failed to apply database migrations", cause }),
}) })
if (result) if (result)
return yield* new MigrationError({ return yield* Effect.fail(
new MigrationError({
message: `Failed to initialize database migrations: ${result.exitCode}`, message: `Failed to initialize database migrations: ${result.exitCode}`,
}) }),
)
yield* Effect.logInfo("database migrations complete").pipe( yield* Effect.logInfo("database migrations complete").pipe(
Effect.annotateLogs({ migrationsDir: settings.migrationsDir }), Effect.annotateLogs({ migrationsDir: settings.migrationsDir }),
) )
@@ -17,6 +17,8 @@ export type StatDimension = "model" | "provider" | "geo" | "geo_model"
export function buildStatsQuery(periodStart: Date, periodEnd: Date, dimension: StatDimension) { export function buildStatsQuery(periodStart: Date, periodEnd: Date, dimension: StatDimension) {
const periodStartValue = sqlString(periodStart.toISOString()) const periodStartValue = sqlString(periodStart.toISOString())
const periodEndValue = sqlString(periodEnd.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] const sourceTable = [Resource.InferenceEvent.catalog, Resource.InferenceEvent.database, Resource.InferenceEvent.table]
.map(sqlIdentifier) .map(sqlIdentifier)
.join(".") .join(".")
@@ -95,6 +97,9 @@ WITH normalized AS (
WHERE event_type = 'completions' WHERE event_type = 'completions'
AND model IS NOT NULL AND model IS NOT NULL
AND model <> '' AND model <> ''
AND source = 'lite'
AND event_date >= ${periodStartDateValue}
AND event_date <= ${periodEndDateValue}
AND event_timestamp >= ${periodStartValue} AND event_timestamp >= ${periodStartValue}
AND event_timestamp < ${periodEndValue} AND event_timestamp < ${periodEndValue}
), filtered AS ( ), filtered AS (
+16 -9
View File
@@ -1,6 +1,6 @@
import { Buffer } from "node:buffer" import { Buffer } from "node:buffer"
import { FirehoseClient, PutRecordBatchCommand } from "@aws-sdk/client-firehose" import { FirehoseClient, PutRecordBatchCommand } from "@aws-sdk/client-firehose"
import { Effect, Layer, Schema } from "effect" import { Effect, Layer } from "effect"
import * as Context from "effect/Context" import * as Context from "effect/Context"
import { Resource } from "sst/resource" import { Resource } from "sst/resource"
@@ -12,11 +12,16 @@ type IngestEvent = Record<string, unknown>
type LakeRoute = { database: string; table: string } type LakeRoute = { database: string; table: string }
type FirehoseRecord = { Data: Uint8Array } type FirehoseRecord = { Data: Uint8Array }
export class IngestError extends Schema.TaggedErrorClass<IngestError>()("IngestError", { export class IngestError extends Error {
message: Schema.String, readonly _tag = "IngestError"
failed: Schema.Number, readonly failed: number
cause: Schema.optional(Schema.Defect()),
}) {} constructor(input: { message: string; failed: number; cause?: unknown }) {
super(input.message, { cause: input.cause })
this.name = "IngestError"
this.failed = input.failed
}
}
export declare namespace Ingest { export declare namespace Ingest {
export interface Service { export interface Service {
@@ -37,10 +42,12 @@ export class Ingest extends Context.Service<Ingest, Ingest.Service>()("@opencode
yield* Effect.logWarning( yield* Effect.logWarning(
`lake ingest rejected ${JSON.stringify({ records: counts.records, unsupported: counts.unsupported })}`, `lake ingest rejected ${JSON.stringify({ records: counts.records, unsupported: counts.unsupported })}`,
) )
return yield* new IngestError({ return yield* Effect.fail(
new IngestError({
message: "Unsupported lake event type", message: "Unsupported lake event type",
failed: counts.unsupported, failed: counts.unsupported,
}) }),
)
} }
if (counts.records === 0) return { records: 0 } if (counts.records === 0) return { records: 0 }
@@ -66,7 +73,7 @@ export class Ingest extends Context.Service<Ingest, Ingest.Service>()("@opencode
if (failed > 0) { if (failed > 0) {
yield* Effect.logWarning(`lake ingest incomplete ${JSON.stringify({ records: counts.records, failed })}`) yield* Effect.logWarning(`lake ingest incomplete ${JSON.stringify({ records: counts.records, failed })}`)
return yield* new IngestError({ message: "Failed to ingest all lake records", failed }) return yield* Effect.fail(new IngestError({ message: "Failed to ingest all lake records", failed }))
} }
yield* Effect.logInfo(`lake ingest complete ${JSON.stringify({ records: counts.records, batches })}`) yield* Effect.logInfo(`lake ingest complete ${JSON.stringify({ records: counts.records, batches })}`)