Stop calling the PostgreSQL store SQLite, and stop guessing where it is
Two things that were only ever confusing, and one of them stopped a command working on the cluster. The store type was named SQLite and kept the name when PostgreSQL was added, so OpenPostgres returned a *SQLite and all seventy-odd methods on the production store were spelled after the database it was not running on. Nothing broke; people did. Earlier today I added three methods to *SQLite for a table that on this cluster only ever lives in PostgreSQL, and the name did not once look wrong. It is a field, which is all it ever was. `llmgw users` and `llmgw keys` defaulted --store to a SQLite path. On a deployment whose store is PostgreSQL every one of them failed with "unable to open database file" while the right answer sat in auth.store in the config file beside them; the conclusion an operator draws is that the tool is broken, and they are not far off. They now ask the server's own configuration before falling back, and the error names which of the three sources chose the address — the failure was never the path, it was which one picked it. The single-machine default still applies last, so a deployment with no config file keeps working with no flags. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
1 parent
3febb6ef63
commit
ebce41e22d
32 files changed
+320
-140
No files matched your search
@@ -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)
|
||||
|
||||
@@ -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"))
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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"))
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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"))
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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},
|
||||
|
||||
@@ -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 == "" {
|
||||
|
||||
@@ -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
|
||||
|
||||
+54
-34
@@ -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),
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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
|
||||
|
||||
+15
-15
@@ -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
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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"},
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in new issue
Block a user