Run the store on PostgreSQL, chosen by what the address looks like
A path is still SQLite and a postgres:// URL is PostgreSQL, so every existing configuration keeps working untouched and moving is one line changed rather than a new setting to find. An address of the form @/path/to/file reads the real one out of that file, which is how a DSN carrying a password stays out of a configuration that gets copied, diffed and pasted into issues. The driver is pgx, which is pure Go. That is not a preference: the whole binary is built with CGO_ENABLED=0 so 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. Tested against a real PostgreSQL, which found two things nothing else could. The schema_migrations table is created in Go rather than in a migration file, so it escaped the port and still said INTEGER — 64 bits in SQLite, 32 in PostgreSQL, and a millisecond timestamp does not fit, so the very first migration failed to record that it had run. And the memory migration uses an FTS5 virtual table with three triggers keeping it in step, which has no counterpart: PostgreSQL derives a tsvector from the row, so there is nothing to keep in step and no trigger to forget. That leaves exactly one query that is not the same sentence in both dialects — the note search, MATCH against @@ — and it is written out twice rather than hidden behind a shared spelling that would mean something subtly different in each. Everything else is one set of queries. The text search uses the 'simple' configuration rather than English. These are an agent's own notes, in whatever language it chose, and stemming English rules over Turkish would quietly stop finding things. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Ilker ManapandClaude Opus 5 committed
1 parent
ccc9aa8b76
commit
a76ebdb79a
18 files changed
+750
-12
No files matched your search
@@ -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
|
||||
|
||||
@@ -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=
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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);
|
||||
@@ -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);
|
||||
@@ -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);
|
||||
@@ -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;
|
||||
@@ -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 ''
|
||||
);
|
||||
@@ -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)
|
||||
);
|
||||
@@ -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);
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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))
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user