File size: 3,596 Bytes
4bbfe8b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 | import { and, eq, inArray } from "drizzle-orm"
import { Context, Effect, Layer } from "effect"
import { DatabaseError, DrizzleClient } from "../database"
import { modelRetention } from "../database/schema"
import { chunks, UPSERT_CHUNK_SIZE } from "./stat"
export type RetentionStatRow = typeof modelRetention.$inferInsert
export type RetentionStatAggregate = {
cohortDate: string
dataset: string
tier: string
provider: string
model: string
eligibleUsers: number
retainedUsers: number
}
export declare namespace RetentionStatRepo {
export interface Service {
readonly available: () => Effect.Effect<boolean, DatabaseError>
readonly replace: (
rows: RetentionStatRow[],
scope: { cohortDates: string[]; dataset: string; tier: string },
) => Effect.Effect<void, DatabaseError>
}
}
export class RetentionStatRepo extends Context.Service<RetentionStatRepo, RetentionStatRepo.Service>()(
"@opencode/stats/RetentionStatRepo",
) {
static readonly layer: Layer.Layer<RetentionStatRepo, never, DrizzleClient> = Layer.effect(
RetentionStatRepo,
Effect.gen(function* () {
const db = yield* DrizzleClient
const available = Effect.fn("RetentionStatRepo.available")(function* () {
return yield* Effect.tryPromise({
try: async () => {
try {
await db.select({ id: modelRetention.id }).from(modelRetention).limit(1)
return true
} catch (cause) {
if (isMissingRetentionTable(cause)) return false
throw cause
}
},
catch: (cause) => DatabaseError.make({ cause }),
})
})
const replace = Effect.fn("RetentionStatRepo.replace")(function* (
rows: RetentionStatRow[],
scope: { cohortDates: string[]; dataset: string; tier: string },
) {
if (scope.cohortDates.length === 0) return
yield* Effect.tryPromise({
try: () =>
db
.delete(modelRetention)
.where(
and(
inArray(modelRetention.cohort_date, scope.cohortDates),
eq(modelRetention.dataset, scope.dataset),
eq(modelRetention.tier, scope.tier),
),
),
catch: (cause) => DatabaseError.make({ cause }),
})
yield* Effect.forEach(
chunks(rows, UPSERT_CHUNK_SIZE),
(chunk) =>
Effect.tryPromise({
try: () => db.insert(modelRetention).values(chunk),
catch: (cause) => DatabaseError.make({ cause }),
}),
{ discard: true },
)
})
return RetentionStatRepo.of({ available, replace })
}),
)
}
export function rowsFromAggregates(aggregates: RetentionStatAggregate[]): RetentionStatRow[] {
return aggregates.map((row) => ({
cohort_date: row.cohortDate,
dataset: row.dataset,
tier: row.tier,
provider: row.provider,
model: row.model,
eligible_users: row.eligibleUsers,
retained_users: row.retainedUsers,
}))
}
export function isMissingRetentionTable(cause: unknown): boolean {
const text = errorText(cause).toLowerCase()
return text.includes("model_retention") && text.includes("exist")
}
function errorText(cause: unknown): string {
if (cause instanceof Error) return `${cause.message} ${errorText((cause as { cause?: unknown }).cause)}`
if (typeof cause === "object" && cause)
return Object.values(cause as Record<string, unknown>)
.map(errorText)
.join(" ")
return String(cause)
}
|