llmgw/internal/store/sqlite.go

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

	"gitea.mynodes.xyz/ilker/llmgw/internal/secrets"
)

// 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

	// sealer seals the store's own secrets; nil keeps them plain.
	sealer secrets.Sealer

	// 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")
	}
	k, err := s.keyTTL(ctx, k)
	if err != nil {
		return err
	}
	_, 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) {
	src, args := usageSince(since)
	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 `+src+` u
		 LEFT JOIN projects p ON p.id = NULLIF(u.project_id, '')
		 LEFT JOIN teams t    ON t.id = p.team_id
		 GROUP BY u.key_id, u.user_id, p.id, t.id, t.org_id`, args...)
	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,
			   guard_action, guard_policy, guard_category, guard_ms, end_user, tags, cache_hit)
			 VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`)
		if err != nil {
			return err
		}
		defer stmt.Close()
		for _, r := range records {
			hit := 0
			if r.CacheHit {
				hit = 1
			}
			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,
				r.GuardAction, r.GuardPolicy, r.GuardCategory, r.GuardMs, r.EndUser, tagsJSON(r.Tags), hit); err != nil {
				return err
			}
		}
		return writeRollup(ctx, tx, records)
	})
}

// SpendSince totals cost per key since a moment.
func (s *DB) SpendSince(ctx context.Context, since time.Time) (map[string]int64, error) {
	src, args := usageSince(since)
	rows, err := s.db.QueryContext(ctx,
		`SELECT u.key_id, SUM(u.cost_micros) FROM `+src+` u GROUP BY u.key_id`, args...)
	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) {
	src, args := usageSince(since)
	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, ''),
		       SUM(u.requests), SUM(u.prompt_tokens), SUM(u.completion_tokens),
		       SUM(u.cost_micros), SUM(u.errors), MAX(u.last_ts)
		FROM `+src+` 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
		-- 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, SUM(u.requests) DESC`,
		args, 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()

	// From the rollup when the slices are whole rollup slices, which every
	// page asks for; from the rows otherwise.
	q, from := `
		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`, millis(since)
	if t, ok := tierForSlices(width); ok {
		q, from = `
		SELECT (bucket / ?) * ?, SUM(requests), SUM(errors),
		       SUM(prompt_tokens), SUM(completion_tokens), SUM(cost_micros),
		       1.0 * SUM(latency_ms_sum) / SUM(requests)
		FROM `+t.table+` WHERE bucket >= ?
		GROUP BY 1 ORDER BY 1`, t.floor(since)
	}
	rows, err := s.db.QueryContext(ctx, q, ms, ms, from)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	found := map[int64]Bucket{}
	for rows.Next() {
		var start int64
		var b Bucket
		var latency sql.NullFloat64
		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.Float64)
		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) {
	src, args := usageSince(since)
	rows, err := s.db.QueryContext(ctx, `
		SELECT u.user_id, COALESCE(us.email, ''), COALESCE(us.name, ''),
		       COUNT(DISTINCT u.key_id),
		       SUM(u.requests), SUM(u.prompt_tokens), SUM(u.completion_tokens),
		       SUM(u.cost_micros), SUM(u.errors), MAX(u.last_ts)
		FROM `+src+` u
		LEFT JOIN users us ON us.id = NULLIF(u.user_id, '')
		-- 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, SUM(u.requests) DESC`, args...)
	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) {
	src, args := usageSince(since)
	rows, err := s.db.QueryContext(ctx, `
		SELECT CASE WHEN served_model != '' THEN served_model ELSE model END,
		       SUM(requests), SUM(prompt_tokens), SUM(completion_tokens),
		       SUM(cost_micros), SUM(errors),
		       1.0 * SUM(latency_ms_sum) / SUM(requests)
		FROM `+src+` u
		GROUP BY 1 ORDER BY SUM(requests) DESC`, args...)
	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
	var tags string
	var hit int
	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,
		        guard_action, guard_policy, guard_category, guard_ms, end_user, tags, cache_hit
		 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, &u.GuardAction, &u.GuardPolicy, &u.GuardCategory, &u.GuardMs,
			&u.EndUser, &tags, &hit)
	if errors.Is(err, sql.ErrNoRows) {
		return Usage{}, false, nil
	}
	if err != nil {
		return Usage{}, false, err
	}
	u.At = time.UnixMilli(ts)
	u.Tags = parseTags(tags)
	u.CacheHit = hit == 1
	return u, true, nil
}