diff --git a/go.mod b/go.mod index ee7bfe6d..d5d4e851 100644 --- a/go.mod +++ b/go.mod @@ -6,6 +6,7 @@ require ( connectrpc.com/connect v1.21.0 github.com/NVIDIA/go-nvml v0.13.4-0 github.com/cespare/xxhash/v2 v2.3.0 + github.com/jackc/pgx/v5 v5.11.0 github.com/knadh/koanf/parsers/yaml v1.1.1 github.com/knadh/koanf/providers/env/v2 v2.0.1 github.com/knadh/koanf/providers/file v1.2.1 @@ -14,6 +15,7 @@ require ( github.com/prometheus/client_model v0.6.3 github.com/prometheus/common v0.71.0 github.com/spf13/cobra v1.10.2 + golang.org/x/crypto v0.57.0 golang.org/x/net v0.59.0 golang.org/x/sync v0.23.0 google.golang.org/protobuf v1.36.12 @@ -27,6 +29,9 @@ require ( github.com/go-viper/mapstructure/v2 v2.4.0 // indirect github.com/google/uuid v1.6.0 // indirect github.com/inconshreveable/mousetrap v1.1.0 // indirect + github.com/jackc/pgpassfile v1.0.0 // indirect + github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect + github.com/jackc/puddle/v2 v2.2.2 // indirect github.com/knadh/koanf/maps v0.1.2 // indirect github.com/mattn/go-isatty v0.0.24 // indirect github.com/mitchellh/copystructure v1.2.0 // indirect @@ -37,7 +42,6 @@ require ( github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect github.com/spf13/pflag v1.0.9 // indirect go.yaml.in/yaml/v3 v3.0.5 // indirect - golang.org/x/crypto v0.57.0 // indirect golang.org/x/sys v0.48.0 // indirect golang.org/x/text v0.42.0 // indirect modernc.org/libc v1.75.7 // indirect diff --git a/go.sum b/go.sum index 04191895..18e313b3 100644 --- a/go.sum +++ b/go.sum @@ -7,6 +7,7 @@ github.com/beorn7/perks v1.0.1/go.mod h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6r github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs= github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= github.com/cpuguy83/go-md2man/v2 v2.0.6/go.mod h1:oOW0eioCTA6cOiMLiUPZOpcVxMig6NIQQ7OS05n1F4g= +github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY= github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto= github.com/fsnotify/fsnotify v1.9.0 h1:2Ml+OJNzbYCTzsxtv8vKSFD9PbJjmhYF14k/jKC7S9k= @@ -23,6 +24,14 @@ github.com/hashicorp/golang-lru/v2 v2.0.7 h1:a+bsQ5rvGLjzHuww6tVxozPZFVghXaHOwFs github.com/hashicorp/golang-lru/v2 v2.0.7/go.mod h1:QeFd9opnmA6QUJc5vARoKUSoFhyfM2/ZepoAG6RGpeM= github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8= github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw= +github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM= +github.com/jackc/pgpassfile v1.0.0/go.mod h1:CEx0iS5ambNFdcRtxPj5JhEz+xB6uRky5eyVu/W2HEg= +github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 h1:iCEnooe7UlwOQYpKFhBabPMi4aNAfoODPEFNiAnClxo= +github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761/go.mod h1:5TJZWKEWniPve33vlWYSoGYefn3gLQRzjfDlhSJ9ZKM= +github.com/jackc/pgx/v5 v5.11.0 h1:IzBBtyK9AHqf98cctWFifYSci2hgQR/cd56wB4p+ogg= +github.com/jackc/pgx/v5 v5.11.0/go.mod h1:mal1tBGAFfLHvZzaYh77YS/eC6IX9OWbRV1QIIM0Jn4= +github.com/jackc/puddle/v2 v2.2.2 h1:PR8nw+E/1w0GLuRFSmiioY6UooMp6KJv0/61nB7icHo= +github.com/jackc/puddle/v2 v2.2.2/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4= github.com/klauspost/compress v1.19.1 h1:VsB4HPswih7mmZ8WleSFQ75c/Ui1M4trX5oAsJnhSlk= github.com/klauspost/compress v1.19.1/go.mod h1:cwPg85FWrGar70rWktvGQj8/hthj3wpl0PGDogxkrSQ= github.com/knadh/koanf/maps v0.1.2 h1:RBfmAW5CnZT+PJ1CVc1QSJKf4Xu9kxfQgYVQSu8hpbo= @@ -47,6 +56,7 @@ github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 h1:C3w9PqII01/Oq github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822/go.mod h1:+n7T8mK8HuQTcFwEeznm/DIxMOiR9yIdICNftLE1DvQ= github.com/ncruces/go-strftime v1.0.0 h1:HMFp8mLCTPp341M/ZnA4qaf7ZlsbTc+miZjCLOFAw7w= github.com/ncruces/go-strftime v1.0.0/go.mod h1:Fwc5htZGVVkseilnfgOVb9mKy6w1naJmn9CehxcKcls= +github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/prometheus/client_golang v1.24.1 h1:JnJkREXzWxUdCuPFpIWZiPispT9xVV59uiuyR2bPlnU= github.com/prometheus/client_golang v1.24.1/go.mod h1:F+oSRECHg4sse5ucfYpYDeIv/hu68Zo0uoHKetWnzcE= github.com/prometheus/client_model v0.6.3 h1:O0jaTVAYNxTHYInEPFJt5I3+sN8zqBtVMPTB1qyxiEo= @@ -62,6 +72,9 @@ github.com/spf13/cobra v1.10.2 h1:DMTTonx5m65Ic0GOoRY2c16WCbHxOOw6xxezuLaBpcU= github.com/spf13/cobra v1.10.2/go.mod h1:7C1pvHqHw5A4vrJfjNwvOdzYu0Gml16OCs2GRiTUUS4= github.com/spf13/pflag v1.0.9 h1:9exaQaMOCwffKiiiYk6/BndUBv+iRViNW+4lEMi0PvY= github.com/spf13/pflag v1.0.9/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg= +github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= +github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI= +github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= github.com/stretchr/testify v1.12.1 h1:EuwCh5fleGS7H32xRwO3wRGT7DxrDhLAT6FF8MpWDWE= github.com/stretchr/testify v1.12.1/go.mod h1:MDEgiDPPsNp5cuIrHPPCyornHKgEVbtFUmoNlxoYthg= go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto= @@ -88,6 +101,7 @@ golang.org/x/tools v0.49.0/go.mod h1:SJNXV9DBKT0UbdttsQjbfJlAE/q+y36++zo3uL3N0Oo google.golang.org/protobuf v1.36.12 h1:pJOKDDOyeXErUroCihFAd5LQuwXBSpVnKGrj5o/fwxc= google.golang.org/protobuf v1.36.12/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= modernc.org/cc/v4 v4.29.2 h1:h6+9ciCnPKutf4I03CvheAvDLX7+IHlqR6Iy6J+cgd8= modernc.org/cc/v4 v4.29.2/go.mod h1:OnovgIhbbMXMu1aISnJ0wvVD1KnW+cAUJkIrAWh+kVI= modernc.org/ccgo/v4 v4.35.0 h1:F+TUsmw09QxLzmi3aeYYGxjAXarmZaKgj3mKQHNaA8w= diff --git a/internal/cli/keys.go b/internal/cli/keys.go index 595b0a26..60380b5d 100644 --- a/internal/cli/keys.go +++ b/internal/cli/keys.go @@ -27,7 +27,7 @@ func newTenancyCmd() *cobra.Command { } cmd.PersistentFlags().StringVar(&dbPath, "store", "/var/lib/llmgw/llmgw.db", "path to the store") - open := func(ctx context.Context) (*store.SQLite, error) { return store.OpenSQLite(ctx, dbPath) } + open := func(ctx context.Context) (*store.SQLite, error) { return store.Open(ctx, dbPath) } cmd.AddCommand( newKeyInitCmd(open), diff --git a/internal/cli/users.go b/internal/cli/users.go index 2ca3715b..370e2096 100644 --- a/internal/cli/users.go +++ b/internal/cli/users.go @@ -24,7 +24,7 @@ func newUsersCmd() *cobra.Command { } cmd.PersistentFlags().StringVar(&dbPath, "store", "/var/lib/llmgw/llmgw.db", "path to the store") open := func(cmd *cobra.Command) (*store.SQLite, error) { - return store.OpenSQLite(cmd.Context(), dbPath) + return store.Open(cmd.Context(), dbPath) } cmd.AddCommand(newUserListCmd(open), newUserApproveCmd(open), diff --git a/internal/router/router.go b/internal/router/router.go index 193119da..b7873f0a 100644 --- a/internal/router/router.go +++ b/internal/router/router.go @@ -189,7 +189,7 @@ func New(cfg config.Server, log *slog.Logger) (*Router, error) { }, log) if cfg.Auth.Enabled { - st, err := storepkg.OpenSQLite(context.Background(), cfg.Auth.Store) + st, err := storepkg.Open(context.Background(), cfg.Auth.Store) if err != nil { return nil, fmt.Errorf("auth store: %w", err) } @@ -234,7 +234,7 @@ func New(cfg config.Server, log *slog.Logger) (*Router, error) { r.recorder, log) r.threads = r.store } else if cfg.Threads.Store != "" { - st, err := storepkg.OpenSQLite(context.Background(), cfg.Threads.Store) + st, err := storepkg.Open(context.Background(), cfg.Threads.Store) if err != nil { return nil, fmt.Errorf("threads store: %w", err) } diff --git a/internal/store/memory.go b/internal/store/memory.go index 68a40df6..d6b47c08 100644 --- a/internal/store/memory.go +++ b/internal/store/memory.go @@ -223,8 +223,20 @@ func (s *SQLite) SearchNotes(ctx context.Context, projectID, agentID, query, tag args := []any{projectID, agentID} if match := ftsQuery(query); match != "" { - q += ` AND id IN (SELECT rowid FROM memory_notes_fts WHERE memory_notes_fts MATCH ?)` - args = append(args, match) + // The one query in this store that is not the same sentence in both + // dialects. SQLite searches an FTS5 table kept in step by triggers; + // PostgreSQL matches a tsvector derived from the row. Writing it out + // twice is honest about that — the alternative was a shared spelling + // that means something subtly different in each. + if s.dialect == PostgresDialect { + q += ` AND search @@ plainto_tsquery('simple', ?)` + // plainto_tsquery takes words, not FTS5 syntax, so the raw query + // goes in rather than the quoted form SQLite needs. + args = append(args, query) + } else { + q += ` AND id IN (SELECT rowid FROM memory_notes_fts WHERE memory_notes_fts MATCH ?)` + args = append(args, match) + } } if tag != "" { t, err := NormalizeTag(tag) diff --git a/internal/store/migrate.go b/internal/store/migrate.go index f62e91c3..69051820 100644 --- a/internal/store/migrate.go +++ b/internal/store/migrate.go @@ -10,7 +10,7 @@ import ( "strings" ) -//go:embed migrations/sqlite/*.sql +//go:embed migrations/sqlite/*.sql migrations/postgres/*.sql var migrationFS embed.FS // migrate applies every unapplied migration in order, each in its own @@ -25,8 +25,17 @@ var migrationFS embed.FS // promise this could not keep. Same version numbers, same meaning, different // files. func migrate(ctx context.Context, db *sql.DB, dialect string) error { + // BIGINT, not INTEGER. SQLite's INTEGER holds 64 bits; PostgreSQL's is + // int4, and a millisecond timestamp does not fit in it — which the first + // migration discovers by failing to record that it ran. This table is + // created here rather than in a migration file, which is exactly how it + // escaped the port of the others. + stamp := "BIGINT" + if dialect == string(SQLiteDialect) { + stamp = "INTEGER" + } if _, err := db.ExecContext(ctx, - `CREATE TABLE IF NOT EXISTS schema_migrations (version TEXT PRIMARY KEY, applied_at INTEGER NOT NULL)`, + `CREATE TABLE IF NOT EXISTS schema_migrations (version TEXT PRIMARY KEY, applied_at `+stamp+` NOT NULL)`, ); err != nil { return fmt.Errorf("creating the migration table: %w", err) } @@ -63,7 +72,7 @@ func migrate(ctx context.Context, db *sql.DB, dialect string) error { if err != nil { return err } - if err := applyOne(ctx, db, name, string(body)); err != nil { + if err := applyOne(ctx, db, dialect, name, string(body)); err != nil { return fmt.Errorf("migration %s: %w", name, err) } } @@ -86,7 +95,7 @@ func migrationNames(dialect string) ([]string, error) { return names, nil } -func applyOne(ctx context.Context, db *sql.DB, name, body string) error { +func applyOne(ctx context.Context, db *sql.DB, dialect, name, body string) error { tx, err := db.BeginTx(ctx, nil) if err != nil { return err @@ -97,7 +106,7 @@ func applyOne(ctx context.Context, db *sql.DB, name, body string) error { return err } if _, err := tx.ExecContext(ctx, - `INSERT INTO schema_migrations (version, applied_at) VALUES (?, ?)`, + Dialect(dialect).rebind(`INSERT INTO schema_migrations (version, applied_at) VALUES (?, ?)`), name, nowMillis(), ); err != nil { return err diff --git a/internal/store/migrations/postgres/0001_init.sql b/internal/store/migrations/postgres/0001_init.sql new file mode 100644 index 00000000..9c5f6ff2 --- /dev/null +++ b/internal/store/migrations/postgres/0001_init.sql @@ -0,0 +1,81 @@ +-- Ported from the SQLite migration of the same number. Same version, same +-- meaning, different spelling: SQLite's storage classes and PostgreSQL's types +-- do not line up, and one shared file would be a promise neither could keep. +-- +-- Times are milliseconds since the epoch in BIGINT, exactly as in SQLite, +-- rather than TIMESTAMPTZ. The Go code reads and writes int64 milliseconds +-- everywhere; converting here would mean two time representations in one +-- codebase, and the conversion would be the thing that is wrong at 2am. + +-- The tenancy tree. Three levels, each owning the next. +CREATE TABLE organizations ( + id TEXT PRIMARY KEY, + name TEXT NOT NULL +); + +CREATE TABLE teams ( + id TEXT PRIMARY KEY, + org_id TEXT NOT NULL REFERENCES organizations(id), + name TEXT NOT NULL +); +CREATE INDEX teams_by_org ON teams(org_id); + +CREATE TABLE projects ( + id TEXT PRIMARY KEY, + team_id TEXT NOT NULL REFERENCES teams(id), + name TEXT NOT NULL +); +CREATE INDEX projects_by_team ON projects(team_id); + +-- Only sha256(key) is stored, so a database leak hands over no working +-- credential. SHA-256 rather than a slow password hash on purpose: the input is +-- 256 random bits with no dictionary to attack, so the cost of argon2 would buy +-- nothing and charge tens of milliseconds on every request (plan §7.9.1). +CREATE TABLE api_keys ( + id TEXT PRIMARY KEY, + key_hash BYTEA NOT NULL UNIQUE, + key_prefix TEXT NOT NULL, + name TEXT NOT NULL, + project_id TEXT NOT NULL REFERENCES projects(id), + band TEXT, + routing_profile TEXT, + created_at BIGINT NOT NULL, + expires_at BIGINT, + -- Revocation is a soft delete. Usage rows reference the key, and an audit + -- trail that loses its subject is not an audit trail. + revoked_at BIGINT, + last_used_at BIGINT +); +CREATE INDEX api_keys_by_project ON api_keys(project_id); + +-- Limits hang off any scope rather than being duplicated onto each one. A NULL +-- column means "no limit at this level", which is not the same as 0. +CREATE TABLE limits ( + scope_type TEXT NOT NULL, + scope_id TEXT NOT NULL, + rpm BIGINT, + tpm BIGINT, + -- Micro-units, not a float: a budget is accumulated by repeated addition, and + -- repeated addition of floats drifts. + budget_micros BIGINT, + budget_period TEXT, + allowed_models TEXT, + PRIMARY KEY (scope_type, scope_id) +); + +-- One row per request, written in batches off the request path. +CREATE TABLE usage ( + ts BIGINT NOT NULL, + key_id TEXT NOT NULL, + project_id TEXT NOT NULL, + model TEXT NOT NULL, + served_model TEXT NOT NULL, + prompt_tokens BIGINT NOT NULL, + completion_tokens BIGINT NOT NULL, + cost_micros BIGINT NOT NULL, + status BIGINT NOT NULL, + latency_ms BIGINT NOT NULL +); +-- Budget restore reads by key within a period, which is this index exactly. +CREATE INDEX usage_by_key_ts ON usage(key_id, ts); +CREATE INDEX usage_by_ts ON usage(ts); diff --git a/internal/store/migrations/postgres/0002_claims.sql b/internal/store/migrations/postgres/0002_claims.sql new file mode 100644 index 00000000..bb043407 --- /dev/null +++ b/internal/store/migrations/postgres/0002_claims.sql @@ -0,0 +1,29 @@ +-- Ported from the SQLite migration of the same number. Same version, same +-- meaning, different spelling: SQLite's storage classes and PostgreSQL's types +-- do not line up, and one shared file would be a promise neither could keep. +-- +-- Times are milliseconds since the epoch in BIGINT, exactly as in SQLite, +-- rather than TIMESTAMPTZ. The Go code reads and writes int64 milliseconds +-- everywhere; converting here would mean two time representations in one +-- codebase, and the conversion would be the thing that is wrong at 2am. + +-- A key has to reach the person it was issued for, and a key pasted into a chat +-- window stays in that chat window forever. A claim is a one-time link instead: +-- the operator sends the link, the person opens it once, and the link dies. +-- +-- The key is not stored here in the clear. It is sealed with a key derived from +-- the claim token, and the token is never written down — only its hash is. So a +-- database leak yields two useless hashes and a blob nothing in the database can +-- open, which keeps the property the api_keys table has. +CREATE TABLE key_claims ( + token_hash BYTEA PRIMARY KEY, + key_id TEXT NOT NULL REFERENCES api_keys(id), + -- The sealed key. Emptied the moment it is claimed, so a claimed link leaves + -- nothing behind even before the row is purged. + sealed BYTEA NOT NULL, + created_at BIGINT NOT NULL, + expires_at BIGINT NOT NULL, + claimed_at BIGINT +); +CREATE INDEX key_claims_by_key ON key_claims(key_id); +CREATE INDEX key_claims_by_expiry ON key_claims(expires_at); diff --git a/internal/store/migrations/postgres/0003_users.sql b/internal/store/migrations/postgres/0003_users.sql new file mode 100644 index 00000000..b0fea479 --- /dev/null +++ b/internal/store/migrations/postgres/0003_users.sql @@ -0,0 +1,56 @@ +-- Ported from the SQLite migration of the same number. Same version, same +-- meaning, different spelling: SQLite's storage classes and PostgreSQL's types +-- do not line up, and one shared file would be a promise neither could keep. +-- +-- Times are milliseconds since the epoch in BIGINT, exactly as in SQLite, +-- rather than TIMESTAMPTZ. The Go code reads and writes int64 milliseconds +-- everywhere; converting here would mean two time representations in one +-- codebase, and the conversion would be the thing that is wrong at 2am. + +-- Accounts, so a person can sign up, be approved, and then issue their own keys +-- without an operator running a command for each one. +CREATE TABLE users ( + id TEXT PRIMARY KEY, + -- Lowercased at the boundary. Two accounts differing only in case would be + -- two accounts to the database and one to every human who looks at them. + email TEXT NOT NULL UNIQUE, + name TEXT NOT NULL, + -- argon2id, encoded with its own parameters so the cost can be raised later + -- without invalidating existing hashes. Slow hashing is right here and wrong + -- for API keys: this input is chosen by a human and has a dictionary. + password_hash TEXT NOT NULL, + -- 'pending' until an operator approves, then 'active'. 'suspended' keeps the + -- account and stops it working. + state TEXT NOT NULL, + -- Assigned at approval: which project this person's keys belong to, and + -- therefore which limits they inherit. + project_id TEXT REFERENCES projects(id), + created_at BIGINT NOT NULL, + approved_at BIGINT, + approved_by TEXT, + last_login_at BIGINT +); +CREATE INDEX users_by_state ON users(state); + +-- Server-side sessions rather than signed cookies, so signing out and +-- suspending an account take effect immediately instead of whenever a token +-- happens to expire. +CREATE TABLE sessions ( + token_hash BYTEA PRIMARY KEY, + user_id TEXT NOT NULL REFERENCES users(id), + -- Per session, submitted with every form. A cookie alone authenticates the + -- browser, not the page that asked, which is what makes a cross-site POST + -- work. + csrf_hash BYTEA NOT NULL, + created_at BIGINT NOT NULL, + expires_at BIGINT NOT NULL, + last_seen BIGINT NOT NULL +); +CREATE INDEX sessions_by_user ON sessions(user_id); +CREATE INDEX sessions_by_expiry ON sessions(expires_at); + +-- A key issued through the profile page belongs to the person who made it, so +-- they can see and revoke their own and nobody else's. NULL means a key an +-- operator issued directly. +ALTER TABLE api_keys ADD COLUMN user_id TEXT REFERENCES users(id); +CREATE INDEX api_keys_by_user ON api_keys(user_id); diff --git a/internal/store/migrations/postgres/0004_admins.sql b/internal/store/migrations/postgres/0004_admins.sql new file mode 100644 index 00000000..eaa58a5c --- /dev/null +++ b/internal/store/migrations/postgres/0004_admins.sql @@ -0,0 +1,20 @@ +-- Ported from the SQLite migration of the same number. Same version, same +-- meaning, different spelling: SQLite's storage classes and PostgreSQL's types +-- do not line up, and one shared file would be a promise neither could keep. +-- +-- Times are milliseconds since the epoch in BIGINT, exactly as in SQLite, +-- rather than TIMESTAMPTZ. The Go code reads and writes int64 milliseconds +-- everywhere; converting here would mean two time representations in one +-- codebase, and the conversion would be the thing that is wrong at 2am. + +-- Who may see the cluster. +-- +-- Until now the answer was "whoever has a shell on the control plane", because +-- the dashboard lived on a loopback-bound port. Serving it over the web needs +-- an actual answer, and this is the smallest one that is not a lie: a flag on +-- an account that already has a password, a session and a CSRF token. +-- +-- Deliberately not a role system. One bit is enough for "can see the cluster" +-- and impossible to get subtly wrong; roles arrive when there is a second +-- distinction worth drawing. +ALTER TABLE users ADD COLUMN is_admin BIGINT NOT NULL DEFAULT 0; diff --git a/internal/store/migrations/postgres/0005_models.sql b/internal/store/migrations/postgres/0005_models.sql new file mode 100644 index 00000000..3ed04ed3 --- /dev/null +++ b/internal/store/migrations/postgres/0005_models.sql @@ -0,0 +1,34 @@ +-- Ported from the SQLite migration of the same number. Same version, same +-- meaning, different spelling: SQLite's storage classes and PostgreSQL's types +-- do not line up, and one shared file would be a promise neither could keep. +-- +-- Times are milliseconds since the epoch in BIGINT, exactly as in SQLite, +-- rather than TIMESTAMPTZ. The Go code reads and writes int64 milliseconds +-- everywhere; converting here would mean two time representations in one +-- codebase, and the conversion would be the thing that is wrong at 2am. + +-- Models added while the cluster is running. +-- +-- The catalog came from a YAML file read once at startup, so adding a model +-- meant editing a file on the control plane and restarting it: a settling +-- window, adoption churn, and on one occasion two minutes of outage. That is a +-- reasonable way to manage a fixed production set and a bad one for a platform +-- whose promise is that people can try things. +-- +-- These live alongside the ones in configuration rather than replacing them. +-- Config-as-code stays declarative — a model in YAML is in the catalog on every +-- start, whatever the database says — and this table holds the ones added +-- afterwards. A model added here may not take a name or alias that +-- configuration already uses, which the catalog's own constructor enforces for +-- free, because it refuses a duplicate whatever the source. +-- +-- The definition is stored as JSON rather than as columns. A profile is a +-- nested, evolving shape, and every field added to it would otherwise be a +-- migration; the catalog already knows how to validate one, and a row that +-- fails validation is reported and skipped rather than blocking startup. +CREATE TABLE models ( + name TEXT PRIMARY KEY, + definition TEXT NOT NULL, + created_at TIMESTAMP NOT NULL, + created_by TEXT NOT NULL DEFAULT '' +); diff --git a/internal/store/migrations/postgres/0006_threads.sql b/internal/store/migrations/postgres/0006_threads.sql new file mode 100644 index 00000000..b6a5879b --- /dev/null +++ b/internal/store/migrations/postgres/0006_threads.sql @@ -0,0 +1,57 @@ +-- Ported from the SQLite migration of the same number. Same version, same +-- meaning, different spelling: SQLite's storage classes and PostgreSQL's types +-- do not line up, and one shared file would be a promise neither could keep. +-- +-- Times are milliseconds since the epoch in BIGINT, exactly as in SQLite, +-- rather than TIMESTAMPTZ. The Go code reads and writes int64 milliseconds +-- everywhere; converting here would mean two time representations in one +-- codebase, and the conversion would be the thing that is wrong at 2am. + +-- Server-side conversations, so a thread can continue on a different model +-- than it started on. +-- +-- Off unless a project turns it on. Until this migration ran, llmgw stored no +-- message content anywhere — usage rows carry token counts and costs, never +-- text — and that stays true for every project that does not ask. An on-prem +-- gateway keeping nothing anyone typed is much easier to put in front of a +-- team than one keeping everything, so storing is a choice a project makes +-- rather than a default it discovers. + +-- retention_hours is the switch and the limit at once. NULL means this project +-- stores nothing, which is every project that existed before this and every +-- one created since without asking. +-- +-- One column rather than a boolean and a duration: a thread store with no +-- expiry is a transcript of everything a team ever typed, kept forever by +-- omission. Turning it on therefore requires saying for how long. +ALTER TABLE projects ADD COLUMN thread_retention_hours BIGINT; + +CREATE TABLE threads ( + id TEXT PRIMARY KEY, + -- Scoped to a project, and checked on every read. Thread ids are opaque, but + -- an unguessable id is a secret and secrets leak into logs, shell history + -- and bug reports; the scope is what makes a leaked id useless to someone + -- outside the project. + project_id TEXT NOT NULL REFERENCES projects(id) ON DELETE CASCADE, + created_at BIGINT NOT NULL, + -- Touched on every append, so retention measures idleness rather than age: + -- expiring a conversation someone is still having would be a strange way to + -- keep a promise about storage. + updated_at BIGINT NOT NULL +); +CREATE INDEX threads_by_project ON threads(project_id); +CREATE INDEX threads_by_updated ON threads(updated_at); + +CREATE TABLE thread_messages ( + thread_id TEXT NOT NULL REFERENCES threads(id) ON DELETE CASCADE, + -- Position rather than a timestamp, because order is the point and two + -- messages can share a millisecond. + seq BIGINT NOT NULL, + role TEXT NOT NULL, + content TEXT NOT NULL, + -- Which model produced an assistant turn. A thread that moved between models + -- otherwise reads as one voice, and when the answers disagree that is the + -- first thing anyone wants to know. + model TEXT, + PRIMARY KEY (thread_id, seq) +); diff --git a/internal/store/migrations/postgres/0007_memory.sql b/internal/store/migrations/postgres/0007_memory.sql new file mode 100644 index 00000000..49f0541f --- /dev/null +++ b/internal/store/migrations/postgres/0007_memory.sql @@ -0,0 +1,41 @@ +-- Ported from the SQLite migration of the same number, and the only one that +-- is not a mechanical translation. +-- +-- SQLite gets full-text search from an FTS5 virtual table with an external +-- content table and three triggers keeping it in step. PostgreSQL derives a +-- tsvector from the row itself, so there is nothing to keep in step and no +-- trigger to forget: the column cannot disagree with the row it is computed +-- from. +-- +-- The search query differs too, and that difference lives in the Go code +-- rather than being hidden here — a query that means "match these words" is +-- spelled MATCH in one and @@ in the other, and pretending otherwise would +-- cost more than admitting it. + +ALTER TABLE projects ADD COLUMN memory_retention_hours BIGINT; + +CREATE TABLE memory_notes ( + id BIGINT GENERATED ALWAYS AS IDENTITY PRIMARY KEY, + -- Scoped to a project and checked in SQL on every read, like threads. + project_id TEXT NOT NULL REFERENCES projects(id) ON DELETE CASCADE, + -- A namespace inside the project, so two agents sharing a key do not read + -- each other's notes. + agent_id TEXT NOT NULL, + key TEXT NOT NULL, + text TEXT NOT NULL, + -- Normalised to ",a,b," so a tag matches whole, never as part of another. + tags TEXT NOT NULL DEFAULT '', + created_at BIGINT NOT NULL, + -- Touched on every write, so retention measures how long a note has gone + -- unrefreshed rather than how old it is. + updated_at BIGINT NOT NULL, + -- 'simple' rather than a language configuration: these are an agent's own + -- notes, in whatever language it chose, and stemming English rules over + -- Turkish text would quietly stop finding things. + search tsvector GENERATED ALWAYS AS ( + to_tsvector('simple', coalesce(key, '') || ' ' || coalesce(text, '') || ' ' || coalesce(tags, '')) + ) STORED, + UNIQUE (project_id, agent_id, key) +); +CREATE INDEX memory_notes_by_updated ON memory_notes(updated_at); +CREATE INDEX memory_notes_search ON memory_notes USING GIN (search); diff --git a/internal/store/open.go b/internal/store/open.go new file mode 100644 index 00000000..cb4132f1 --- /dev/null +++ b/internal/store/open.go @@ -0,0 +1,61 @@ +package store + +import ( + "context" + "fmt" + "os" + "strings" +) + +// Open picks a backend from what the address looks like. +// +// A path is SQLite and a postgres:// URL is PostgreSQL, so an existing +// configuration keeps working untouched and moving to PostgreSQL is one line +// changed rather than a new setting to discover. A gateway on one machine has +// no reason to run a database server, and a gateway that has outgrown that +// should not have to be reconfigured twice. +// +// 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) { + address, err := resolveAddress(address) + if err != nil { + return nil, err + } + if isPostgres(address) { + return OpenPostgres(ctx, address) + } + return OpenSQLite(ctx, address) +} + +// isPostgres reports whether an address names a PostgreSQL server. +// +// By scheme, because that is the only part of a DSN that is reliably there. +// A bare host:port would be ambiguous with a relative file path, so it is not +// accepted: an address this cannot classify should fail loudly rather than +// quietly create a SQLite file named after someone's database server. +func isPostgres(address string) bool { + return strings.HasPrefix(address, "postgres://") || + strings.HasPrefix(address, "postgresql://") +} + +// resolveAddress expands an @file indirection. +// +// The file is read whole and trimmed, so a trailing newline — which every +// editor and every `echo` adds — does not become part of a password. +func resolveAddress(address string) (string, error) { + if !strings.HasPrefix(address, "@") { + return address, nil + } + path := strings.TrimPrefix(address, "@") + b, err := os.ReadFile(path) + if err != nil { + return "", fmt.Errorf("reading the store address from %s: %w", path, err) + } + out := strings.TrimSpace(string(b)) + if out == "" { + return "", fmt.Errorf("%s is empty, so there is no store address in it", path) + } + return out, nil +} diff --git a/internal/store/open_test.go b/internal/store/open_test.go new file mode 100644 index 00000000..248d0cfb --- /dev/null +++ b/internal/store/open_test.go @@ -0,0 +1,84 @@ +package store + +import ( + "context" + "os" + "path/filepath" + "strings" + "testing" +) + +func TestAPathIsStillSQLite(t *testing.T) { + // Every existing configuration names a file, and every one of them has to + // keep working without being touched. + for _, addr := range []string{ + "/var/lib/llmgw/llmgw.db", "llmgw.db", "./data/llmgw.db", + } { + if isPostgres(addr) { + t.Errorf("%q was read as a PostgreSQL address", addr) + } + } +} + +func TestAPostgresURLIsRecognised(t *testing.T) { + for _, addr := range []string{ + "postgres://u:p@h:5432/db", "postgresql://u@h/db?sslmode=disable", + } { + if !isPostgres(addr) { + t.Errorf("%q was not read as a PostgreSQL address", addr) + } + } +} + +func TestAnAddressFileKeepsThePasswordOutOfTheConfig(t *testing.T) { + // A DSN carries a password, and a configuration file gets copied, diffed + // and pasted into issues. The indirection is what lets the secret live + // somewhere with its own permissions. + dir := t.TempDir() + p := filepath.Join(dir, "dsn") + // With the trailing newline every editor and every echo adds, which must + // not become part of the password. + if err := os.WriteFile(p, []byte("postgres://u:secret@h/db\n"), 0o600); err != nil { + t.Fatal(err) + } + + got, err := resolveAddress("@" + p) + if err != nil { + t.Fatal(err) + } + if got != "postgres://u:secret@h/db" { + t.Errorf("resolved to %q", got) + } +} + +func TestAMissingOrEmptyAddressFileSaysSo(t *testing.T) { + // Rather than falling back to creating a SQLite file named "@/etc/...", + // which is what a silent failure here would produce. + if _, err := resolveAddress("@/nonexistent/dsn"); err == nil { + t.Error("a missing address file was accepted") + } + dir := t.TempDir() + p := filepath.Join(dir, "empty") + if err := os.WriteFile(p, []byte(" \n"), 0o600); err != nil { + t.Fatal(err) + } + _, err := resolveAddress("@" + p) + if err == nil || !strings.Contains(err.Error(), "empty") { + t.Errorf("an empty address file gave %v", err) + } +} + +func TestOpenStillCreatesASQLiteFile(t *testing.T) { + p := filepath.Join(t.TempDir(), "llmgw.db") + s, err := Open(context.Background(), p) + if err != nil { + t.Fatal(err) + } + defer s.Close() + if s.dialect != SQLiteDialect { + t.Errorf("a path opened as %q", s.dialect) + } + if _, err := os.Stat(p); err != nil { + t.Errorf("no database file was created: %v", err) + } +} diff --git a/internal/store/postgres.go b/internal/store/postgres.go new file mode 100644 index 00000000..934f6200 --- /dev/null +++ b/internal/store/postgres.go @@ -0,0 +1,63 @@ +package store + +import ( + "context" + "database/sql" + "fmt" + "time" + + // The pure-Go PostgreSQL driver. Pure-Go is not a preference here: the + // whole binary is built with CGO_ENABLED=0 so that it runs on a node whose + // libc is older than the builder's, and a cgo driver would take that away + // from every deployment, including the ones still on SQLite. + _ "github.com/jackc/pgx/v5/stdlib" +) + +// OpenPostgres opens the store against PostgreSQL. +// +// The same queries run here as on SQLite; only the placeholders differ, and +// the dialect handles that. What differs operationally is concurrency: SQLite +// is held to one connection because more would only produce lock contention +// that reads as a deadlock, while PostgreSQL is the reason someone moves to +// it and wants a pool. +// +// The DSN carries a password, so it is taken as a string the caller read from +// 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) { + if dsn == "" { + return nil, fmt.Errorf("store: no PostgreSQL address configured") + } + db, err := sql.Open("pgx", dsn) + if err != nil { + // Deliberately not wrapped with the DSN: it holds a password, and an + // error string ends up in logs, tickets and screenshots. + return nil, fmt.Errorf("opening the PostgreSQL store: %w", err) + } + + // Sized for a control plane rather than for a fleet of application + // servers. The request path barely touches the store — usage is batched, + // budgets are counted in memory, authentication is cached — so the pool + // exists for conversation turns and reporting queries, not for inference. + db.SetMaxOpenConns(16) + db.SetMaxIdleConns(4) + // Long enough that a quiet cluster keeps its connections, short enough + // that a restarted database is not answered with dead ones for an hour. + db.SetConnMaxLifetime(30 * time.Minute) + db.SetConnMaxIdleTime(5 * time.Minute) + + if err := db.PingContext(ctx); err != nil { + _ = db.Close() + return nil, fmt.Errorf("reaching the PostgreSQL store: %w", err) + } + if err := migrate(ctx, db, string(PostgresDialect)); err != nil { + _ = db.Close() + return nil, err + } + return &SQLite{ + raw: db, + dialect: PostgresDialect, + db: bound{c: PostgresDialect, to: db}, + }, nil +} diff --git a/internal/store/postgres_test.go b/internal/store/postgres_test.go new file mode 100644 index 00000000..1b72d153 --- /dev/null +++ b/internal/store/postgres_test.go @@ -0,0 +1,173 @@ +package store + +import ( + "context" + "errors" + "os" + "testing" + "time" +) + +// openPG opens the test database named by LLMGW_TEST_POSTGRES, or skips. +// +// Env-gated rather than spun up here, because a store test that needs a +// container is a test nobody runs. Pointed at a database, it runs the same +// assertions the SQLite tests make — which is the only thing that can catch +// the two backends drifting apart. +// +// 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 { + t.Helper() + dsn := os.Getenv("LLMGW_TEST_POSTGRES") + if dsn == "" { + t.Skip("set LLMGW_TEST_POSTGRES to a database DSN to run these") + } + s, err := OpenPostgres(context.Background(), dsn) + if err != nil { + t.Fatalf("opening the test database: %v", err) + } + t.Cleanup(func() { _ = s.Close() }) + + // CASCADE because the tables reference each other, and TRUNCATE rather + // than DELETE so identity columns restart too. + _, err = s.db.ExecContext(context.Background(), + `TRUNCATE thread_messages, threads, memory_notes, usage, api_keys, + key_claims, limits, models, sessions, users, projects, teams, + organizations RESTART IDENTITY CASCADE`) + if err != nil { + t.Fatalf("clearing the test database: %v", err) + } + return s +} + +func TestPostgresCarriesTheWholeTenancyTree(t *testing.T) { + ctx := context.Background() + s := openPG(t) + + proj := tree(t, s) + secret, k := mustKey(t, s, "k1", proj) + _ = secret + + got, err := s.KeyByHash(ctx, k.Hash) + if err != nil { + t.Fatalf("reading the key back: %v", err) + } + if got.ID != k.ID || got.ProjectID != proj { + t.Errorf("key came back as %+v", got) + } +} + +func TestPostgresKeepsUsageAndSpend(t *testing.T) { + // The accounting path, which is the one that decides whether a budget + // means anything after a restart. + ctx := context.Background() + s := openPG(t) + proj := tree(t, s) + _, k := mustKey(t, s, "k1", proj) + + now := time.Now() + err := s.WriteUsage(ctx, []Usage{ + {At: now, KeyID: k.ID, ProjectID: proj, Model: "m", ServedModel: "m", + PromptTokens: 10, CompletionTokens: 20, CostMicros: 500, Status: 200, LatencyMs: 42}, + {At: now, KeyID: k.ID, ProjectID: proj, Model: "m", ServedModel: "m", + PromptTokens: 1, CompletionTokens: 2, CostMicros: 250, Status: 200, LatencyMs: 7}, + }) + if err != nil { + t.Fatalf("writing usage: %v", err) + } + + spend, err := s.SpendSince(ctx, now.Add(-time.Hour)) + if err != nil { + t.Fatal(err) + } + if spend[k.ID] != 750 { + t.Errorf("spend = %d, want 750", spend[k.ID]) + } +} + +func TestPostgresKeepsThreadsInOrderAndScoped(t *testing.T) { + // The newest tables, and the ones with the most SQL per operation — + // a transaction, a MAX(seq), an insert per turn. + ctx := context.Background() + s := openPG(t) + proj := tree(t, s) + if err := s.SetThreadRetention(ctx, proj, 24); err != nil { + t.Fatal(err) + } + + for _, m := range []Message{ + {Role: "user", Content: "first"}, + {Role: "assistant", Content: "answer", Model: "small"}, + {Role: "user", Content: "second"}, + } { + if err := s.AppendThread(ctx, proj, "th1", []Message{m}); err != nil { + t.Fatal(err) + } + } + + got, err := s.Thread(ctx, proj, "th1") + if err != nil { + t.Fatal(err) + } + if len(got) != 3 || got[0].Content != "first" || got[2].Content != "second" { + t.Fatalf("thread came back as %+v", got) + } + if got[1].Model != "small" { + t.Errorf("the model that answered was lost: %+v", got[1]) + } + + // And the scope holds here exactly as it does on SQLite. + if err := s.PutProject(ctx, Project{ID: "proj2", TeamID: "team1", Name: "Other"}); err != nil { + t.Fatal(err) + } + if _, err := s.Thread(ctx, "proj2", "th1"); !errors.Is(err, ErrNoThread) { + t.Errorf("reading across projects returned %v", err) + } +} + +func TestPostgresEnforcesLimitsTheSameWay(t *testing.T) { + ctx := context.Background() + s := openPG(t) + proj := tree(t, s) + _, k := mustKey(t, s, "k1", proj) + + rpm := 7 + budget := int64(9000) + if err := s.PutLimits(ctx, Limits{ + Scope: ScopeProject, ScopeID: proj, RPM: &rpm, + Budget: &budget, BudgetPeriod: PeriodMonth, + }); err != nil { + t.Fatal(err) + } + + eff, err := s.Resolve(ctx, k.ID) + if err != nil { + t.Fatal(err) + } + if eff.RPM == nil || *eff.RPM != 7 { + t.Errorf("rpm did not resolve down the tree: %+v", eff.RPM) + } + if eff.Budget == nil || *eff.Budget != 9000 { + t.Errorf("budget did not resolve down the tree: %+v", eff.Budget) + } +} + +func TestPostgresUpsertsRatherThanDuplicating(t *testing.T) { + // ON CONFLICT is the one non-trivial construct shared by both dialects. + // If it behaved differently here, applying a configuration twice would + // start failing instead of being a no-op. + ctx := context.Background() + s := openPG(t) + tree(t, s) + tree(t, s) + + orgs, err := s.Orgs(ctx) + if err != nil { + t.Fatal(err) + } + if len(orgs) != 1 { + t.Errorf("applying the tree twice produced %d orgs", len(orgs)) + } +}