fix(stats): unblock stats sync
This commit is contained in:
@@ -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(
|
||||||
message: `Athena stats query ${status.State.toLowerCase()}: ${status.StateChangeReason ?? "unknown reason"}`,
|
new AthenaQueryError({
|
||||||
queryExecutionId,
|
message: `Athena stats query ${status.State.toLowerCase()}: ${status.StateChangeReason ?? "unknown reason"}`,
|
||||||
})
|
queryExecutionId,
|
||||||
|
}),
|
||||||
|
)
|
||||||
|
|
||||||
if (attempt >= ATHENA_MAX_POLL_ATTEMPTS - 1)
|
if (attempt >= ATHENA_MAX_POLL_ATTEMPTS - 1)
|
||||||
return yield* new AthenaQueryTimeoutError({
|
return yield* Effect.fail(
|
||||||
message: `Athena stats query ${queryExecutionId} did not complete`,
|
new AthenaQueryTimeoutError({
|
||||||
queryExecutionId,
|
message: `Athena stats query ${queryExecutionId} did not complete`,
|
||||||
})
|
queryExecutionId,
|
||||||
|
}),
|
||||||
|
)
|
||||||
|
|
||||||
return yield* poll(client, queryExecutionId, attempt + 1)
|
return yield* poll(client, queryExecutionId, attempt + 1)
|
||||||
})
|
})
|
||||||
|
|||||||
@@ -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(
|
||||||
message: `Failed to initialize database migrations: ${result.exitCode}`,
|
new MigrationError({
|
||||||
})
|
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 (
|
||||||
|
|||||||
@@ -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(
|
||||||
message: "Unsupported lake event type",
|
new IngestError({
|
||||||
failed: counts.unsupported,
|
message: "Unsupported lake event type",
|
||||||
})
|
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 })}`)
|
||||||
|
|||||||
Reference in New Issue
Block a user