Move session usage backfill to data migration
This commit is contained in:
@@ -3,24 +3,4 @@ ALTER TABLE `session` ADD `tokens_input` integer DEFAULT 0 NOT NULL;--> statemen
|
||||
ALTER TABLE `session` ADD `tokens_output` integer DEFAULT 0 NOT NULL;--> statement-breakpoint
|
||||
ALTER TABLE `session` ADD `tokens_reasoning` integer DEFAULT 0 NOT NULL;--> statement-breakpoint
|
||||
ALTER TABLE `session` ADD `tokens_cache_read` integer DEFAULT 0 NOT NULL;--> statement-breakpoint
|
||||
ALTER TABLE `session` ADD `tokens_cache_write` integer DEFAULT 0 NOT NULL;--> statement-breakpoint
|
||||
CREATE TEMP TABLE `session_usage` AS
|
||||
SELECT
|
||||
`part`.`session_id`,
|
||||
COALESCE(SUM(COALESCE(json_extract(`part`.`data`, '$.cost'), 0)), 0) AS `cost`,
|
||||
COALESCE(SUM(COALESCE(json_extract(`part`.`data`, '$.tokens.input'), 0)), 0) AS `tokens_input`,
|
||||
COALESCE(SUM(COALESCE(json_extract(`part`.`data`, '$.tokens.output'), 0)), 0) AS `tokens_output`,
|
||||
COALESCE(SUM(COALESCE(json_extract(`part`.`data`, '$.tokens.reasoning'), 0)), 0) AS `tokens_reasoning`,
|
||||
COALESCE(SUM(COALESCE(json_extract(`part`.`data`, '$.tokens.cache.read'), 0)), 0) AS `tokens_cache_read`,
|
||||
COALESCE(SUM(COALESCE(json_extract(`part`.`data`, '$.tokens.cache.write'), 0)), 0) AS `tokens_cache_write`
|
||||
FROM `part`
|
||||
WHERE json_extract(`part`.`data`, '$.type') = 'step-finish'
|
||||
GROUP BY `part`.`session_id`;--> statement-breakpoint
|
||||
UPDATE `session` SET
|
||||
`cost` = COALESCE((SELECT `cost` FROM `session_usage` WHERE `session_usage`.`session_id` = `session`.`id`), 0),
|
||||
`tokens_input` = COALESCE((SELECT `tokens_input` FROM `session_usage` WHERE `session_usage`.`session_id` = `session`.`id`), 0),
|
||||
`tokens_output` = COALESCE((SELECT `tokens_output` FROM `session_usage` WHERE `session_usage`.`session_id` = `session`.`id`), 0),
|
||||
`tokens_reasoning` = COALESCE((SELECT `tokens_reasoning` FROM `session_usage` WHERE `session_usage`.`session_id` = `session`.`id`), 0),
|
||||
`tokens_cache_read` = COALESCE((SELECT `tokens_cache_read` FROM `session_usage` WHERE `session_usage`.`session_id` = `session`.`id`), 0),
|
||||
`tokens_cache_write` = COALESCE((SELECT `tokens_cache_write` FROM `session_usage` WHERE `session_usage`.`session_id` = `session`.`id`), 0);--> statement-breakpoint
|
||||
DROP TABLE `session_usage`;
|
||||
ALTER TABLE `session` ADD `tokens_cache_write` integer DEFAULT 0 NOT NULL;
|
||||
|
||||
@@ -2,7 +2,9 @@ import { Context, Effect, Layer } from "effect"
|
||||
import { Database } from "./storage/db"
|
||||
import { DataMigrationTable } from "./data-migration.sql"
|
||||
import * as Log from "@opencode-ai/core/util/log"
|
||||
import { eq } from "drizzle-orm"
|
||||
import { asc, eq, gt, inArray } from "drizzle-orm"
|
||||
import { PartTable, SessionTable } from "./session/session.sql"
|
||||
import type { SessionID } from "./session/schema"
|
||||
|
||||
export type Migration<R = never> = {
|
||||
name: string
|
||||
@@ -18,7 +20,84 @@ export class Service extends Context.Service<Service, Interface>()("@opencode/Da
|
||||
export const layer = Layer.effect(
|
||||
Service,
|
||||
Effect.gen(function* () {
|
||||
const migrations: Migration[] = []
|
||||
const migrations: Migration[] = [
|
||||
{
|
||||
name: "session_usage_from_parts",
|
||||
run: Effect.sync(() =>
|
||||
Database.transaction((db) => {
|
||||
type Usage = {
|
||||
cost: number
|
||||
tokens: {
|
||||
input: number
|
||||
output: number
|
||||
reasoning: number
|
||||
cache: { read: number; write: number }
|
||||
}
|
||||
}
|
||||
|
||||
for (let cursor: SessionID | undefined; ; ) {
|
||||
const sessions = db
|
||||
.select({ id: SessionTable.id })
|
||||
.from(SessionTable)
|
||||
.where(cursor ? gt(SessionTable.id, cursor) : undefined)
|
||||
.orderBy(asc(SessionTable.id))
|
||||
.limit(500)
|
||||
.all()
|
||||
if (sessions.length === 0) return
|
||||
|
||||
const usageBySession = new Map<SessionID, Usage>(
|
||||
sessions.map((session) => [
|
||||
session.id,
|
||||
{ cost: 0, tokens: { input: 0, output: 0, reasoning: 0, cache: { read: 0, write: 0 } } },
|
||||
]),
|
||||
)
|
||||
|
||||
for (const row of db
|
||||
.select({ session_id: PartTable.session_id, data: PartTable.data })
|
||||
.from(PartTable)
|
||||
.where(inArray(PartTable.session_id, sessions.map((session) => session.id)))
|
||||
.all()) {
|
||||
const part = row.data
|
||||
if (part.type !== "step-finish") continue
|
||||
if (!("cost" in part) || typeof part.cost !== "number") continue
|
||||
if (!("tokens" in part) || typeof part.tokens !== "object" || part.tokens === null) continue
|
||||
if (!("input" in part.tokens) || typeof part.tokens.input !== "number") continue
|
||||
if (!("output" in part.tokens) || typeof part.tokens.output !== "number") continue
|
||||
if (!("reasoning" in part.tokens) || typeof part.tokens.reasoning !== "number") continue
|
||||
if (!("cache" in part.tokens) || typeof part.tokens.cache !== "object" || part.tokens.cache === null) continue
|
||||
if (!("read" in part.tokens.cache) || typeof part.tokens.cache.read !== "number") continue
|
||||
if (!("write" in part.tokens.cache) || typeof part.tokens.cache.write !== "number") continue
|
||||
|
||||
const current = usageBySession.get(row.session_id)
|
||||
if (!current) continue
|
||||
current.cost += part.cost
|
||||
current.tokens.input += part.tokens.input
|
||||
current.tokens.output += part.tokens.output
|
||||
current.tokens.reasoning += part.tokens.reasoning
|
||||
current.tokens.cache.read += part.tokens.cache.read
|
||||
current.tokens.cache.write += part.tokens.cache.write
|
||||
}
|
||||
|
||||
for (const [sessionID, value] of usageBySession) {
|
||||
db.update(SessionTable)
|
||||
.set({
|
||||
cost: value.cost,
|
||||
tokens_input: value.tokens.input,
|
||||
tokens_output: value.tokens.output,
|
||||
tokens_reasoning: value.tokens.reasoning,
|
||||
tokens_cache_read: value.tokens.cache.read,
|
||||
tokens_cache_write: value.tokens.cache.write,
|
||||
})
|
||||
.where(eq(SessionTable.id, sessionID))
|
||||
.run()
|
||||
}
|
||||
|
||||
cursor = sessions.at(-1)?.id
|
||||
}
|
||||
}),
|
||||
),
|
||||
},
|
||||
]
|
||||
|
||||
yield* Effect.gen(function* () {
|
||||
if (migrations.length === 0) return
|
||||
|
||||
Reference in New Issue
Block a user