diff --git a/internal/account/account_test.go b/internal/account/account_test.go index eea84b6..7256dfa 100644 --- a/internal/account/account_test.go +++ b/internal/account/account_test.go @@ -12,7 +12,7 @@ import ( "gitea.mynodes.xyz/ilker/llmgw/internal/store" ) -func newStore(t *testing.T) *store.SQLite { +func newStore(t *testing.T) *store.DB { t.Helper() s, err := store.OpenSQLite(context.Background(), filepath.Join(t.TempDir(), "a.db")) if err != nil { @@ -33,7 +33,7 @@ func newStore(t *testing.T) *store.SQLite { return s } -func newService(t *testing.T) (*Service, *store.SQLite) { +func newService(t *testing.T) (*Service, *store.DB) { t.Helper() s := newStore(t) return New(s, DefaultConfig(), slog.New(slog.DiscardHandler)), s @@ -113,7 +113,7 @@ func TestApprovalNeedsAProject(t *testing.T) { } } -func newStoreAndUser(t *testing.T) (*Service, *store.SQLite) { +func newStoreAndUser(t *testing.T) (*Service, *store.DB) { t.Helper() svc, st := newService(t) hash, _ := HashPassword(goodPassword) diff --git a/internal/auth/auth_test.go b/internal/auth/auth_test.go index b1f4d9e..65e3f6f 100644 --- a/internal/auth/auth_test.go +++ b/internal/auth/auth_test.go @@ -16,7 +16,7 @@ import ( ) // fixture builds a store with one org/team/project and one key. -func fixture(t *testing.T) (*store.SQLite, string, string) { +func fixture(t *testing.T) (*store.DB, string, string) { t.Helper() ctx := context.Background() s, err := store.OpenSQLite(ctx, filepath.Join(t.TempDir(), "t.db")) diff --git a/internal/cli/keys.go b/internal/cli/keys.go index 5e702df..42eb7ba 100644 --- a/internal/cli/keys.go +++ b/internal/cli/keys.go @@ -16,7 +16,7 @@ import ( ) func newTenancyCmd() *cobra.Command { - var dbPath string + var sf storeFlags cmd := &cobra.Command{ Use: "keys", @@ -25,9 +25,11 @@ func newTenancyCmd() *cobra.Command { "can be issued before the server is running and without exposing an endpoint\n" + "that mints credentials.", } - cmd.PersistentFlags().StringVar(&dbPath, "store", "/var/lib/llmgw/llmgw.db", "path to the store") + sf.bind(cmd) - open := func(ctx context.Context) (*store.SQLite, error) { return store.Open(ctx, dbPath) } + // Takes a context rather than the command, because some of these run + // before the command tree is entered. + open := func(ctx context.Context) (*store.DB, error) { return sf.openCtx(ctx) } cmd.AddCommand( newKeyInitCmd(open), @@ -42,7 +44,7 @@ func newTenancyCmd() *cobra.Command { return cmd } -type opener func(context.Context) (*store.SQLite, error) +type opener func(context.Context) (*store.DB, error) func newKeyInitCmd(open opener) *cobra.Command { var org, team, project string diff --git a/internal/cli/server.go b/internal/cli/server.go index 3f4c36c..afdf7cd 100644 --- a/internal/cli/server.go +++ b/internal/cli/server.go @@ -362,7 +362,7 @@ func startControlPlane(ctx context.Context, cfg config.Server, cat *catalog.Live // the sink is the router when the router is serving, and nothing when a // control plane runs on its own. var lead *control.Leadership - if h, ok := sink.(interface{ Store() *storepkg.SQLite }); ok { + if h, ok := sink.(interface{ Store() *storepkg.DB }); ok { if st := h.Store(); st != nil { lead = control.NewLeadership(st, cfg.Control.LeaseTTL, log) ctl.SetLeadership(lead) diff --git a/internal/cli/storeflag.go b/internal/cli/storeflag.go new file mode 100644 index 0000000..e806dd6 --- /dev/null +++ b/internal/cli/storeflag.go @@ -0,0 +1,76 @@ +package cli + +import ( + "context" + "fmt" + + "github.com/spf13/cobra" + + "gitea.mynodes.xyz/ilker/llmgw/internal/config" + "gitea.mynodes.xyz/ilker/llmgw/internal/store" +) + +// defaultStorePath is where a single-machine deployment keeps its store. +const defaultStorePath = "/var/lib/llmgw/llmgw.db" + +// defaultConfigPath is where the control plane's configuration usually lives. +// +// Consulted rather than required. The commands that work on the store run on +// the same machine as the server far more often than not, and the server +// already knows where the store is — asking it is better than making the +// operator repeat an answer that is written down three lines away. +const defaultConfigPath = "/etc/llmgw/server.yaml" + +// storeFlags is how a command that works directly on the store finds it. +type storeFlags struct { + address string + config string +} + +// bind adds --store and --config to a command. +func (f *storeFlags) bind(cmd *cobra.Command) { + cmd.PersistentFlags().StringVar(&f.address, "store", "", + "store address: a SQLite path, a postgres:// DSN, or @file holding one "+ + "(default: auth.store from --config, else "+defaultStorePath+")") + cmd.PersistentFlags().StringVar(&f.config, "config", defaultConfigPath, + "server configuration to read the store address from") +} + +// open resolves the store and opens it. +// +// Explicit --store wins; otherwise the server's own configuration is asked, +// and only then does the single-machine default apply. That order exists +// because the default was a SQLite path and nothing else: on a cluster whose +// store is PostgreSQL, every one of these commands failed with "unable to open +// database file" while the right answer sat in the config file beside it. The +// operator's conclusion was that the tool was broken, and they were not far off. +func (f *storeFlags) open(cmd *cobra.Command) (*store.DB, error) { + return f.openCtx(cmd.Context()) +} + +// openCtx is open for a caller that has a context rather than a command. +func (f *storeFlags) openCtx(ctx context.Context) (*store.DB, error) { + address, from := f.resolve() + s, err := store.Open(ctx, address) + if err != nil { + return nil, fmt.Errorf("opening the store at %s (%s): %w", address, from, err) + } + return s, nil +} + +// resolve returns the address and where it came from, for the error message. +// +// Where it came from is half the message. "unable to open database file" says +// nothing about which of three sources chose that path, and the whole problem +// is that the wrong one did. +func (f *storeFlags) resolve() (address, from string) { + if f.address != "" { + return f.address, "--store" + } + if f.config != "" { + if cfg, err := config.Load(f.config); err == nil && cfg.Auth.Store != "" { + return cfg.Auth.Store, "auth.store in " + f.config + } + } + return defaultStorePath, "the default; no --store and nothing usable in " + f.config +} diff --git a/internal/cli/storeflag_test.go b/internal/cli/storeflag_test.go new file mode 100644 index 0000000..9ac0738 --- /dev/null +++ b/internal/cli/storeflag_test.go @@ -0,0 +1,82 @@ +package cli + +import ( + "os" + "path/filepath" + "strings" + "testing" +) + +func writeConfig(t *testing.T, body string) string { + t.Helper() + p := filepath.Join(t.TempDir(), "server.yaml") + if err := os.WriteFile(p, []byte(body), 0o600); err != nil { + t.Fatal(err) + } + return p +} + +func TestAnExplicitStoreWins(t *testing.T) { + cfg := writeConfig(t, "listen: \":8080\"\nauth:\n enabled: true\n store: postgres://from-config/db\n") + f := storeFlags{address: "postgres://explicit/db", config: cfg} + + got, from := f.resolve() + if got != "postgres://explicit/db" || !strings.Contains(from, "--store") { + t.Errorf("an explicit --store should win; got %q from %q", got, from) + } +} + +func TestTheServerConfigIsAskedBeforeTheDefault(t *testing.T) { + cfg := writeConfig(t, "listen: \":8080\"\nauth:\n enabled: true\n store: postgres://from-config/db\n") + f := storeFlags{config: cfg} + + got, from := f.resolve() + // The whole defect: on a cluster whose store is PostgreSQL, these commands + // reached for a SQLite path and failed with "unable to open database file" + // while the right answer sat in the config beside them. + if got != "postgres://from-config/db" { + t.Errorf("expected the address from the config, got %q", got) + } + if !strings.Contains(from, "auth.store") { + t.Errorf("the source should be named for the error message; got %q", from) + } +} + +func TestTheSingleMachineDefaultStillApplies(t *testing.T) { + f := storeFlags{config: filepath.Join(t.TempDir(), "absent.yaml")} + + got, from := f.resolve() + // A deployment with no config file is the single-machine one, and it worked + // before this existed. It has to keep working with no flags at all. + if got != defaultStorePath { + t.Errorf("expected the default path, got %q", got) + } + if !strings.Contains(from, "default") { + t.Errorf("the source should say it fell back; got %q", from) + } +} + +func TestAConfigWithNoStoreFallsThrough(t *testing.T) { + // Authentication off means no store at all. Reading an empty value as the + // answer would open a database named "". + cfg := writeConfig(t, "listen: \":8080\"\n") + f := storeFlags{config: cfg} + + if got, _ := f.resolve(); got != defaultStorePath { + t.Errorf("expected the default when the config names no store, got %q", got) + } +} + +func TestTheErrorSaysWhichSourceChoseThePath(t *testing.T) { + f := storeFlags{address: filepath.Join(t.TempDir(), "nope", "x.db")} + + _, err := f.openCtx(t.Context()) + if err == nil { + t.Fatal("expected opening a store in a missing directory to fail") + } + // "unable to open database file" says nothing about which of three sources + // picked that path, and the whole problem is that the wrong one did. + if !strings.Contains(err.Error(), "--store") { + t.Errorf("the error should name where the address came from; got %q", err) + } +} diff --git a/internal/cli/users.go b/internal/cli/users.go index 370e209..e4f17ae 100644 --- a/internal/cli/users.go +++ b/internal/cli/users.go @@ -12,7 +12,7 @@ import ( ) func newUsersCmd() *cobra.Command { - var dbPath string + var sf storeFlags cmd := &cobra.Command{ Use: "users", @@ -22,9 +22,9 @@ func newUsersCmd() *cobra.Command { "the step that decides who gets to spend the cluster, and it needs a shell\n" + "on this machine rather than a session cookie.", } - cmd.PersistentFlags().StringVar(&dbPath, "store", "/var/lib/llmgw/llmgw.db", "path to the store") - open := func(cmd *cobra.Command) (*store.SQLite, error) { - return store.Open(cmd.Context(), dbPath) + sf.bind(cmd) + open := func(cmd *cobra.Command) (*store.DB, error) { + return sf.open(cmd) } cmd.AddCommand(newUserListCmd(open), newUserApproveCmd(open), @@ -38,7 +38,7 @@ func newUsersCmd() *cobra.Command { return cmd } -type userOpener func(*cobra.Command) (*store.SQLite, error) +type userOpener func(*cobra.Command) (*store.DB, error) func newUserListCmd(open userOpener) *cobra.Command { var pending bool diff --git a/internal/router/router.go b/internal/router/router.go index ee928d8..38ec0d1 100644 --- a/internal/router/router.go +++ b/internal/router/router.go @@ -74,7 +74,7 @@ type Router struct { // disabled. auth *auth.Authenticator accounts *account.Service - store *storepkg.SQLite + store *storepkg.DB acct *accounting.Accountant // threads holds conversations and memory notes. It is store when auth is on, @@ -82,7 +82,7 @@ type Router struct { // nothing is kept. Separate from store because a non-nil store also switches // on the accounts, web pages and usage recording, none of which belong to a // gateway that only wants to remember things. - threads *storepkg.SQLite + threads *storepkg.DB // anonymous means threads has no project to scope to: everything shares the // one pool named by anonymousProject. anonymous bool @@ -124,7 +124,7 @@ func (r *Router) SetMirror(m web.MirrorAdmin) { r.mirror = m } func (r *Router) Auth() *auth.Authenticator { return r.auth } // Store exposes the key store, or nil when authentication is off. -func (r *Router) Store() *storepkg.SQLite { return r.store } +func (r *Router) Store() *storepkg.DB { return r.store } // New builds a router from configuration. Start must be called to begin scraping. func New(cfg config.Server, log *slog.Logger) (*Router, error) { diff --git a/internal/store/capacity.go b/internal/store/capacity.go index 1cd57bc..e0a3688 100644 --- a/internal/store/capacity.go +++ b/internal/store/capacity.go @@ -27,7 +27,7 @@ type StoredCapacity struct { // point: hardware and engine builds change, and keeping every run would leave a // reader choosing between numbers with no basis for the choice. The run itself // is printed when it happens, for anyone who wants the ramp. -func (s *SQLite) PutCapacity(ctx context.Context, c StoredCapacity) error { +func (s *DB) PutCapacity(ctx context.Context, c StoredCapacity) error { if c.Model == "" || c.Profile == "" { return fmt.Errorf("a capacity measurement needs a model and a profile") } @@ -53,7 +53,7 @@ func (s *SQLite) PutCapacity(ctx context.Context, c StoredCapacity) error { // Newest first because the question a reader usually has is "what do we know // lately", and an old measurement is the one most likely to be describing // hardware that has since changed. -func (s *SQLite) Capacities(ctx context.Context) ([]StoredCapacity, error) { +func (s *DB) Capacities(ctx context.Context) ([]StoredCapacity, error) { rows, err := s.db.QueryContext(ctx, `SELECT model, profile, measurement, measured_at, measured_by FROM capacity ORDER BY measured_at DESC, model, profile`) @@ -81,7 +81,7 @@ func (s *SQLite) Capacities(ctx context.Context) ([]StoredCapacity, error) { // hardware has been replaced is worse than no measurement, because it looks // like knowledge. Deleting one that is not there is not an error — the caller // wanted it gone, and it is. -func (s *SQLite) DeleteCapacity(ctx context.Context, model, profile string) error { +func (s *DB) DeleteCapacity(ctx context.Context, model, profile string) error { _, err := s.db.ExecContext(ctx, `DELETE FROM capacity WHERE model = ? AND profile = ?`, model, profile) return err diff --git a/internal/store/capacity_test.go b/internal/store/capacity_test.go index 0016b4c..fa3c0ca 100644 --- a/internal/store/capacity_test.go +++ b/internal/store/capacity_test.go @@ -7,7 +7,7 @@ import ( "time" ) -func capacityStore(t *testing.T) *SQLite { +func capacityStore(t *testing.T) *DB { t.Helper() s, err := OpenSQLite(context.Background(), filepath.Join(t.TempDir(), "cap.db")) if err != nil { diff --git a/internal/store/claims.go b/internal/store/claims.go index 79a092b..ad952c8 100644 --- a/internal/store/claims.go +++ b/internal/store/claims.go @@ -41,7 +41,7 @@ func NewClaimToken() (string, error) { // // The secret is sealed with the token, which is not stored, so what lands in // the database cannot be opened by anything that can read the database. -func (s *SQLite) CreateClaim(ctx context.Context, keyID, token, secret string, ttl time.Duration) error { +func (s *DB) CreateClaim(ctx context.Context, keyID, token, secret string, ttl time.Duration) error { sealed, err := seal(token, []byte(secret)) if err != nil { return err @@ -62,7 +62,7 @@ func (s *SQLite) CreateClaim(ctx context.Context, keyID, token, secret string, t // The page behind the link is shown before it is used, because link previewers // follow URLs: a GET that consumed the claim would mean the first chat client // to unfurl it burns the key. -func (s *SQLite) PeekClaim(ctx context.Context, token string) (Claim, error) { +func (s *DB) PeekClaim(ctx context.Context, token string) (Claim, error) { row := s.db.QueryRowContext(ctx, `SELECT key_id, expires_at, claimed_at FROM key_claims WHERE token_hash = ?`, HashKey(token)) @@ -88,7 +88,7 @@ func (s *SQLite) PeekClaim(ctx context.Context, token string) (Claim, error) { // The update is the lock: exactly one caller can move claimed_at from NULL, so // two people opening the link at the same moment produce one winner rather than // two copies of the key. -func (s *SQLite) ConsumeClaim(ctx context.Context, token string) (string, error) { +func (s *DB) ConsumeClaim(ctx context.Context, token string) (string, error) { hash := HashKey(token) now := time.Now() @@ -140,7 +140,7 @@ func (s *SQLite) ConsumeClaim(ctx context.Context, token string) (string, error) // Housekeeping rather than security — an expired row is already refused and its // contents are already unopenable — but a table that only grows is its own // problem. -func (s *SQLite) PurgeClaims(ctx context.Context, before time.Time) (int64, error) { +func (s *DB) PurgeClaims(ctx context.Context, before time.Time) (int64, error) { res, err := s.db.ExecContext(ctx, `DELETE FROM key_claims WHERE expires_at < ? OR claimed_at IS NOT NULL`, millis(before)) if err != nil { diff --git a/internal/store/claims_test.go b/internal/store/claims_test.go index b9b64b9..2af5c1e 100644 --- a/internal/store/claims_test.go +++ b/internal/store/claims_test.go @@ -9,7 +9,7 @@ import ( "time" ) -func claimFixture(t *testing.T) (*SQLite, string, string) { +func claimFixture(t *testing.T) (*DB, string, string) { t.Helper() ctx := context.Background() s, err := OpenSQLite(ctx, filepath.Join(t.TempDir(), "c.db")) diff --git a/internal/store/copy.go b/internal/store/copy.go index d6d1c01..367ded1 100644 --- a/internal/store/copy.go +++ b/internal/store/copy.go @@ -70,7 +70,7 @@ type CopyResult struct { // a migration runs out of memory on the machine least able to spare it. A // failure therefore leaves a partial copy, which is why the destination has to // have been empty: starting again is dropping and recreating it. -func (s *SQLite) CopyTo(ctx context.Context, dst *SQLite) ([]CopyResult, error) { +func (s *DB) CopyTo(ctx context.Context, dst *DB) ([]CopyResult, error) { empty, where, err := dst.isEmpty(ctx) if err != nil { return nil, err @@ -96,7 +96,7 @@ func (s *SQLite) CopyTo(ctx context.Context, dst *SQLite) ([]CopyResult, error) } // isEmpty reports whether any table has rows, and which one first does. -func (s *SQLite) isEmpty(ctx context.Context) (bool, string, error) { +func (s *DB) isEmpty(ctx context.Context) (bool, string, error) { for _, table := range copyOrder { var n int64 //nolint:gosec // table comes from copyOrder, not from input @@ -116,7 +116,7 @@ func (s *SQLite) isEmpty(ctx context.Context) (bool, string, error) { // migration adds one: a hand-written column list would copy the old shape and // silently drop the new field, and nothing would notice until someone looked // for a setting that was there before the move. -func (s *SQLite) copyTable(ctx context.Context, dst *SQLite, table string) (int64, error) { +func (s *DB) copyTable(ctx context.Context, dst *DB, table string) (int64, error) { //nolint:gosec // table comes from copyOrder, not from input rows, err := s.db.QueryContext(ctx, `SELECT * FROM `+table) if err != nil { @@ -196,7 +196,7 @@ func quoteAll(cols []string) []string { // copied in: the sequence still starts at 1 while the table already holds a // thousand notes, and the failure arrives later, to somebody else, looking // like a bug in the feature rather than in the move. -func (s *SQLite) resetIdentities(ctx context.Context) error { +func (s *DB) resetIdentities(ctx context.Context) error { if s.dialect != PostgresDialect { return nil } diff --git a/internal/store/leadership.go b/internal/store/leadership.go index 1b999c9..5331cec 100644 --- a/internal/store/leadership.go +++ b/internal/store/leadership.go @@ -30,7 +30,7 @@ type Lease struct { // spelled differently in each dialect. Two control planes on machines whose // clocks disagree by more than the lease will fight, so keep them on NTP — // which any cluster running InfiniBand already has to be. -func (s *SQLite) AcquireLeadership(ctx context.Context, role, owner string, ttl time.Duration) (Lease, error) { +func (s *DB) AcquireLeadership(ctx context.Context, role, owner string, ttl time.Duration) (Lease, error) { now := time.Now() expires := now.Add(ttl) @@ -67,7 +67,7 @@ func (s *SQLite) AcquireLeadership(ctx context.Context, role, owner string, ttl } // Leadership reports who holds a role's lease. -func (s *SQLite) Leadership(ctx context.Context, role string) (Lease, error) { +func (s *DB) Leadership(ctx context.Context, role string) (Lease, error) { var owner string var expires, acquired int64 err := s.db.QueryRowContext(ctx, @@ -89,7 +89,7 @@ func (s *SQLite) Leadership(ctx context.Context, role string) (Lease, error) { // Only the holder can release it: a process that has already lost the lease // must not be able to take it from whoever has it now, which is exactly what // an unconditional delete during a slow shutdown would do. -func (s *SQLite) ReleaseLeadership(ctx context.Context, role, owner string) error { +func (s *DB) ReleaseLeadership(ctx context.Context, role, owner string) error { _, err := s.db.ExecContext(ctx, `DELETE FROM leadership WHERE role = ? AND owner = ?`, role, owner) return err diff --git a/internal/store/leadership_test.go b/internal/store/leadership_test.go index d12746d..7c7b9ea 100644 --- a/internal/store/leadership_test.go +++ b/internal/store/leadership_test.go @@ -11,7 +11,7 @@ import ( // leadershipSuite is the same assertions against whichever store it is given, // because the whole point of the lease is that it behaves identically on both // — and the one that matters for HA is the one it was never tested on. -func leadershipSuite(t *testing.T, open func(*testing.T) *SQLite) { +func leadershipSuite(t *testing.T, open func(*testing.T) *DB) { ctx := context.Background() t.Run("the first to ask gets it", func(t *testing.T) { @@ -143,7 +143,7 @@ func leadershipSuite(t *testing.T, open func(*testing.T) *SQLite) { } func TestLeadershipOnSQLite(t *testing.T) { - leadershipSuite(t, func(t *testing.T) *SQLite { return openDB(t) }) + leadershipSuite(t, func(t *testing.T) *DB { return openDB(t) }) } func TestLeadershipOnPostgres(t *testing.T) { diff --git a/internal/store/limits_owner_test.go b/internal/store/limits_owner_test.go index e14a711..7c97fdc 100644 --- a/internal/store/limits_owner_test.go +++ b/internal/store/limits_owner_test.go @@ -7,7 +7,7 @@ import ( "time" ) -func treeStore(t *testing.T) *SQLite { +func treeStore(t *testing.T) *DB { t.Helper() ctx := context.Background() s, err := OpenSQLite(ctx, filepath.Join(t.TempDir(), "t.db")) diff --git a/internal/store/memory.go b/internal/store/memory.go index d6b47c0..e9d20b5 100644 --- a/internal/store/memory.go +++ b/internal/store/memory.go @@ -84,7 +84,7 @@ func decodeTags(s string) []string { } // MemoryRetention reports how long a project keeps notes. Zero means none. -func (s *SQLite) MemoryRetention(ctx context.Context, projectID string) (time.Duration, error) { +func (s *DB) MemoryRetention(ctx context.Context, projectID string) (time.Duration, error) { var hours sql.NullInt64 err := s.db.QueryRowContext(ctx, `SELECT memory_retention_hours FROM projects WHERE id = ?`, projectID).Scan(&hours) @@ -101,7 +101,7 @@ func (s *SQLite) MemoryRetention(ctx context.Context, projectID string) (time.Du } // SetMemoryRetention turns note storage on for a project, or off. -func (s *SQLite) SetMemoryRetention(ctx context.Context, projectID string, hours int) error { +func (s *DB) SetMemoryRetention(ctx context.Context, projectID string, hours int) error { var v any if hours > 0 { v = hours @@ -119,7 +119,7 @@ func (s *SQLite) SetMemoryRetention(ctx context.Context, projectID string, hours // PutNote writes a note, replacing any note the agent already holds under the // same key. -func (s *SQLite) PutNote(ctx context.Context, projectID, agentID string, n Note) error { +func (s *DB) PutNote(ctx context.Context, projectID, agentID string, n Note) error { if n.Key == "" || len(n.Key) > MaxNoteKey { return fmt.Errorf("key must be 1 to %d bytes", MaxNoteKey) } @@ -191,7 +191,7 @@ func scanNote(rows interface{ Scan(...any) error }) (Note, error) { const noteColumns = `key, text, tags, created_at, updated_at` // GetNote reads one note. Scoped to the project and agent in the query itself. -func (s *SQLite) GetNote(ctx context.Context, projectID, agentID, key string) (Note, error) { +func (s *DB) GetNote(ctx context.Context, projectID, agentID, key string) (Note, error) { n, err := scanNote(s.db.QueryRowContext(ctx, `SELECT `+noteColumns+` FROM memory_notes WHERE project_id = ? AND agent_id = ? AND key = ?`, projectID, agentID, key)) @@ -215,7 +215,7 @@ func ftsQuery(q string) string { // SearchNotes finds an agent's notes by words and/or tag, newest first. With // neither it lists the most recent. -func (s *SQLite) SearchNotes(ctx context.Context, projectID, agentID, query, tag string, limit int) ([]Note, error) { +func (s *DB) SearchNotes(ctx context.Context, projectID, agentID, query, tag string, limit int) ([]Note, error) { if limit <= 0 || limit > MaxNoteResults { limit = MaxNoteResults } @@ -253,7 +253,7 @@ func (s *SQLite) SearchNotes(ctx context.Context, projectID, agentID, query, tag } // DeleteNote forgets one note. -func (s *SQLite) DeleteNote(ctx context.Context, projectID, agentID, key string) error { +func (s *DB) DeleteNote(ctx context.Context, projectID, agentID, key string) error { res, err := s.db.ExecContext(ctx, `DELETE FROM memory_notes WHERE project_id = ? AND agent_id = ? AND key = ?`, projectID, agentID, key) @@ -267,7 +267,7 @@ func (s *SQLite) DeleteNote(ctx context.Context, projectID, agentID, key string) } // PruneMemory deletes notes not refreshed within their project's retention. -func (s *SQLite) PruneMemory(ctx context.Context, now time.Time) (int, error) { +func (s *DB) PruneMemory(ctx context.Context, now time.Time) (int, error) { type window struct { project string hours int64 diff --git a/internal/store/memory_test.go b/internal/store/memory_test.go index fd52065..7f561c1 100644 --- a/internal/store/memory_test.go +++ b/internal/store/memory_test.go @@ -9,7 +9,7 @@ import ( "time" ) -func memoryProject(t *testing.T, s *SQLite) string { +func memoryProject(t *testing.T, s *DB) string { t.Helper() proj := tree(t, s) if err := s.SetMemoryRetention(context.Background(), proj, 24); err != nil { diff --git a/internal/store/models.go b/internal/store/models.go index 9c88d4d..19e0bf4 100644 --- a/internal/store/models.go +++ b/internal/store/models.go @@ -19,7 +19,7 @@ type StoredModel struct { } // PutModel adds or replaces a runtime model. -func (s *SQLite) PutModel(ctx context.Context, m StoredModel) error { +func (s *DB) PutModel(ctx context.Context, m StoredModel) error { if m.Name == "" { return fmt.Errorf("a model needs a name") } @@ -39,14 +39,14 @@ func (s *SQLite) PutModel(ctx context.Context, m StoredModel) error { // DeleteModel removes a runtime model. Removing one that is not there is not an // error: the caller wanted it gone, and it is. -func (s *SQLite) DeleteModel(ctx context.Context, name string) error { +func (s *DB) DeleteModel(ctx context.Context, name string) error { _, err := s.db.ExecContext(ctx, `DELETE FROM models WHERE name = ?`, name) return err } // Models lists the runtime models, oldest first so the catalog is built in a // stable order and two starts produce the same result. -func (s *SQLite) Models(ctx context.Context) ([]StoredModel, error) { +func (s *DB) Models(ctx context.Context) ([]StoredModel, error) { rows, err := s.db.QueryContext(ctx, `SELECT name, definition, created_at, created_by FROM models ORDER BY created_at, name`) if err != nil { diff --git a/internal/store/open.go b/internal/store/open.go index cb4132f..2637014 100644 --- a/internal/store/open.go +++ b/internal/store/open.go @@ -18,7 +18,7 @@ import ( // An address of the form @/path/to/file reads the real address out of that // file, which is how a DSN carrying a password stays out of a configuration // file that gets copied, diffed and pasted into issues. -func Open(ctx context.Context, address string) (*SQLite, error) { +func Open(ctx context.Context, address string) (*DB, error) { address, err := resolveAddress(address) if err != nil { return nil, err diff --git a/internal/store/postgres.go b/internal/store/postgres.go index 934f620..7c1ef32 100644 --- a/internal/store/postgres.go +++ b/internal/store/postgres.go @@ -25,7 +25,7 @@ import ( // somewhere protected rather than as a path this package knows how to find — // and it is never logged, which is why errors below name the database and not // the DSN. -func OpenPostgres(ctx context.Context, dsn string) (*SQLite, error) { +func OpenPostgres(ctx context.Context, dsn string) (*DB, error) { if dsn == "" { return nil, fmt.Errorf("store: no PostgreSQL address configured") } @@ -55,7 +55,7 @@ func OpenPostgres(ctx context.Context, dsn string) (*SQLite, error) { _ = db.Close() return nil, err } - return &SQLite{ + return &DB{ raw: db, dialect: PostgresDialect, db: bound{c: PostgresDialect, to: db}, diff --git a/internal/store/postgres_test.go b/internal/store/postgres_test.go index 57465f7..bee55f8 100644 --- a/internal/store/postgres_test.go +++ b/internal/store/postgres_test.go @@ -20,7 +20,7 @@ import ( // Every table is emptied first. The alternative is a test database that // accumulates rows until an assertion about counts starts failing for reasons // that have nothing to do with the change being tested. -func openPG(t *testing.T) *SQLite { +func openPG(t *testing.T) *DB { t.Helper() dsn := os.Getenv("LLMGW_TEST_POSTGRES") if dsn == "" { diff --git a/internal/store/sessions.go b/internal/store/sessions.go index c664310..6d45115 100644 --- a/internal/store/sessions.go +++ b/internal/store/sessions.go @@ -32,7 +32,7 @@ func NewSessionToken() (string, error) { // CreateSession opens a session and returns nothing: the caller already holds // the tokens, and neither is stored in the clear. -func (s *SQLite) CreateSession(ctx context.Context, userID, token, csrf string, ttl time.Duration) error { +func (s *DB) CreateSession(ctx context.Context, userID, token, csrf string, ttl time.Duration) error { if ttl <= 0 { ttl = SessionTTL } @@ -49,7 +49,7 @@ func (s *SQLite) CreateSession(ctx context.Context, userID, token, csrf string, // One query rather than two: a session whose user has been suspended must stop // working immediately, and checking that in a second round trip is a check // somebody eventually forgets. -func (s *SQLite) SessionByToken(ctx context.Context, token string) (Session, User, error) { +func (s *DB) SessionByToken(ctx context.Context, token string) (Session, User, error) { row := s.db.QueryRowContext(ctx, `SELECT s.user_id, s.csrf_hash, s.expires_at, u.id, u.email, u.name, u.password_hash, u.state, COALESCE(u.project_id, ''), @@ -94,7 +94,7 @@ func (s *SQLite) SessionByToken(ctx context.Context, token string) (Session, Use } // TouchSession records activity, at most once a minute. -func (s *SQLite) TouchSession(ctx context.Context, token string, at time.Time) error { +func (s *DB) TouchSession(ctx context.Context, token string, at time.Time) error { _, err := s.db.ExecContext(ctx, `UPDATE sessions SET last_seen = ? WHERE token_hash = ? AND last_seen < ?`, millis(at), HashKey(token), millis(at.Add(-time.Minute))) @@ -102,7 +102,7 @@ func (s *SQLite) TouchSession(ctx context.Context, token string, at time.Time) e } // DeleteSession signs one browser out. -func (s *SQLite) DeleteSession(ctx context.Context, token string) error { +func (s *DB) DeleteSession(ctx context.Context, token string) error { _, err := s.db.ExecContext(ctx, `DELETE FROM sessions WHERE token_hash = ?`, HashKey(token)) return err } @@ -111,13 +111,13 @@ func (s *SQLite) DeleteSession(ctx context.Context, token string) error { // // What suspension and a password change both need: leaving old sessions open // would make either of them advisory. -func (s *SQLite) DeleteUserSessions(ctx context.Context, userID string) error { +func (s *DB) DeleteUserSessions(ctx context.Context, userID string) error { _, err := s.db.ExecContext(ctx, `DELETE FROM sessions WHERE user_id = ?`, userID) return err } // PurgeSessions removes expired rows. -func (s *SQLite) PurgeSessions(ctx context.Context, before time.Time) (int64, error) { +func (s *DB) PurgeSessions(ctx context.Context, before time.Time) (int64, error) { res, err := s.db.ExecContext(ctx, `DELETE FROM sessions WHERE expires_at < ?`, millis(before)) if err != nil { return 0, err diff --git a/internal/store/sqlite.go b/internal/store/sqlite.go index ed2cf49..96c8057 100644 --- a/internal/store/sqlite.go +++ b/internal/store/sqlite.go @@ -15,8 +15,17 @@ import ( _ "modernc.org/sqlite" // the cgo-free driver, so the binary stays static ) -// SQLite is the single-instance store. -type SQLite struct { +// 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 @@ -41,7 +50,7 @@ type SQLite struct { // 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) { +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 { @@ -70,16 +79,16 @@ func OpenSQLite(ctx context.Context, path string) (*SQLite, error) { return nil, fmt.Errorf("securing %s: %w", f, err) } } - return &SQLite{raw: db, dialect: SQLiteDialect, db: bound{c: SQLiteDialect, to: db}}, nil + return &DB{raw: db, dialect: SQLiteDialect, db: bound{c: SQLiteDialect, to: db}}, nil } // Close releases the database. -func (s *SQLite) Close() error { return s.raw.Close() } +func (s *DB) Close() error { return s.raw.Close() } // --- tenancy --- // PutOrg upserts an organization. -func (s *SQLite) PutOrg(ctx context.Context, o Org) error { +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) @@ -87,7 +96,7 @@ func (s *SQLite) PutOrg(ctx context.Context, o Org) error { } // PutTeam upserts a team. -func (s *SQLite) PutTeam(ctx context.Context, t Team) error { +func (s *DB) 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`, @@ -96,7 +105,7 @@ func (s *SQLite) PutTeam(ctx context.Context, t Team) error { } // PutProject upserts a project. -func (s *SQLite) PutProject(ctx context.Context, p Project) error { +func (s *DB) 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`, @@ -105,7 +114,7 @@ func (s *SQLite) PutProject(ctx context.Context, p Project) error { } // Orgs lists every organization. -func (s *SQLite) Orgs(ctx context.Context) ([]Org, error) { +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 @@ -114,7 +123,7 @@ func (s *SQLite) Orgs(ctx context.Context) ([]Org, error) { } // Teams lists an organization's teams. -func (s *SQLite) Teams(ctx context.Context, orgID string) ([]Team, error) { +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 @@ -123,7 +132,7 @@ func (s *SQLite) Teams(ctx context.Context, orgID string) ([]Team, error) { } // Projects lists a team's projects. -func (s *SQLite) Projects(ctx context.Context, teamID string) ([]Project, error) { +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 @@ -138,7 +147,7 @@ const keyColumns = `id, key_hash, key_prefix, name, project_id, band, routing_pr // 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 { +func (s *DB) CreateKey(ctx context.Context, k Key) error { if len(k.Hash) == 0 { return errors.New("store: a key needs a hash") } @@ -151,7 +160,7 @@ func (s *SQLite) CreateKey(ctx context.Context, k Key) error { } // KeyByHash finds a key by the hash of the presented secret. -func (s *SQLite) KeyByHash(ctx context.Context, hash []byte) (Key, error) { +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) @@ -162,14 +171,14 @@ func (s *SQLite) KeyByHash(ctx context.Context, hash []byte) (Key, error) { } // Keys lists a project's keys, revoked ones included. -func (s *SQLite) Keys(ctx context.Context, projectID string) ([]Key, error) { +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 *SQLite) RevokeKey(ctx context.Context, id string, at time.Time) error { +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) @@ -196,7 +205,7 @@ func (s *SQLite) RevokeKey(ctx context.Context, id string, at time.Time) error { // 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 { +func (s *DB) TouchKeys(ctx context.Context, seen map[string]time.Time) error { if len(seen) == 0 { return nil } @@ -219,7 +228,7 @@ func (s *SQLite) TouchKeys(ctx context.Context, seen map[string]time.Time) error // --- limits --- // PutLimits upserts the limits at one scope. -func (s *SQLite) PutLimits(ctx context.Context, l Limits) error { +func (s *DB) PutLimits(ctx context.Context, l Limits) error { models, err := encodeModels(l.AllowedModels) if err != nil { return err @@ -237,7 +246,7 @@ func (s *SQLite) PutLimits(ctx context.Context, l Limits) error { } // LimitsFor reads the limits set at one scope, without inheritance. -func (s *SQLite) LimitsFor(ctx context.Context, scope Scope, id string) (Limits, error) { +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) @@ -276,7 +285,7 @@ func (s *SQLite) LimitsFor(ctx context.Context, scope Scope, id string) (Limits, // // 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) { +func (s *DB) 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 @@ -330,7 +339,7 @@ func (s *SQLite) Resolve(ctx context.Context, keyID string) (Effective, error) { // 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 *SQLite) SetDefaultLimits(rpm, tpm *int) { +func (s *DB) SetDefaultLimits(rpm, tpm *int) { s.mu.Lock() defer s.mu.Unlock() s.defaultRPM, s.defaultTPM = rpm, tpm @@ -342,7 +351,7 @@ func (s *SQLite) SetDefaultLimits(rpm, tpm *int) { // 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 *SQLite) applyDefaults(e *Effective) { +func (s *DB) applyDefaults(e *Effective) { s.mu.RLock() defer s.mu.RUnlock() own := LimitOwner{Scope: ScopeKey, ID: e.KeyID} @@ -357,7 +366,7 @@ func (s *SQLite) applyDefaults(e *Effective) { } // SpendSinceByScope totals cost per key, with the key's place in the tree. -func (s *SQLite) SpendSinceByScope(ctx context.Context, since time.Time) ([]ScopeSpend, error) { +func (s *DB) SpendSinceByScope(ctx context.Context, since time.Time) ([]ScopeSpend, error) { rows, err := s.db.QueryContext(ctx, `SELECT k.id, p.id, t.id, o.id, SUM(u.cost_micros) FROM usage u @@ -428,7 +437,7 @@ func (e *Effective) tighten(l Limits) { // --- accounting --- // WriteUsage appends a batch of records in one transaction. -func (s *SQLite) WriteUsage(ctx context.Context, records []Usage) error { +func (s *DB) WriteUsage(ctx context.Context, records []Usage) error { if len(records) == 0 { return nil } @@ -453,7 +462,7 @@ func (s *SQLite) WriteUsage(ctx context.Context, records []Usage) error { } // SpendSince totals cost per key since a moment. -func (s *SQLite) SpendSince(ctx context.Context, since time.Time) (map[string]int64, error) { +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 { @@ -478,7 +487,7 @@ func (s *SQLite) SpendSince(ctx context.Context, since time.Time) (map[string]in // 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 *SQLite) begin(ctx context.Context) (conn, *sql.Tx, error) { +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 @@ -488,7 +497,7 @@ func (s *SQLite) begin(ctx context.Context) (conn, *sql.Tx, error) { // inTx runs fn in a transaction, with the same placeholder rewriting the // outside-transaction path gets. -func (s *SQLite) inTx(ctx context.Context, fn func(conn) error) error { +func (s *DB) inTx(ctx context.Context, fn func(conn) error) error { tx, err := s.raw.BeginTx(ctx, nil) if err != nil { return err @@ -620,14 +629,14 @@ func fromNullMillis(v sql.NullInt64) time.Time { func nowMillis() int64 { return time.Now().UnixMilli() } // SQLite satisfies Store. -var _ Store = (*SQLite)(nil) +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 *SQLite) UsageSince(ctx context.Context, since time.Time) ([]KeyUsage, error) { +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, @@ -638,7 +647,9 @@ func (s *SQLite) UsageSince(ctx context.Context, since time.Time) ([]KeyUsage, e 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 + -- k.name and k.key_prefix are selected, so PostgreSQL wants them + -- grouped. Both are determined by the key id already in the list. + GROUP BY u.key_id, u.project_id, k.name, k.key_prefix ORDER BY SUM(u.cost_micros) DESC, COUNT(*) DESC`, []any{millis(since)}, func(rows *sql.Rows) (KeyUsage, error) { var k KeyUsage @@ -655,7 +666,7 @@ func (s *SQLite) UsageSince(ctx context.Context, since time.Time) ([]KeyUsage, e // 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 *SQLite) UsageBuckets(ctx context.Context, since time.Time, width time.Duration) ([]Bucket, error) { +func (s *DB) UsageBuckets(ctx context.Context, since time.Time, width time.Duration) ([]Bucket, error) { if width <= 0 { width = time.Hour } @@ -667,7 +678,11 @@ func (s *SQLite) UsageBuckets(ctx context.Context, since time.Time, width time.D SUM(prompt_tokens), SUM(completion_tokens), SUM(cost_micros), AVG(latency_ms) FROM usage WHERE ts >= ? - GROUP BY ts / ? ORDER BY 1`, ms, ms, millis(since), ms) + -- 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 } @@ -728,7 +743,7 @@ func (u UserUsage) Tokens() int64 { return u.Prompt + u.Completion } // 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 *SQLite) UsageByUser(ctx context.Context, since time.Time) ([]UserUsage, error) { +func (s *DB) UsageByUser(ctx context.Context, since time.Time) ([]UserUsage, error) { rows, err := s.db.QueryContext(ctx, ` SELECT COALESCE(k.user_id, ''), COALESCE(us.email, ''), COALESCE(us.name, ''), COUNT(DISTINCT u.key_id), @@ -740,7 +755,12 @@ func (s *SQLite) UsageByUser(ctx context.Context, since time.Time) ([]UserUsage, LEFT JOIN api_keys k ON k.id = u.key_id LEFT JOIN users us ON us.id = k.user_id WHERE u.ts >= ? - GROUP BY COALESCE(k.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 COALESCE(k.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 @@ -781,7 +801,7 @@ func (m ModelUsage) Tokens() int64 { return m.Prompt + m.Completion } // 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 *SQLite) UsageByModel(ctx context.Context, since time.Time) ([]ModelUsage, error) { +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), diff --git a/internal/store/store_test.go b/internal/store/store_test.go index b222215..5797582 100644 --- a/internal/store/store_test.go +++ b/internal/store/store_test.go @@ -11,7 +11,7 @@ import ( "time" ) -func openDB(t *testing.T) *SQLite { +func openDB(t *testing.T) *DB { t.Helper() s, err := OpenSQLite(context.Background(), filepath.Join(t.TempDir(), "llmgw.db")) if err != nil { @@ -22,7 +22,7 @@ func openDB(t *testing.T) *SQLite { } // tree builds org → team → project and returns the project id. -func tree(t *testing.T, s *SQLite) string { +func tree(t *testing.T, s *DB) string { t.Helper() ctx := context.Background() if err := s.PutOrg(ctx, Org{ID: "org1", Name: "Acme"}); err != nil { @@ -37,7 +37,7 @@ func tree(t *testing.T, s *SQLite) string { return "proj1" } -func mustKey(t *testing.T, s *SQLite, id, project string) (string, Key) { +func mustKey(t *testing.T, s *DB, id, project string) (string, Key) { t.Helper() secret, k, err := NewKey(id, id, project) if err != nil { diff --git a/internal/store/threads.go b/internal/store/threads.go index fca9548..8c7506d 100644 --- a/internal/store/threads.go +++ b/internal/store/threads.go @@ -39,7 +39,7 @@ type Message struct { // A real project rather than a special case in every query: threads and notes // are scoped by project id and enforce it in SQL, and an empty id would have to // be an exception in each of those places. -func (s *SQLite) EnsureAnonymousProject(ctx context.Context, id string, retainHours int) error { +func (s *DB) EnsureAnonymousProject(ctx context.Context, id string, retainHours int) error { if err := s.PutOrg(ctx, Org{ID: id, Name: id}); err != nil { return err } @@ -59,7 +59,7 @@ func (s *SQLite) EnsureAnonymousProject(ctx context.Context, id string, retainHo // // Zero means it keeps none, which is the default and was the only behaviour // before threads existed. -func (s *SQLite) ThreadRetention(ctx context.Context, projectID string) (time.Duration, error) { +func (s *DB) ThreadRetention(ctx context.Context, projectID string) (time.Duration, error) { var hours sql.NullInt64 err := s.db.QueryRowContext(ctx, `SELECT thread_retention_hours FROM projects WHERE id = ?`, projectID).Scan(&hours) @@ -80,7 +80,7 @@ func (s *SQLite) ThreadRetention(ctx context.Context, projectID string) (time.Du // Hours rather than a boolean: a thread store with no expiry is a transcript // of everything a team ever typed, kept forever because nobody chose a // number. Turning it on means saying for how long. -func (s *SQLite) SetThreadRetention(ctx context.Context, projectID string, hours int) error { +func (s *DB) SetThreadRetention(ctx context.Context, projectID string, hours int) error { var v any if hours > 0 { v = hours @@ -100,7 +100,7 @@ func (s *SQLite) SetThreadRetention(ctx context.Context, projectID string, hours // // The project is part of the lookup rather than checked afterwards, so there // is no path that reads another project's conversation even briefly. -func (s *SQLite) Thread(ctx context.Context, projectID, threadID string) ([]Message, error) { +func (s *DB) Thread(ctx context.Context, projectID, threadID string) ([]Message, error) { var owner string err := s.db.QueryRowContext(ctx, `SELECT project_id FROM threads WHERE id = ?`, threadID).Scan(&owner) @@ -125,7 +125,7 @@ func (s *SQLite) Thread(ctx context.Context, projectID, threadID string) ([]Mess // One transaction, because a half-written exchange is worse than a lost one: // a stored question with no answer would be replayed to the next model as // though the assistant had ignored it. -func (s *SQLite) AppendThread(ctx context.Context, projectID, threadID string, msgs []Message) error { +func (s *DB) AppendThread(ctx context.Context, projectID, threadID string, msgs []Message) error { if len(msgs) == 0 { return nil } @@ -191,7 +191,7 @@ func (s *SQLite) AppendThread(ctx context.Context, projectID, threadID string, m // // Somebody has to be able to say "forget that", and a retention window is not // an answer to a message sent by mistake. -func (s *SQLite) DeleteThread(ctx context.Context, projectID, threadID string) error { +func (s *DB) DeleteThread(ctx context.Context, projectID, threadID string) error { res, err := s.db.ExecContext(ctx, `DELETE FROM threads WHERE id = ? AND project_id = ?`, threadID, projectID) if err != nil { @@ -208,7 +208,7 @@ func (s *SQLite) DeleteThread(ctx context.Context, projectID, threadID string) e // Per project, because the window is per project. A single sweep with one // cutoff would apply the shortest retention to everyone or the longest to // nobody, and the number a project chose is the promise made to it. -func (s *SQLite) PruneThreads(ctx context.Context, now time.Time) (int, error) { +func (s *DB) PruneThreads(ctx context.Context, now time.Time) (int, error) { type window struct { project string hours int64 diff --git a/internal/store/users.go b/internal/store/users.go index 0fc139a..217ccae 100644 --- a/internal/store/users.go +++ b/internal/store/users.go @@ -59,7 +59,7 @@ const userColumns = `id, email, name, password_hash, state, COALESCE(project_id, created_at, approved_at, approved_by, last_login_at, is_admin` // CreateUser records a new, unapproved account. -func (s *SQLite) CreateUser(ctx context.Context, u User) error { +func (s *DB) CreateUser(ctx context.Context, u User) error { u.Email = NormalizeEmail(u.Email) _, err := s.db.ExecContext(ctx, `INSERT INTO users (id, email, name, password_hash, state, project_id, created_at) @@ -73,19 +73,19 @@ func (s *SQLite) CreateUser(ctx context.Context, u User) error { } // UserByEmail finds an account by its address. -func (s *SQLite) UserByEmail(ctx context.Context, email string) (User, error) { +func (s *DB) UserByEmail(ctx context.Context, email string) (User, error) { return s.scanUser(s.db.QueryRowContext(ctx, `SELECT `+userColumns+` FROM users WHERE email = ?`, NormalizeEmail(email))) } // UserByID finds an account. -func (s *SQLite) UserByID(ctx context.Context, id string) (User, error) { +func (s *DB) UserByID(ctx context.Context, id string) (User, error) { return s.scanUser(s.db.QueryRowContext(ctx, `SELECT `+userColumns+` FROM users WHERE id = ?`, id)) } // Users lists accounts, optionally filtered by state. -func (s *SQLite) Users(ctx context.Context, state UserState) ([]User, error) { +func (s *DB) Users(ctx context.Context, state UserState) ([]User, error) { query := `SELECT ` + userColumns + ` FROM users` var args []any if state != "" { @@ -102,7 +102,7 @@ func (s *SQLite) Users(ctx context.Context, state UserState) ([]User, error) { // // The project is required: an approved account with nowhere to put its keys // would pass every check and then fail at the first key it tried to make. -func (s *SQLite) ApproveUser(ctx context.Context, id, projectID, by string, at time.Time) error { +func (s *DB) ApproveUser(ctx context.Context, id, projectID, by string, at time.Time) error { if projectID == "" { return errors.New("approving an account needs a project for its keys") } @@ -120,7 +120,7 @@ func (s *SQLite) ApproveUser(ctx context.Context, id, projectID, by string, at t } // SetUserState suspends or reinstates an account. -func (s *SQLite) SetUserState(ctx context.Context, id string, state UserState) error { +func (s *DB) SetUserState(ctx context.Context, id string, state UserState) error { res, err := s.db.ExecContext(ctx, `UPDATE users SET state = ? WHERE id = ?`, string(state), id) if err != nil { return err @@ -137,7 +137,7 @@ func (s *SQLite) SetUserState(ctx context.Context, id string, state UserState) e // worst a stale session buys is a few more seconds of looking. Suspending the // account is what closes sessions, and that is the lever for someone who should // not be here at all. -func (s *SQLite) SetAdmin(ctx context.Context, id string, admin bool) error { +func (s *DB) SetAdmin(ctx context.Context, id string, admin bool) error { res, err := s.db.ExecContext(ctx, `UPDATE users SET is_admin = ? WHERE id = ?`, admin, id) if err != nil { return err @@ -149,21 +149,21 @@ func (s *SQLite) SetAdmin(ctx context.Context, id string, admin bool) error { } // TouchLogin records a successful sign-in. -func (s *SQLite) TouchLogin(ctx context.Context, id string, at time.Time) error { +func (s *DB) TouchLogin(ctx context.Context, id string, at time.Time) error { _, err := s.db.ExecContext(ctx, `UPDATE users SET last_login_at = ? WHERE id = ?`, millis(at), id) return err } // KeysForUser lists the keys a person made for themselves. -func (s *SQLite) KeysForUser(ctx context.Context, userID string) ([]Key, error) { +func (s *DB) KeysForUser(ctx context.Context, userID string) ([]Key, error) { return queryAll(ctx, s.db, `SELECT `+keyColumns+` FROM api_keys WHERE user_id = ? ORDER BY created_at DESC`, []any{userID}, func(rows *sql.Rows) (Key, error) { return scanKey(rows) }) } // CreateUserKey issues a key owned by a person. -func (s *SQLite) CreateUserKey(ctx context.Context, userID string, k Key) error { +func (s *DB) CreateUserKey(ctx context.Context, userID string, k Key) error { if len(k.Hash) == 0 { return errors.New("store: a key needs a hash") } @@ -182,7 +182,7 @@ func (s *SQLite) CreateUserKey(ctx context.Context, userID string, k Key) error // queries would leave a window where a key changes hands between them, and more // importantly it would be one forgotten check away from letting anyone revoke // anyone's key by id. -func (s *SQLite) RevokeUserKey(ctx context.Context, userID, keyID string, at time.Time) error { +func (s *DB) RevokeUserKey(ctx context.Context, userID, keyID string, at time.Time) error { res, err := s.db.ExecContext(ctx, `UPDATE api_keys SET revoked_at = ? WHERE id = ? AND user_id = ? AND revoked_at IS NULL`, @@ -196,7 +196,7 @@ func (s *SQLite) RevokeUserKey(ctx context.Context, userID, keyID string, at tim return nil } -func (s *SQLite) scanUser(sc scanner) (User, error) { +func (s *DB) scanUser(sc scanner) (User, error) { var u User var state string var created int64 @@ -224,7 +224,7 @@ func (s *SQLite) scanUser(sc scanner) (User, error) { // reasons to change a password want different things. Someone rotating theirs // wants their other devices signed out. An operator resetting a forgotten one // wants the same, but also has nobody to keep signed in. -func (s *SQLite) SetPassword(ctx context.Context, userID, hash string) error { +func (s *DB) SetPassword(ctx context.Context, userID, hash string) error { res, err := s.db.ExecContext(ctx, `UPDATE users SET password_hash = ? WHERE id = ?`, hash, userID) if err != nil { @@ -240,7 +240,7 @@ func (s *SQLite) SetPassword(ctx context.Context, userID, hash string) error { // // For the approval form, which has to offer somewhere to put a new account's // keys. Walking the tree in the page would mean three round trips per render. -func (s *SQLite) AllProjects(ctx context.Context) ([]ProjectPath, error) { +func (s *DB) AllProjects(ctx context.Context) ([]ProjectPath, error) { return queryAll(ctx, s.db, ` SELECT p.id, o.name || ' / ' || t.name || ' / ' || p.name FROM projects p @@ -263,7 +263,7 @@ type ProjectPath struct { // has no key in their hand — the key was shown once and is theirs to keep — so // their identity there is their project, and it must be resolved the same way // or the page would be a hole in the limits everything else obeys. -func (s *SQLite) ResolveProject(ctx context.Context, projectID string) (Effective, error) { +func (s *DB) ResolveProject(ctx context.Context, projectID string) (Effective, error) { var e Effective row := s.db.QueryRowContext(ctx, `SELECT p.id, t.id, o.id diff --git a/test/e2e/accounting_test.go b/test/e2e/accounting_test.go index 000149c..d76516b 100644 --- a/test/e2e/accounting_test.go +++ b/test/e2e/accounting_test.go @@ -74,7 +74,7 @@ func TestUsageReachesTheStoreWithoutBlockingRequests(t *testing.T) { func TestRequestRateLimitsBindThroughTheGateway(t *testing.T) { path, mint := provision(t) key := mint(t, "burst") - withStore(t, path, func(s *store.SQLite) { + withStore(t, path, func(s *store.DB) { rpm := 5 if err := s.PutLimits(context.Background(), store.Limits{ Scope: store.ScopeKey, ScopeID: "key-burst", RPM: &rpm, @@ -107,7 +107,7 @@ func TestRequestRateLimitsBindThroughTheGateway(t *testing.T) { func TestASpentBudgetRefusesAndDoesNotSuggestRetrying(t *testing.T) { path, mint := provision(t) key := mint(t, "broke") - withStore(t, path, func(s *store.SQLite) { + withStore(t, path, func(s *store.DB) { // One micro-unit: the first request exhausts it. budget := int64(1) if err := s.PutLimits(context.Background(), store.Limits{ @@ -162,7 +162,7 @@ func TestBudgetSurvivesAGatewayRestart(t *testing.T) { // A second gateway over the same store, with a budget just under what was // already spent: it must refuse immediately rather than starting fresh. - withStore(t, path, func(s *store.SQLite) { + withStore(t, path, func(s *store.DB) { budget := spent[0].CostMicros if err := s.PutLimits(context.Background(), store.Limits{ Scope: store.ScopeKey, ScopeID: "key-persistent", @@ -238,7 +238,7 @@ func TestASpentBudgetIsServedAnAffordableModel(t *testing.T) { // guardrail rather than a wall. path, mint := provision(t) key := mint(t, "broke") - withStore(t, path, func(s *store.SQLite) { + withStore(t, path, func(s *store.DB) { budget := int64(1) // one micro-unit: the first request exhausts it if err := s.PutLimits(context.Background(), store.Limits{ Scope: store.ScopeKey, ScopeID: "key-broke", @@ -276,7 +276,7 @@ func TestASpentBudgetWithNoFallbackStillRefuses(t *testing.T) { // refused exactly as before, or a budget would stop meaning anything. path, mint := provision(t) key := mint(t, "broke") - withStore(t, path, func(s *store.SQLite) { + withStore(t, path, func(s *store.DB) { budget := int64(1) if err := s.PutLimits(context.Background(), store.Limits{ Scope: store.ScopeKey, ScopeID: "key-broke", diff --git a/test/e2e/auth_test.go b/test/e2e/auth_test.go index e494410..0e1f44b 100644 --- a/test/e2e/auth_test.go +++ b/test/e2e/auth_test.go @@ -66,7 +66,7 @@ func provision(t *testing.T) (string, func(t *testing.T, name string, tune ...fu } // withStore sets limits on an existing store. -func withStore(t *testing.T, path string, fn func(*store.SQLite)) { +func withStore(t *testing.T, path string, fn func(*store.DB)) { t.Helper() s, err := store.OpenSQLite(context.Background(), path) if err != nil { @@ -140,7 +140,7 @@ func TestRevocationTakesEffectWithinTheCacheTTL(t *testing.T) { t.Fatal("the key did not work to begin with") } - withStore(t, path, func(s *store.SQLite) { + withStore(t, path, func(s *store.DB) { if err := s.RevokeKey(context.Background(), "key-doomed", time.Now()); err != nil { t.Fatal(err) } @@ -166,7 +166,7 @@ func TestAnAllowlistCoversEveryNameOfAModel(t *testing.T) { byAlias := mint(t, "alias") other := mint(t, "other") - withStore(t, path, func(s *store.SQLite) { + withStore(t, path, func(s *store.DB) { for _, l := range []store.Limits{ {Scope: store.ScopeKey, ScopeID: "key-canonical", AllowedModels: []string{testModel}}, {Scope: store.ScopeKey, ScopeID: "key-alias", AllowedModels: []string{"alias-for-test-model"}}, @@ -317,7 +317,7 @@ func TestOnboardingAUserEndToEnd(t *testing.T) { // The operator's side. var link string - withStore(t, path, func(s *store.SQLite) { + withStore(t, path, func(s *store.DB) { secret, k, err := store.NewKey("key-new", "ayse-laptop", "proj") if err != nil { t.Fatal(err) @@ -396,7 +396,7 @@ func TestTheSelfDescriptionEndpointAnswersWithoutAnOperator(t *testing.T) { // with access to the dashboard. path, mint := provision(t) key := mint(t, "curious") - withStore(t, path, func(s *store.SQLite) { + withStore(t, path, func(s *store.DB) { rpm := 42 if err := s.PutLimits(context.Background(), store.Limits{ Scope: store.ScopeKey, ScopeID: "key-curious", diff --git a/test/e2e/fallback_context_test.go b/test/e2e/fallback_context_test.go index b94b8aa..77c9733 100644 --- a/test/e2e/fallback_context_test.go +++ b/test/e2e/fallback_context_test.go @@ -101,7 +101,7 @@ func contextFallbackConfig(path string) func(*config.Server) { // spendTheBudget exhausts a key so the next request falls back. func spendTheBudget(t *testing.T, path, id string) { t.Helper() - withStore(t, path, func(s *store.SQLite) { + withStore(t, path, func(s *store.DB) { budget := int64(1) if err := s.PutLimits(context.Background(), store.Limits{ Scope: store.ScopeKey, ScopeID: id, diff --git a/test/e2e/signup_test.go b/test/e2e/signup_test.go index 890200e..b57aa3c 100644 --- a/test/e2e/signup_test.go +++ b/test/e2e/signup_test.go @@ -120,7 +120,7 @@ func TestAnAccountIsUselessUntilItIsApproved(t *testing.T) { } // Approve, and it all works. - withStore(t, path, func(s *store.SQLite) { + withStore(t, path, func(s *store.DB) { users, err := s.Users(context.Background(), store.UserPending) if err != nil || len(users) != 1 { t.Fatalf("pending users: %v %v", users, err) @@ -181,7 +181,7 @@ func TestAUserCannotRevokeSomeoneElsesKey(t *testing.T) { ayseKey := extractKey(t, body) var ayseKeyID string - withStore(t, path, func(s *store.SQLite) { + withStore(t, path, func(s *store.DB) { users, _ := s.Users(context.Background(), store.UserActive) for _, u := range users { if u.Email == "ayse@example.com" { @@ -249,7 +249,7 @@ func TestSuspensionClosesOpenSessions(t *testing.T) { if code, _ := b.get("/profile"); code != http.StatusOK { t.Fatalf("profile before suspension: %d", code) } - withStore(t, path, func(s *store.SQLite) { + withStore(t, path, func(s *store.DB) { if err := s.SetUserState(context.Background(), id, store.UserSuspended); err != nil { t.Fatal(err) } @@ -277,7 +277,7 @@ func TestSignupCanBeClosed(t *testing.T) { }); code == http.StatusOK || code == http.StatusSeeOther { t.Errorf("signup accepted a post although it is closed: %d", code) } - withStore(t, path, func(s *store.SQLite) { + withStore(t, path, func(s *store.DB) { users, err := s.Users(context.Background(), "") if err != nil { t.Fatal(err) @@ -298,7 +298,7 @@ func signUpAndApprove(t *testing.T, b *browser, path, email string) string { } var id string - withStore(t, path, func(s *store.SQLite) { + withStore(t, path, func(s *store.DB) { users, err := s.Users(context.Background(), store.UserPending) if err != nil { t.Fatal(err) @@ -437,7 +437,7 @@ func TestTheClusterPageIsOnlyForAdmins(t *testing.T) { t.Error("a non-admin is shown a link to the cluster") } - withStore(t, path, func(s *store.SQLite) { + withStore(t, path, func(s *store.DB) { if err := s.SetAdmin(context.Background(), id, true); err != nil { t.Fatal(err) } @@ -555,7 +555,7 @@ func TestChatObeysTheModelAllowlist(t *testing.T) { b := newBrowser(t, h.router.URL) signUpAndApprove(t, b, path, "ayse@example.com") - withStore(t, path, func(s *store.SQLite) { + withStore(t, path, func(s *store.DB) { if err := s.PutLimits(context.Background(), store.Limits{ Scope: store.ScopeProject, ScopeID: "proj", AllowedModels: []string{"something-else"}, diff --git a/test/e2e/threads_test.go b/test/e2e/threads_test.go index 0212e04..b30298f 100644 --- a/test/e2e/threads_test.go +++ b/test/e2e/threads_test.go @@ -57,7 +57,7 @@ func authConfig(path string) func(*config.Server) { func enableThreads(t *testing.T, path string, hours int) { t.Helper() - withStore(t, path, func(s *store.SQLite) { + withStore(t, path, func(s *store.DB) { if err := s.SetThreadRetention(context.Background(), "proj", hours); err != nil { t.Fatal(err) } @@ -67,7 +67,7 @@ func enableThreads(t *testing.T, path string, hours int) { func storedTurns(t *testing.T, path, thread string) []store.Message { t.Helper() var out []store.Message - withStore(t, path, func(s *store.SQLite) { + withStore(t, path, func(s *store.DB) { msgs, err := s.Thread(context.Background(), "proj", thread) if err != nil { t.Fatalf("reading thread %q: %v", thread, err) @@ -164,7 +164,7 @@ func TestOneProjectCannotReadAnothersThread(t *testing.T) { enableThreads(t, path, 24) // A key in a second project. - withStore(t, path, func(s *store.SQLite) { + withStore(t, path, func(s *store.DB) { if err := s.PutProject(context.Background(), store.Project{ID: "other", TeamID: "team", Name: "Other"}); err != nil { t.Fatal(err) @@ -211,7 +211,7 @@ func TestWithoutAThreadNothingIsStoredAtAll(t *testing.T) { t.Fatalf("%d %s", code, msg) } - withStore(t, path, func(s *store.SQLite) { + withStore(t, path, func(s *store.DB) { // Pruning with a cutoff far past every retention window deletes every // thread there is, so a count of zero means none were created. n, err := s.PruneThreads(context.Background(), time.Now().Add(1000*time.Hour)) @@ -258,7 +258,7 @@ func TestSomeoneCanUnsayWhatTheySaid(t *testing.T) { if code := h.forget(t, key, "th1"); code != http.StatusNoContent { t.Fatalf("deleting the thread returned %d", code) } - withStore(t, path, func(s *store.SQLite) { + withStore(t, path, func(s *store.DB) { if _, err := s.Thread(context.Background(), "proj", "th1"); err == nil { t.Error("the conversation survived being deleted") } @@ -272,7 +272,7 @@ func TestDeletingSomeoneElsesThreadLooksLikeDeletingNothing(t *testing.T) { path, mint := provision(t) mine := mint(t, "mine") enableThreads(t, path, 24) - withStore(t, path, func(s *store.SQLite) { + withStore(t, path, func(s *store.DB) { if err := s.PutProject(context.Background(), store.Project{ID: "other", TeamID: "team", Name: "Other"}); err != nil { t.Fatal(err)