llmgw/internal/store/sqlite.go

package store

import (
	"context"
	"database/sql"
	"encoding/json"
	"errors"
	"fmt"
	"slices"
	"strings"
	"time"

	_ "modernc.org/sqlite" // the cgo-free driver, so the binary stays static
)

// SQLite is the single-instance store.
type SQLite struct {
	db *sql.DB
}

// 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) (*SQLite, 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
	}
	return &SQLite{db: db}, nil
}

// Close releases the database.
func (s *SQLite) Close() error { return s.db.Close() }

// --- tenancy ---

// PutOrg upserts an organization.
func (s *SQLite) 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.
func (s *SQLite) PutTeam(ctx context.Context, t Team) error {
	_, 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.
func (s *SQLite) PutProject(ctx context.Context, p Project) error {
	_, 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 *SQLite) 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 *SQLite) 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 *SQLite) 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`

// CreateKey stores a key. Only its hash is written; the key itself is the
// caller's to show once and then forget.
func (s *SQLite) 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, k.ProjectID,
		nullString(k.Band), nullString(k.RoutingProfile),
		millis(k.CreatedAt), nullMillis(k.ExpiresAt), nullMillis(k.RevokedAt), nullMillis(k.LastUsedAt))
	return err
}

// KeyByHash finds a key by the hash of the presented secret.
func (s *SQLite) 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 *SQLite) 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 *SQLite) 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 *SQLite) TouchKeys(ctx context.Context, seen map[string]time.Time) error {
	if len(seen) == 0 {
		return nil
	}
	return s.inTx(ctx, func(tx *sql.Tx) 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 *SQLite) 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 *SQLite) 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
}

// Resolve walks a key's chain and returns the most restrictive of each limit.
//
// One query rather than four lookups: the chain is a fixed three-level join, and
// this is called on a cache miss during authentication.
func (s *SQLite) Resolve(ctx context.Context, keyID string) (Effective, error) {
	var e Effective
	row := s.db.QueryRowContext(ctx,
		`SELECT k.id, k.band, k.routing_profile, p.id, t.id, o.id
		 FROM api_keys k
		 JOIN projects p ON p.id = k.project_id
		 JOIN teams t    ON t.id = p.team_id
		 JOIN organizations o ON o.id = t.org_id
		 WHERE k.id = ?`, keyID)

	var band, profile sql.NullString
	if err := row.Scan(&e.KeyID, &band, &profile, &e.ProjectID, &e.TeamID, &e.OrgID); err != nil {
		if errors.Is(err, sql.ErrNoRows) {
			return Effective{}, ErrNotFound
		}
		return Effective{}, err
	}
	e.Band, e.RoutingProfile = band.String, profile.String

	// Widest first, so a narrower scope can only tighten what it inherits.
	chain := []struct {
		scope Scope
		id    string
	}{
		{ScopeOrg, e.OrgID},
		{ScopeTeam, e.TeamID},
		{ScopeProject, e.ProjectID},
		{ScopeKey, e.KeyID},
	}
	for _, step := range chain {
		l, err := s.LimitsFor(ctx, step.scope, step.id)
		if errors.Is(err, ErrNotFound) {
			continue
		}
		if err != nil {
			return Effective{}, err
		}
		e.tighten(l)
	}
	return e, nil
}

// tighten applies one scope's limits, keeping whichever is stricter.
func (e *Effective) tighten(l Limits) {
	e.RPM = minPtr(e.RPM, l.RPM)
	e.TPM = minPtr(e.TPM, l.TPM)
	if l.Budget != nil && (e.Budget == nil || *l.Budget < *e.Budget) {
		e.Budget = l.Budget
		e.BudgetPeriod = l.BudgetPeriod
	}
	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 *SQLite) WriteUsage(ctx context.Context, records []Usage) error {
	if len(records) == 0 {
		return nil
	}
	return s.inTx(ctx, func(tx *sql.Tx) 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)
			 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); err != nil {
				return err
			}
		}
		return nil
	})
}

// SpendSince totals cost per key since a moment.
func (s *SQLite) 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 ---

func (s *SQLite) inTx(ctx context.Context, fn func(*sql.Tx) error) error {
	tx, err := s.db.BeginTx(ctx, nil)
	if err != nil {
		return err
	}
	defer func() { _ = tx.Rollback() }()
	if err := fn(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 sql.NullString
	var expires, revoked, lastUsed sql.NullInt64
	var created int64
	if err := sc.Scan(&k.ID, &k.Hash, &k.Prefix, &k.Name, &k.ProjectID,
		&band, &profile, &created, &expires, &revoked, &lastUsed); err != nil {
		return Key{}, err
	}
	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 *sql.DB, 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 = (*SQLite)(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 *SQLite) 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,
		       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
		WHERE u.ts >= ?
		GROUP BY u.key_id, u.project_id
		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.Requests, &k.Prompt, &k.Completion, &k.CostMicros, &k.Errors, &last)
			k.LastAt = fromMillis(last)
			return k, err
		})
}