Add the usage table in user schema migration 9 -> 10, zero the stats_* columns

Claude-Session: https://claude.ai/code/session_01DeXuhSjq5Gta5qPJgSceFe

user/manager.go

@@ -1959,6 +1959,12 @@ func (a *Manager) maybeProvisionTokens(tx *sql.Tx, provisionUsernames []string,
 	return nil
 }
 
+// DB returns the manager's database. It also holds the usage table, which package limit reads and
+// writes (see the usage table in the schema files).
+func (a *Manager) DB() *db.DB {
+	return a.db
+}
+
 // Close stops background goroutines and closes the underlying database.
 func (a *Manager) Close() error {
 	select {

user/manager_postgres_schema.go

@@ -123,14 +123,38 @@ const (
 	`
 )
 
+// The usage table holds the daily usage counts, one row per key, reused across the key's 24h
+// periods (see package limit, which reads and writes it). It lives in the user schema so it shares
+// its migrations; in a cluster, every node uses it.
+const (
+	postgresCreateUsageTableQuery = `
+		CREATE TABLE IF NOT EXISTS usage (
+			key TEXT NOT NULL PRIMARY KEY,
+			period_start BIGINT NOT NULL,
+			messages BIGINT NOT NULL DEFAULT 0,
+			emails BIGINT NOT NULL DEFAULT 0,
+			calls BIGINT NOT NULL DEFAULT 0,
+			updated_at BIGINT NOT NULL
+		);
+		CREATE INDEX IF NOT EXISTS idx_usage_updated_at ON usage (updated_at);
+	`
+
+	// 9 -> 10: The usage table replaces the stats_* columns, which stay in the schema but are
+	// zeroed, so nothing reads their stale numbers
+	postgresMigrate9To10UpdateQueries = postgresCreateUsageTableQuery + `
+		UPDATE "user" SET stats_messages = 0, stats_emails = 0, stats_calls = 0;
+	`
+)
+
 var (
-	postgresCreateTables = schema.AsMigrateFunc(postgresCreateTablesQueries)
+	postgresCreateTables = schema.AsMigrateFunc(postgresCreateTablesQueries + postgresCreateUsageTableQuery)
 
 	// postgresMigrations maps a schema version to the migration upgrading it to the next
 	// version. Always append migrations at the end, never insert in the middle.
 	postgresMigrations = schema.Migrations{
-		{From: 6, To: 7}: schema.Query(postgresMigrate6To7UpdateQueries),
-		{From: 7, To: 8}: schema.Query(postgresMigrate7To8UpdateQueries),
-		{From: 8, To: 9}: schema.Nop, // 8 -> 9 repairs a SQLite-only foreign key defect; nothing to do on Postgres
+		{From: 6, To: 7}:  schema.Query(postgresMigrate6To7UpdateQueries),
+		{From: 7, To: 8}:  schema.Query(postgresMigrate7To8UpdateQueries),
+		{From: 8, To: 9}:  schema.Nop, // 8 -> 9 repairs a SQLite-only foreign key defect; nothing to do on Postgres
+		{From: 9, To: 10}: schema.Query(postgresMigrate9To10UpdateQueries),
 	}
 )

user/manager_sqlite_schema.go

@@ -364,20 +364,44 @@ const (
 	`
 )
 
+// The usage table holds the daily usage counts, one row per key, reused across the key's 24h
+// periods (see package limit, which reads and writes it). It lives here so a single-node server
+// keeps it in its user database, under the user schema's migrations.
+const (
+	sqliteCreateUsageTableQuery = `
+		CREATE TABLE IF NOT EXISTS usage (
+			key TEXT NOT NULL PRIMARY KEY,
+			period_start INT NOT NULL,
+			messages INT NOT NULL DEFAULT 0,
+			emails INT NOT NULL DEFAULT 0,
+			calls INT NOT NULL DEFAULT 0,
+			updated_at INT NOT NULL
+		);
+		CREATE INDEX IF NOT EXISTS idx_usage_updated_at ON usage (updated_at);
+	`
+
+	// 9 -> 10: The usage table replaces the stats_* columns, which stay in the schema but are
+	// zeroed, so nothing reads their stale numbers
+	sqliteMigrate9To10UpdateQueries = sqliteCreateUsageTableQuery + `
+		UPDATE user SET stats_messages = 0, stats_emails = 0, stats_calls = 0;
+	`
+)
+
 var (
-	sqliteCreateTables = schema.AsMigrateFunc(sqliteCreateTablesQueries)
+	sqliteCreateTables = schema.AsMigrateFunc(sqliteCreateTablesQueries + sqliteCreateUsageTableQuery)
 
 	// sqliteMigrations maps a schema version to the migration upgrading it to the next
 	// version. Always append migrations at the end, never insert in the middle.
 	sqliteMigrations = schema.Migrations{
-		{From: 1, To: 2}: schema.Func(sqliteMigrateFrom1),
-		{From: 2, To: 3}: schema.Query(sqliteMigrate2To3UpdateQueries),
-		{From: 3, To: 4}: schema.Query(sqliteMigrate3To4UpdateQueries),
-		{From: 4, To: 5}: schema.Query(sqliteMigrate4To5UpdateQueries),
-		{From: 5, To: 6}: schema.Query(sqliteMigrate5To6UpdateQueries),
-		{From: 6, To: 7}: schema.Query(sqliteMigrate6To7UpdateQueries),
-		{From: 7, To: 8}: schema.Query(sqliteMigrate7To8UpdateQueries),
-		{From: 8, To: 9}: schema.Query(sqliteMigrate8To9UpdateQueries),
+		{From: 1, To: 2}:  schema.Func(sqliteMigrateFrom1),
+		{From: 2, To: 3}:  schema.Query(sqliteMigrate2To3UpdateQueries),
+		{From: 3, To: 4}:  schema.Query(sqliteMigrate3To4UpdateQueries),
+		{From: 4, To: 5}:  schema.Query(sqliteMigrate4To5UpdateQueries),
+		{From: 5, To: 6}:  schema.Query(sqliteMigrate5To6UpdateQueries),
+		{From: 6, To: 7}:  schema.Query(sqliteMigrate6To7UpdateQueries),
+		{From: 7, To: 8}:  schema.Query(sqliteMigrate7To8UpdateQueries),
+		{From: 8, To: 9}:  schema.Query(sqliteMigrate8To9UpdateQueries),
+		{From: 9, To: 10}: schema.Query(sqliteMigrate9To10UpdateQueries),
 	}
 )
 

user/manager_test.go

@@ -1909,7 +1909,7 @@ func TestMigrationFrom6Postgres(t *testing.T) {
 	require.Nil(t, err)
 	var version int
 	require.Nil(t, testDB.QueryRow(`SELECT version FROM schema_version WHERE store = 'user'`).Scan(&version))
-	require.Equal(t, 9, version)
+	require.Equal(t, 10, version)
 	// The manager works against the migrated schema
 	require.Nil(t, a.AddUser("phil", "mypass", RoleUser, false))
 	u, err := a.User("phil")
@@ -1951,7 +1951,7 @@ func checkSchemaVersion(t *testing.T, d *db.DB) {
 
 	var schemaVersion int
 	require.Nil(t, rows.Scan(&schemaVersion))
-	require.Equal(t, 9, schemaVersion)
+	require.Equal(t, 10, schemaVersion)
 	require.Nil(t, rows.Close())
 }
 
@@ -3732,3 +3732,55 @@ func TestManager_EnqueueStats_DeltasAccumulate(t *testing.T) {
 		require.Equal(t, int64(4), u.Stats.Calls)
 	})
 }
+
+func TestMigrationFrom9To10SQLite(t *testing.T) {
+	filename := filepath.Join(t.TempDir(), "user.db")
+	a := newTestManagerFromFile(t, filename, "", PermissionDenyAll, bcrypt.MinCost, DefaultUserStatsQueueWriterInterval)
+	testMigrationFrom9To10(t, a, `UPDATE schemaVersion SET version = 9 WHERE id = 1`, `UPDATE user SET stats_messages = 5, stats_emails = 2, stats_calls = 1`)
+	a = newTestManagerFromFile(t, filename, "", PermissionDenyAll, bcrypt.MinCost, DefaultUserStatsQueueWriterInterval)
+	checkMigrated9To10(t, a, `SELECT version FROM schemaVersion WHERE id = 1`, `SELECT stats_messages, stats_emails, stats_calls FROM user`)
+	checkMigratedSqliteSchema(t, filename)
+}
+
+func TestMigrationFrom9To10Postgres(t *testing.T) {
+	schemaDSN := dbtest.CreateTestPostgresSchema(t)
+	open := func() *Manager {
+		host, err := pg.Open(schemaDSN)
+		require.Nil(t, err)
+		a, err := NewPostgresManager(db.New(host, nil), &Config{DefaultAccess: PermissionDenyAll, BcryptCost: bcrypt.MinCost, QueueWriterInterval: DefaultUserStatsQueueWriterInterval})
+		require.Nil(t, err)
+		t.Cleanup(func() { a.Close() })
+		return a
+	}
+	testMigrationFrom9To10(t, open(), `UPDATE schema_version SET version = 9 WHERE store = 'user'`, `UPDATE "user" SET stats_messages = 5, stats_emails = 2, stats_calls = 1`)
+	checkMigrated9To10(t, open(), `SELECT version FROM schema_version WHERE store = 'user'`, `SELECT stats_messages, stats_emails, stats_calls FROM "user"`)
+}
+
+// testMigrationFrom9To10 turns the manager's database back into a version 9 one with stats, and closes it
+func testMigrationFrom9To10(t *testing.T, a *Manager, setVersionQuery, setStatsQuery string) {
+	t.Helper()
+	require.Nil(t, a.AddUser("phil", "phil", RoleUser, false))
+	_, err := a.DB().Exec(`DROP TABLE IF EXISTS usage`)
+	require.Nil(t, err)
+	_, err = a.DB().Exec(setVersionQuery)
+	require.Nil(t, err)
+	_, err = a.DB().Exec(setStatsQuery)
+	require.Nil(t, err)
+	require.Nil(t, a.Close())
+}
+
+// checkMigrated9To10 checks that the migration zeroed the stats and created the usage table
+func checkMigrated9To10(t *testing.T, a *Manager, selectVersionQuery, selectStatsQuery string) {
+	t.Helper()
+	var version int
+	require.Nil(t, a.DB().QueryRow(selectVersionQuery).Scan(&version))
+	require.Equal(t, 10, version)
+	var messages, emails, calls int64
+	require.Nil(t, a.DB().QueryRow(selectStatsQuery).Scan(&messages, &emails, &calls))
+	require.Equal(t, []int64{0, 0, 0}, []int64{messages, emails, calls})
+	_, err := a.DB().Exec(`INSERT INTO usage (key, period_start, messages, emails, calls, updated_at) VALUES ('ip:1.2.3.4', 1793344380, 37, 0, 0, 1793380805)`)
+	require.Nil(t, err)
+	var count int64
+	require.Nil(t, a.DB().QueryRow(`SELECT messages FROM usage WHERE key = 'ip:1.2.3.4' AND updated_at >= 1793380805`).Scan(&count))
+	require.Equal(t, int64(37), count)
+}