package store import ( "context" "database/sql" "encoding/json" "errors" "fmt" "os" "slices" "strings" "sync" "time" _ "modernc.org/sqlite" // the cgo-free driver, so the binary stays static ) // DB is the store, on either dialect. // // It was called SQLite, and kept the name after PostgreSQL was added, so // OpenPostgres returned a *SQLite and every one of the seventy-odd methods on // the production store was spelled after the database it was not running on. // Nothing broke; people did. A method written for a table that only ever lives // in PostgreSQL still read as SQLite code, and the deployment that matters most // looked like the one nobody uses. // // The dialect is a field, which is all it ever was. type DB struct { // raw is for the things that are not queries: transactions, pings, close. raw *sql.DB dialect Dialect // db is the same handle with placeholders rewritten for the dialect. Every // query goes through it, which is the point — a call site that reached the // driver directly would work on SQLite and fail on PostgreSQL, and only on // the path nobody exercised before deploying. db bound // mu guards the default limits, which are set once at wiring but read on // every authentication. mu sync.RWMutex defaultRPM *int defaultTPM *int // personal is what a person's own use is held to when nobody wrote a // limit for them; see SetPersonalDefaults. personal *Limits } // OpenSQLite opens (creating if needed) a database and brings it up to date. // // The pragmas are the ones that matter. WAL gives concurrent readers alongside // the single writer, which is the access pattern exactly: many reads for // authentication, one batched writer for usage. foreign_keys is off by default // in SQLite, which would silently let the tenancy tree grow orphans. busy_timeout // turns "database is locked" from an error into a wait, which is what a caller // wants when the batch writer holds the write lock for a few milliseconds. func OpenSQLite(ctx context.Context, path string) (*DB, error) { dsn := path + "?_pragma=journal_mode(WAL)&_pragma=foreign_keys(1)&_pragma=busy_timeout(5000)&_pragma=synchronous(NORMAL)" db, err := sql.Open("sqlite", dsn) if err != nil { return nil, err } // One writer is what SQLite supports; more connections would only produce // lock contention that looks like a deadlock under load. db.SetMaxOpenConns(1) if err := db.PingContext(ctx); err != nil { _ = db.Close() return nil, fmt.Errorf("opening %s: %w", path, err) } if err := migrate(ctx, db, "sqlite"); err != nil { _ = db.Close() return nil, err } // SQLite creates the file with the process umask, which on a stock system // means world-readable. What is in here is not a working credential — keys // are stored hashed — but it is the whole tenancy tree and every request // anyone made, and that should not be readable by every account on the // machine because nobody remembered to chmod it. for _, f := range []string{path, path + "-wal", path + "-shm"} { if err := os.Chmod(f, 0o600); err != nil && !os.IsNotExist(err) { _ = db.Close() return nil, fmt.Errorf("securing %s: %w", f, err) } } return &DB{raw: db, dialect: SQLiteDialect, db: bound{c: SQLiteDialect, to: db}}, nil } // Close releases the database. func (s *DB) Close() error { return s.raw.Close() } // --- tenancy --- // PutOrg upserts an organization. func (s *DB) PutOrg(ctx context.Context, o Org) error { _, err := s.db.ExecContext(ctx, `INSERT INTO organizations (id, name) VALUES (?, ?) ON CONFLICT (id) DO UPDATE SET name = excluded.name`, o.ID, o.Name) return err } // PutTeam upserts a team. // // It will rename a team but not move it: an upsert that rewrote the parent // let `keys init --org B --team platform` take organization A's team without // a word. Moving is MoveTeam, asked for by name. func (s *DB) PutTeam(ctx context.Context, t Team) error { var org string err := s.db.QueryRowContext(ctx, `SELECT org_id FROM teams WHERE id = ?`, t.ID).Scan(&org) if err == nil && org != t.OrgID { return fmt.Errorf("team %s belongs to organization %s; move it rather than redeclare it", t.ID, org) } _, err = s.db.ExecContext(ctx, `INSERT INTO teams (id, org_id, name) VALUES (?, ?, ?) ON CONFLICT (id) DO UPDATE SET org_id = excluded.org_id, name = excluded.name`, t.ID, t.OrgID, t.Name) return err } // PutProject upserts a project. Like PutTeam, it renames but never moves. func (s *DB) PutProject(ctx context.Context, p Project) error { var team string err := s.db.QueryRowContext(ctx, `SELECT team_id FROM projects WHERE id = ?`, p.ID).Scan(&team) if err == nil && team != p.TeamID { return fmt.Errorf("project %s belongs to team %s; move it rather than redeclare it", p.ID, team) } _, err = s.db.ExecContext(ctx, `INSERT INTO projects (id, team_id, name) VALUES (?, ?, ?) ON CONFLICT (id) DO UPDATE SET team_id = excluded.team_id, name = excluded.name`, p.ID, p.TeamID, p.Name) return err } // Orgs lists every organization. func (s *DB) Orgs(ctx context.Context) ([]Org, error) { return queryAll(ctx, s.db, `SELECT id, name FROM organizations ORDER BY name`, nil, func(rows *sql.Rows) (Org, error) { var o Org return o, rows.Scan(&o.ID, &o.Name) }) } // Teams lists an organization's teams. func (s *DB) Teams(ctx context.Context, orgID string) ([]Team, error) { return queryAll(ctx, s.db, `SELECT id, org_id, name FROM teams WHERE org_id = ? ORDER BY name`, []any{orgID}, func(rows *sql.Rows) (Team, error) { var t Team return t, rows.Scan(&t.ID, &t.OrgID, &t.Name) }) } // Projects lists a team's projects. func (s *DB) Projects(ctx context.Context, teamID string) ([]Project, error) { return queryAll(ctx, s.db, `SELECT id, team_id, name FROM projects WHERE team_id = ? ORDER BY name`, []any{teamID}, func(rows *sql.Rows) (Project, error) { var p Project return p, rows.Scan(&p.ID, &p.TeamID, &p.Name) }) } // --- keys --- const keyColumns = `id, key_hash, key_prefix, name, project_id, band, routing_profile, created_at, expires_at, revoked_at, last_used_at, user_id` // CreateKey stores a key. Only its hash is written; the key itself is the // caller's to show once and then forget. func (s *DB) CreateKey(ctx context.Context, k Key) error { if len(k.Hash) == 0 { return errors.New("store: a key needs a hash") } _, err := s.db.ExecContext(ctx, `INSERT INTO api_keys (`+keyColumns+`) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, k.ID, k.Hash, k.Prefix, k.Name, nullString(k.ProjectID), nullString(k.Band), nullString(k.RoutingProfile), millis(k.CreatedAt), nullMillis(k.ExpiresAt), nullMillis(k.RevokedAt), nullMillis(k.LastUsedAt), nullString(k.UserID)) return err } // KeyByHash finds a key by the hash of the presented secret. func (s *DB) KeyByHash(ctx context.Context, hash []byte) (Key, error) { row := s.db.QueryRowContext(ctx, `SELECT `+keyColumns+` FROM api_keys WHERE key_hash = ?`, hash) k, err := scanKey(row) if errors.Is(err, sql.ErrNoRows) { return Key{}, ErrNotFound } return k, err } // Keys lists a project's keys, revoked ones included. func (s *DB) Keys(ctx context.Context, projectID string) ([]Key, error) { return queryAll(ctx, s.db, `SELECT `+keyColumns+` FROM api_keys WHERE project_id = ? ORDER BY created_at`, []any{projectID}, func(rows *sql.Rows) (Key, error) { return scanKey(rows) }) } // RevokeKey marks a key unusable without deleting it. func (s *DB) RevokeKey(ctx context.Context, id string, at time.Time) error { res, err := s.db.ExecContext(ctx, `UPDATE api_keys SET revoked_at = ? WHERE id = ? AND revoked_at IS NULL`, millis(at), id) if err != nil { return err } if n, _ := res.RowsAffected(); n == 0 { // Either it does not exist or it was already revoked. Both mean the key // cannot authenticate, so only the first is worth an error. var exists int if err := s.db.QueryRowContext(ctx, `SELECT COUNT(*) FROM api_keys WHERE id = ?`, id). Scan(&exists); err != nil { return err } if exists == 0 { return ErrNotFound } } return nil } // TouchKeys records when keys were last used. // // Batched, and at most once a minute per key, because nothing reads this field // in real time and writing it per request would put a database write back on // the request path. func (s *DB) TouchKeys(ctx context.Context, seen map[string]time.Time) error { if len(seen) == 0 { return nil } return s.inTx(ctx, func(tx conn) error { stmt, err := tx.PrepareContext(ctx, `UPDATE api_keys SET last_used_at = ? WHERE id = ? AND (last_used_at IS NULL OR last_used_at < ?)`) if err != nil { return err } defer stmt.Close() for id, at := range seen { if _, err := stmt.ExecContext(ctx, millis(at), id, millis(at)); err != nil { return err } } return nil }) } // --- limits --- // PutLimits upserts the limits at one scope. func (s *DB) PutLimits(ctx context.Context, l Limits) error { models, err := encodeModels(l.AllowedModels) if err != nil { return err } _, err = s.db.ExecContext(ctx, `INSERT INTO limits (scope_type, scope_id, rpm, tpm, budget_micros, budget_period, allowed_models) VALUES (?, ?, ?, ?, ?, ?, ?) ON CONFLICT (scope_type, scope_id) DO UPDATE SET rpm = excluded.rpm, tpm = excluded.tpm, budget_micros = excluded.budget_micros, budget_period = excluded.budget_period, allowed_models = excluded.allowed_models`, string(l.Scope), l.ScopeID, nullInt(l.RPM), nullInt(l.TPM), nullInt64(l.Budget), nullString(string(l.BudgetPeriod)), models) return err } // LimitsFor reads the limits set at one scope, without inheritance. func (s *DB) LimitsFor(ctx context.Context, scope Scope, id string) (Limits, error) { row := s.db.QueryRowContext(ctx, `SELECT rpm, tpm, budget_micros, budget_period, allowed_models FROM limits WHERE scope_type = ? AND scope_id = ?`, string(scope), id) l := Limits{Scope: scope, ScopeID: id} var rpm, tpm sql.NullInt64 var budget sql.NullInt64 var period, models sql.NullString if err := row.Scan(&rpm, &tpm, &budget, &period, &models); err != nil { if errors.Is(err, sql.ErrNoRows) { return Limits{}, ErrNotFound } return Limits{}, err } if rpm.Valid { v := int(rpm.Int64) l.RPM = &v } if tpm.Valid { v := int(tpm.Int64) l.TPM = &v } if budget.Valid { v := budget.Int64 l.Budget = &v } l.BudgetPeriod = Period(period.String) var err error if l.AllowedModels, err = decodeModels(models); err != nil { return Limits{}, err } return l, nil } // SetDefaultLimits sets what applies to a key no scope gives a limit to. // // A default is a fallback, not a ceiling, so it is applied after the chain and // only where the chain left a limit unset — an explicit project limit above the // default must win, which it would not if this were folded in as the widest // scope and tightened. // // It exists because unlimited was the default. A project created without // limits — which is the easy path, and what the "default" project on this // cluster still is — metered nothing at all, and the first sign would have // been the bill or the outage. func (s *DB) SetDefaultLimits(rpm, tpm *int) { s.mu.Lock() defer s.mu.Unlock() s.defaultRPM, s.defaultTPM = rpm, tpm } // applyDefaults fills in what the chain left unset. // // Owned by the key, not by any scope. A cluster-wide default is a fallback for // each caller that was given no limit of its own — read as a shared pool it // would become one cluster-wide rate limit that every key in every // organization competes for, which is not a default, it is an outage. func (s *DB) applyDefaults(e *Effective) { s.mu.RLock() defer s.mu.RUnlock() own := LimitOwner{Scope: ScopeKey, ID: e.KeyID} fallback := Limits{Scope: ScopeKey, ScopeID: e.KeyID} if e.RPM == nil && s.defaultRPM != nil { v := *s.defaultRPM e.RPM, e.RPMOwner = &v, own fallback.RPM = &v } if e.TPM == nil && s.defaultTPM != nil { v := *s.defaultTPM e.TPM, e.TPMOwner = &v, own fallback.TPM = &v } // Appended so enforcement sees it like any other level. Left out, a cluster // default would be reported as the ceiling and then never checked. if fallback.RPM != nil || fallback.TPM != nil { e.Chain = append(e.Chain, fallback) } } // SpendSinceByScope totals cost per key and person, with the key's place in // the tree. Personal keys and the chat page's requests have no project and are // kept: dropping them is what would reset a person's budget on every restart. func (s *DB) SpendSinceByScope(ctx context.Context, since time.Time) ([]ScopeSpend, error) { rows, err := s.db.QueryContext(ctx, `SELECT u.key_id, u.user_id, COALESCE(p.id, ''), COALESCE(t.id, ''), COALESCE(t.org_id, ''), SUM(u.cost_micros) FROM usage u LEFT JOIN projects p ON p.id = NULLIF(u.project_id, '') LEFT JOIN teams t ON t.id = p.team_id WHERE u.ts >= ? GROUP BY u.key_id, u.user_id, p.id, t.id, t.org_id`, millis(since)) if err != nil { return nil, err } defer rows.Close() var out []ScopeSpend for rows.Next() { var sp ScopeSpend var total sql.NullInt64 if err := rows.Scan(&sp.KeyID, &sp.UserID, &sp.ProjectID, &sp.TeamID, &sp.OrgID, &total); err != nil { return nil, err } sp.Micros = total.Int64 out = append(out, sp) } return out, rows.Err() } // tighten applies one scope's limits, keeping whichever is stricter. // // Each value that wins records the scope it came from. That is what lets // consumption be counted against the scope that set the limit: a budget written // on a team is one pool the team shares, not a copy of the number handed to // every key beneath it. func (e *Effective) tighten(l Limits) { owner := LimitOwner{Scope: l.Scope, ID: l.ScopeID} if next := minPtr(e.RPM, l.RPM); next != e.RPM { e.RPM, e.RPMOwner = next, owner } if next := minPtr(e.TPM, l.TPM); next != e.TPM { e.TPM, e.TPMOwner = next, owner } if l.Budget != nil && (e.Budget == nil || *l.Budget < *e.Budget) { e.Budget = l.Budget e.BudgetPeriod = l.BudgetPeriod e.BudgetOwner = owner } if l.AllowedModels != nil { if e.AllowedModels == nil { e.AllowedModels = slices.Clone(l.AllowedModels) } else { // Intersection: a narrower scope cannot add a model its parent // withheld, or an allowlist would be advisory. var keep []string for _, m := range e.AllowedModels { if slices.Contains(l.AllowedModels, m) { keep = append(keep, m) } } // An empty intersection is still an allowlist, and an empty one. if keep == nil { keep = []string{} } e.AllowedModels = keep } } } // --- accounting --- // WriteUsage appends a batch of records in one transaction. func (s *DB) WriteUsage(ctx context.Context, records []Usage) error { if len(records) == 0 { return nil } return s.inTx(ctx, func(tx conn) error { stmt, err := tx.PrepareContext(ctx, `INSERT INTO usage (ts, key_id, project_id, model, served_model, prompt_tokens, completion_tokens, cost_micros, status, latency_ms, user_id, request_id, replica, queue_ms, ttft_ms) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`) if err != nil { return err } defer stmt.Close() for _, r := range records { if _, err := stmt.ExecContext(ctx, millis(r.At), r.KeyID, r.ProjectID, r.Model, r.ServedModel, r.PromptTokens, r.CompletionTokens, r.CostMicros, r.Status, r.LatencyMs, r.UserID, r.RequestID, r.Replica, r.QueueMs, r.TTFTMs); err != nil { return err } } return nil }) } // SpendSince totals cost per key since a moment. func (s *DB) SpendSince(ctx context.Context, since time.Time) (map[string]int64, error) { rows, err := s.db.QueryContext(ctx, `SELECT key_id, SUM(cost_micros) FROM usage WHERE ts >= ? GROUP BY key_id`, millis(since)) if err != nil { return nil, err } defer rows.Close() out := map[string]int64{} for rows.Next() { var id string var total sql.NullInt64 if err := rows.Scan(&id, &total); err != nil { return nil, err } out[id] = total.Int64 } return out, rows.Err() } // --- helpers --- // begin starts a transaction and returns it twice: bound for queries, raw for // commit and rollback. The bound half is what callers must use, so a statement // inside a transaction is rewritten exactly as one outside it is. func (s *DB) begin(ctx context.Context) (conn, *sql.Tx, error) { tx, err := s.raw.BeginTx(ctx, nil) if err != nil { return nil, nil, err } return bound{c: s.dialect, to: tx}, tx, nil } // inTx runs fn in a transaction, with the same placeholder rewriting the // outside-transaction path gets. func (s *DB) inTx(ctx context.Context, fn func(conn) error) error { tx, err := s.raw.BeginTx(ctx, nil) if err != nil { return err } defer func() { _ = tx.Rollback() }() if err := fn(bound{c: s.dialect, to: tx}); err != nil { return err } return tx.Commit() } // scanner is what both *sql.Row and *sql.Rows satisfy. type scanner interface{ Scan(...any) error } func scanKey(sc scanner) (Key, error) { var k Key var band, profile, project, user sql.NullString var expires, revoked, lastUsed sql.NullInt64 var created int64 if err := sc.Scan(&k.ID, &k.Hash, &k.Prefix, &k.Name, &project, &band, &profile, &created, &expires, &revoked, &lastUsed, &user); err != nil { return Key{}, err } k.ProjectID, k.UserID = project.String, user.String k.Band, k.RoutingProfile = band.String, profile.String k.CreatedAt = fromMillis(created) k.ExpiresAt = fromNullMillis(expires) k.RevokedAt = fromNullMillis(revoked) k.LastUsedAt = fromNullMillis(lastUsed) return k, nil } func queryAll[T any](ctx context.Context, db conn, query string, args []any, scan func(*sql.Rows) (T, error), ) ([]T, error) { rows, err := db.QueryContext(ctx, query, args...) if err != nil { return nil, err } defer rows.Close() var out []T for rows.Next() { v, err := scan(rows) if err != nil { return nil, err } out = append(out, v) } return out, rows.Err() } func encodeModels(models []string) (any, error) { if models == nil { return nil, nil } b, err := json.Marshal(models) if err != nil { return nil, err } return string(b), nil } func decodeModels(v sql.NullString) ([]string, error) { if !v.Valid || strings.TrimSpace(v.String) == "" { return nil, nil } var out []string if err := json.Unmarshal([]byte(v.String), &out); err != nil { return nil, fmt.Errorf("allowed_models is not a JSON array: %w", err) } if out == nil { out = []string{} } return out, nil } func minPtr(a, b *int) *int { switch { case b == nil: return a case a == nil: return b case *b < *a: return b default: return a } } func nullString(s string) any { if s == "" { return nil } return s } func nullInt(v *int) any { if v == nil { return nil } return *v } func nullInt64(v *int64) any { if v == nil { return nil } return *v } func millis(t time.Time) int64 { return t.UnixMilli() } func nullMillis(t time.Time) any { if t.IsZero() { return nil } return t.UnixMilli() } func fromMillis(v int64) time.Time { return time.UnixMilli(v).UTC() } func fromNullMillis(v sql.NullInt64) time.Time { if !v.Valid { return time.Time{} } return time.UnixMilli(v.Int64).UTC() } func nowMillis() int64 { return time.Now().UnixMilli() } // SQLite satisfies Store. var _ Store = (*DB)(nil) // UsageSince summarizes activity per key. // // The join brings in the key's name and prefix, which is what a dashboard needs // to show: a row identified only by an opaque id tells an operator nothing. // Revoked keys still appear, because their usage still happened. func (s *DB) UsageSince(ctx context.Context, since time.Time) ([]KeyUsage, error) { return queryAll(ctx, s.db, ` SELECT u.key_id, COALESCE(k.name, ''), COALESCE(k.key_prefix, ''), u.project_id, COALESCE(us.email, ''), COALESCE(o.name || ' / ' || t.name || ' / ' || p.name, ''), COUNT(*), SUM(u.prompt_tokens), SUM(u.completion_tokens), SUM(u.cost_micros), SUM(CASE WHEN u.status >= 400 THEN 1 ELSE 0 END), MAX(u.ts) FROM usage u LEFT JOIN api_keys k ON k.id = u.key_id -- Who, as recorded with the spend; the owner of the key for rows -- written before the person was. LEFT JOIN users us ON us.id = COALESCE(NULLIF(u.user_id, ''), k.user_id) LEFT JOIN projects p ON p.id = NULLIF(u.project_id, '') LEFT JOIN teams t ON t.id = p.team_id LEFT JOIN organizations o ON o.id = t.org_id WHERE u.ts >= ? -- Everything selected beside the sums is grouped, which PostgreSQL -- wants; each is determined by the key and project already there. GROUP BY u.key_id, u.project_id, k.name, k.key_prefix, us.email, o.name, t.name, p.name ORDER BY SUM(u.cost_micros) DESC, COUNT(*) DESC`, []any{millis(since)}, func(rows *sql.Rows) (KeyUsage, error) { var k KeyUsage var last int64 err := rows.Scan(&k.KeyID, &k.KeyName, &k.KeyPrefix, &k.ProjectID, &k.Email, &k.ProjectPath, &k.Requests, &k.Prompt, &k.Completion, &k.CostMicros, &k.Errors, &last) k.LastAt = fromMillis(last) return k, err }) } // UsageBuckets groups usage into fixed time slices. // // Empty slices are filled in, which a GROUP BY cannot do: a chart with hours // missing where nothing happened reads as a chart with hours missing from the // data, and the two mean very different things. func (s *DB) UsageBuckets(ctx context.Context, since time.Time, width time.Duration) ([]Bucket, error) { if width <= 0 { width = time.Hour } ms := width.Milliseconds() rows, err := s.db.QueryContext(ctx, ` SELECT (ts / ?) * ?, COUNT(*), SUM(CASE WHEN status >= 400 THEN 1 ELSE 0 END), SUM(prompt_tokens), SUM(completion_tokens), SUM(cost_micros), AVG(latency_ms) FROM usage WHERE ts >= ? -- By ordinal, because PostgreSQL matches a grouped expression to a -- selected one syntactically: "ts / ?" and "(ts / ?) * ?" are not the -- same text, so it saw a bare ts and refused. The ordinal is the same -- expression by construction, and both dialects accept it. GROUP BY 1 ORDER BY 1`, ms, ms, millis(since)) if err != nil { return nil, err } defer rows.Close() found := map[int64]Bucket{} for rows.Next() { var start int64 var b Bucket var latency float64 if err := rows.Scan(&start, &b.Requests, &b.Errors, &b.Prompt, &b.Completion, &b.CostMicros, &latency); err != nil { return nil, err } b.Start = fromMillis(start) b.LatencyMs = int64(latency) found[start] = b } if err := rows.Err(); err != nil { return nil, err } var out []Bucket first := (millis(since) / ms) * ms last := (millis(time.Now()) / ms) * ms for t := first; t <= last; t += ms { if b, ok := found[t]; ok { out = append(out, b) continue } out = append(out, Bucket{Start: fromMillis(t)}) } return out, nil } // UserUsage is one person's consumption, across every key they hold. type UserUsage struct { UserID string Email string Name string Keys int64 Requests int64 Prompt int64 Completion int64 CostMicros int64 Errors int64 LastSeen time.Time } // Tokens is what the person actually consumed, which is the figure they // recognise — prompt and completion together. func (u UserUsage) Tokens() int64 { return u.Prompt + u.Completion } // UsageByUser aggregates usage per person rather than per key. // // Per key is the right unit for revoking something and the wrong one for // answering "who is using this cluster": one person holds a laptop key, a CI // key and a notebook key, and reading three rows to add them up is a chore that // gets skipped. Keys with no owner — issued by an operator before accounts // existed — are grouped under an empty user, which is honest rather than tidy. func (s *DB) UsageByUser(ctx context.Context, since time.Time) ([]UserUsage, error) { rows, err := s.db.QueryContext(ctx, ` SELECT u.user_id, COALESCE(us.email, ''), COALESCE(us.name, ''), COUNT(DISTINCT u.key_id), COUNT(*), SUM(u.prompt_tokens), SUM(u.completion_tokens), SUM(u.cost_micros), SUM(CASE WHEN u.status >= 400 THEN 1 ELSE 0 END), MAX(u.ts) FROM usage u LEFT JOIN users us ON us.id = NULLIF(u.user_id, '') WHERE u.ts >= ? -- Every selected column that is not an aggregate has to be grouped. -- SQLite picks an arbitrary row for the rest and carries on; PostgreSQL -- refuses, so this query worked in tests and returned nothing but an -- error on the cluster. The email and name are functionally dependent -- on the user id, so grouping by them changes no result. GROUP BY u.user_id, COALESCE(us.email, ''), COALESCE(us.name, '') ORDER BY SUM(u.cost_micros) DESC, COUNT(*) DESC`, millis(since)) if err != nil { return nil, err } defer rows.Close() var out []UserUsage for rows.Next() { var u UserUsage var last int64 if err := rows.Scan(&u.UserID, &u.Email, &u.Name, &u.Keys, &u.Requests, &u.Prompt, &u.Completion, &u.CostMicros, &u.Errors, &last); err != nil { return nil, err } u.LastSeen = time.UnixMilli(last).UTC() out = append(out, u) } return out, rows.Err() } // ModelUsage is one model's share of the traffic. type ModelUsage struct { Model string Requests int64 Prompt int64 Completion int64 CostMicros int64 Errors int64 AvgLatency float64 } // Tokens is prompt and completion together. func (m ModelUsage) Tokens() int64 { return m.Prompt + m.Completion } // UsageByModel aggregates usage per model. // // By served_model rather than the requested name, because those differ whenever // a fallback answers — and "what actually ran" is the question being asked of a // cluster. The requested name is kept in the same table for anyone who needs to // ask the other question. func (s *DB) UsageByModel(ctx context.Context, since time.Time) ([]ModelUsage, error) { rows, err := s.db.QueryContext(ctx, ` SELECT CASE WHEN served_model != '' THEN served_model ELSE model END, COUNT(*), SUM(prompt_tokens), SUM(completion_tokens), SUM(cost_micros), SUM(CASE WHEN status >= 400 THEN 1 ELSE 0 END), AVG(latency_ms) FROM usage WHERE ts >= ? GROUP BY 1 ORDER BY COUNT(*) DESC`, millis(since)) if err != nil { return nil, err } defer rows.Close() var out []ModelUsage for rows.Next() { var m ModelUsage var avg sql.NullFloat64 if err := rows.Scan(&m.Model, &m.Requests, &m.Prompt, &m.Completion, &m.CostMicros, &m.Errors, &avg); err != nil { return nil, err } m.AvgLatency = avg.Float64 out = append(out, m) } return out, rows.Err() } // SnapshotTo writes a consistent copy of a SQLite store to path, while it is // in use. VACUUM INTO takes a read transaction, so the copy is of one moment // even with writes going on, which copying the file is not. func (s *DB) SnapshotTo(ctx context.Context, path string) error { if s.dialect != SQLiteDialect { return fmt.Errorf("a snapshot is for SQLite; back up PostgreSQL with pg_dump") } _, err := s.raw.ExecContext(ctx, `VACUUM INTO ?`, path) return err } // UsageByRequest finds the record of one request by the id it was given, or // ok=false when none has that id. func (s *DB) UsageByRequest(ctx context.Context, requestID string) (Usage, bool, error) { var u Usage var ts int64 err := s.db.QueryRowContext(ctx, `SELECT ts, key_id, project_id, COALESCE(user_id, ''), model, served_model, prompt_tokens, completion_tokens, cost_micros, status, latency_ms, request_id, replica, queue_ms, ttft_ms FROM usage WHERE request_id = ? ORDER BY ts DESC LIMIT 1`, requestID). Scan(&ts, &u.KeyID, &u.ProjectID, &u.UserID, &u.Model, &u.ServedModel, &u.PromptTokens, &u.CompletionTokens, &u.CostMicros, &u.Status, &u.LatencyMs, &u.RequestID, &u.Replica, &u.QueueMs, &u.TTFTMs) if errors.Is(err, sql.ErrNoRows) { return Usage{}, false, nil } if err != nil { return Usage{}, false, err } u.At = time.UnixMilli(ts) return u, true, nil }