atepg: Guardrail to disallow changes to atespace tables that limit future sharding

cmd/ateapi/internal/store/atepg/guardrails_test.go

@@ -0,0 +1,125 @@
+// Copyright 2026 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package atepg
+
+// Fixtures shared by the schema guardrail tests, TestActorsTablePartitionable
+// and TestAtespaceTablesShardable: a migrated copy of the schema to alter,
+// and the catalog queries that describe it.
+
+import (
+	"context"
+	"fmt"
+	"testing"
+
+	"github.com/jackc/pgx/v5"
+	"github.com/jackc/pgx/v5/pgxpool"
+)
+
+// migratedPool opens a pool on a fresh schema with the migrations applied.
+func migratedPool(t *testing.T, schema string) *pgxpool.Pool {
+	t.Helper()
+	admin := requirePool(t)
+	quoted := pgx.Identifier{schema}.Sanitize()
+	if _, err := admin.Exec(t.Context(), `DROP SCHEMA IF EXISTS `+quoted+` CASCADE; CREATE SCHEMA `+quoted); err != nil {
+		t.Fatal(err)
+	}
+	t.Cleanup(func() { _, _ = admin.Exec(context.Background(), `DROP SCHEMA IF EXISTS `+quoted+` CASCADE`) })
+	pool := openPool(t, schema, nil)
+	p, err := NewPersistence(t.Context(), pool)
+	if err != nil {
+		t.Fatal(err)
+	}
+	p.Close()
+	return pool
+}
+
+// openPool opens a pool on schema, tracing every statement with tracer.
+func openPool(t *testing.T, schema string, tracer pgx.QueryTracer) *pgxpool.Pool {
+	t.Helper()
+	cfg, err := pgxpool.ParseConfig(containerDSN)
+	if err != nil {
+		t.Fatal(err)
+	}
+	cfg.ConnConfig.RuntimeParams["search_path"] = pgx.Identifier{schema}.Sanitize()
+	cfg.ConnConfig.Tracer = tracer
+	pool, err := pgxpool.NewWithConfig(t.Context(), cfg)
+	if err != nil {
+		t.Fatal(err)
+	}
+	t.Cleanup(pool.Close)
+	return pool
+}
+
+// schemaTables lists the tables in the pool's schema, leaving out the
+// partitions of a partitioned table.
+func schemaTables(t *testing.T, pool *pgxpool.Pool) []string {
+	t.Helper()
+	return collectStrings(t, pool, `
+		SELECT c.relname FROM pg_class c
+		JOIN pg_namespace n ON n.oid = c.relnamespace
+		WHERE n.nspname = current_schema() AND c.relkind IN ('r', 'p') AND NOT c.relispartition
+		ORDER BY c.relname`)
+}
+
+// tablesWithColumn lists the tables in the pool's schema that have column.
+func tablesWithColumn(t *testing.T, pool *pgxpool.Pool, column string) []string {
+	t.Helper()
+	tables := collectStrings(t, pool, `
+		SELECT table_name FROM information_schema.columns
+		WHERE table_schema = current_schema() AND column_name = $1
+		ORDER BY table_name`, column)
+	if len(tables) == 0 {
+		t.Fatalf("no table has a %s column", column)
+	}
+	return tables
+}
+
+func collectStrings(t *testing.T, pool *pgxpool.Pool, sql string, args ...any) []string {
+	t.Helper()
+	rows, err := pool.Query(t.Context(), sql, args...)
+	if err != nil {
+		t.Fatal(err)
+	}
+	values, err := pgx.CollectRows(rows, pgx.RowTo[string])
+	if err != nil {
+		t.Fatal(err)
+	}
+	return values
+}
+
+// foreignKey is a foreign key from one table of the pool's schema to another.
+type foreignKey struct {
+	Name, From, To string
+	Definition     string // as pg_get_constraintdef renders it
+}
+
+// ddl is the statement that adds the foreign key to its table again.
+func (fk foreignKey) ddl() string {
+	return fmt.Sprintf("ALTER TABLE %s ADD CONSTRAINT %s %s", fk.From, pgx.Identifier{fk.Name}.Sanitize(), fk.Definition)
+}
+
+// foreignKeys lists the foreign keys of the pool's schema, by name.
+func foreignKeys(ctx context.Context, pool *pgxpool.Pool) ([]foreignKey, error) {
+	rows, err := pool.Query(ctx, `
+		SELECT conname, conrelid::regclass::text, confrelid::regclass::text, pg_get_constraintdef(c.oid)
+		FROM pg_constraint c
+		JOIN pg_namespace n ON n.oid = c.connamespace
+		WHERE n.nspname = current_schema() AND contype = 'f' AND conparentid = 0
+		ORDER BY conname`)
+	if err != nil {
+		return nil, err
+	}
+	return pgx.CollectRows(rows, pgx.RowToStructByPos[foreignKey])
+}

cmd/ateapi/internal/store/atepg/partition_test.go

@@ -155,14 +155,7 @@ const partitions = 9
 // hash-partitioned table on key, keeping its indexes, constraints and
 // foreign keys. PostgreSQL rejects any of them that omits key.
 func partitionTable(ctx context.Context, pool *pgxpool.Pool, table, key string) error {
-	rows, err := pool.Query(ctx, `
-		SELECT format('ALTER TABLE %s ADD CONSTRAINT %I %s', conrelid::regclass, conname, pg_get_constraintdef(oid))
-		FROM pg_constraint
-		WHERE contype = 'f' AND conparentid = 0 AND $1::regclass IN (conrelid, confrelid)`, table)
-	if err != nil {
-		return err
-	}
-	foreignKeys, err := pgx.CollectRows(rows, pgx.RowTo[string])
+	fks, err := foreignKeys(ctx, pool)
 	if err != nil {
 		return err
 	}
@@ -174,9 +167,14 @@ func partitionTable(ctx context.Context, pool *pgxpool.Pool, table, key string)
 	if _, err := pool.Exec(ctx, strings.Join(ddl, ";\n")); err != nil {
 		return fmt.Errorf("%s cannot be partitioned by %s: %w", table, key, err)
 	}
-	for _, fk := range foreignKeys {
-		if _, err := pool.Exec(ctx, fk); err != nil {
-			return fmt.Errorf("foreign key cannot reference %s partitioned by %s: %s: %w", table, key, fk, err)
+	// LIKE copies no foreign key, and DROP CASCADE removed the ones onto the
+	// table, so add back every one that touches it.
+	for _, fk := range fks {
+		if fk.From != table && fk.To != table {
+			continue
+		}
+		if _, err := pool.Exec(ctx, fk.ddl()); err != nil {
+			return fmt.Errorf("foreign key cannot reference %s partitioned by %s: %s: %w", table, key, fk.ddl(), err)
 		}
 	}
 	return nil
@@ -297,24 +295,6 @@ func normalizeSQL(sql string) string {
 	return strings.Join(strings.Fields(sql), " ")
 }
 
-// migratedPool opens a pool on a fresh schema with the migrations applied.
-func migratedPool(t *testing.T, schema string) *pgxpool.Pool {
-	t.Helper()
-	admin := requirePool(t)
-	quoted := pgx.Identifier{schema}.Sanitize()
-	if _, err := admin.Exec(t.Context(), `DROP SCHEMA IF EXISTS `+quoted+` CASCADE; CREATE SCHEMA `+quoted); err != nil {
-		t.Fatal(err)
-	}
-	t.Cleanup(func() { _, _ = admin.Exec(context.Background(), `DROP SCHEMA IF EXISTS `+quoted+` CASCADE`) })
-	pool := openPool(t, schema, nil)
-	p, err := NewPersistence(t.Context(), pool)
-	if err != nil {
-		t.Fatal(err)
-	}
-	p.Close()
-	return pool
-}
-
 // partitionedPool opens a pool on a fresh schema with tables partitioned on
 // key, and returns the tables it partitioned. A nil tables partitions every
 // table that has a column named key.
@@ -331,40 +311,3 @@ func partitionedPool(t *testing.T, schema, key string, tables []string) (*pgxpoo
 	}
 	return pool, tables
 }
-
-// tablesWithColumn lists the tables in the pool's schema that have column.
-func tablesWithColumn(t *testing.T, pool *pgxpool.Pool, column string) []string {
-	t.Helper()
-	rows, err := pool.Query(t.Context(), `
-		SELECT table_name FROM information_schema.columns
-		WHERE table_schema = current_schema() AND column_name = $1
-		ORDER BY table_name`, column)
-	if err != nil {
-		t.Fatal(err)
-	}
-	tables, err := pgx.CollectRows(rows, pgx.RowTo[string])
-	if err != nil {
-		t.Fatal(err)
-	}
-	if len(tables) == 0 {
-		t.Fatalf("no table has a %s column", column)
-	}
-	return tables
-}
-
-// openPool opens a pool on schema, tracing every statement with tracer.
-func openPool(t *testing.T, schema string, tracer pgx.QueryTracer) *pgxpool.Pool {
-	t.Helper()
-	cfg, err := pgxpool.ParseConfig(containerDSN)
-	if err != nil {
-		t.Fatal(err)
-	}
-	cfg.ConnConfig.RuntimeParams["search_path"] = pgx.Identifier{schema}.Sanitize()
-	cfg.ConnConfig.Tracer = tracer
-	pool, err := pgxpool.NewWithConfig(t.Context(), cfg)
-	if err != nil {
-		t.Fatal(err)
-	}
-	t.Cleanup(pool.Close)
-	return pool
-}

cmd/ateapi/internal/store/atepg/sharding_test.go

@@ -0,0 +1,105 @@
+// Copyright 2026 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package atepg
+
+import (
+	"fmt"
+	"slices"
+	"strings"
+	"testing"
+
+	"github.com/jackc/pgx/v5/pgxpool"
+)
+
+// TestAtespaceTablesShardable exists to keep it possible to move each
+// atespace's tables to a database of their own later. It fails on a foreign
+// key between those tables and the global ones, and on a new table that is
+// not classified as one or the other.
+func TestAtespaceTablesShardable(t *testing.T) {
+	// exemptions are foreign keys allowed to cross between an atespace's
+	// tables and the global ones. Do not add one without discussion and
+	// agreement in the community.
+	var exemptions []string
+
+	for _, v := range shardingViolations(t, migratedPool(t, "shardable"), exemptions) {
+		t.Error(v)
+	}
+
+	t.Run("rejects a foreign key to a global table", func(t *testing.T) {
+		pool := migratedPool(t, "shardable-fk")
+		if _, err := pool.Exec(t.Context(), `ALTER TABLE actors ADD COLUMN worker text REFERENCES workers (name)`); err != nil {
+			t.Fatal(err)
+		}
+		violations := shardingViolations(t, pool, nil)
+		if len(violations) != 1 || !strings.Contains(violations[0], "actors_worker_fkey") {
+			t.Fatalf("violations = %q, want exactly the actors to workers foreign key", violations)
+		}
+		t.Log(violations[0])
+	})
+	t.Run("classifies a table by its atespace column", func(t *testing.T) {
+		pool := migratedPool(t, "shardable-column")
+		if _, err := pool.Exec(t.Context(), `CREATE TABLE notes (atespace text, worker text REFERENCES workers (name))`); err != nil {
+			t.Fatal(err)
+		}
+		violations := shardingViolations(t, pool, nil)
+		if len(violations) != 1 || !strings.Contains(violations[0], "notes_worker_fkey") {
+			t.Fatalf("violations = %q, want exactly the notes to workers foreign key", violations)
+		}
+		t.Log(violations[0])
+	})
+	t.Run("rejects an unclassified table", func(t *testing.T) {
+		pool := migratedPool(t, "shardable-table")
+		if _, err := pool.Exec(t.Context(), `CREATE TABLE notes (id int PRIMARY KEY)`); err != nil {
+			t.Fatal(err)
+		}
+		violations := shardingViolations(t, pool, nil)
+		if len(violations) != 1 || !strings.Contains(violations[0], "table notes is neither") {
+			t.Fatalf("violations = %q, want exactly the unclassified notes table", violations)
+		}
+		t.Log(violations[0])
+	})
+}
+
+// globalTables hold state shared by every atespace. Every other table is an
+// atespace's: atespaces itself, and the tables with an atespace column, which
+// TestActorsTablePartitionable partitions by it.
+var globalTables = []string{"workers", "worker_assignments", "worker_outbox", "worker_outbox_trim", "leases", migrationTableName}
+
+// shardingViolations reports every table that is neither an atespace table
+// nor a global one, and every foreign key between the two sets other than
+// the exempt ones.
+func shardingViolations(t *testing.T, pool *pgxpool.Pool, exempt []string) []string {
+	t.Helper()
+	atespaceTables := append([]string{"atespaces"}, tablesWithColumn(t, pool, "atespace")...)
+	var violations []string
+	for _, table := range schemaTables(t, pool) {
+		if !slices.Contains(atespaceTables, table) && !slices.Contains(globalTables, table) {
+			violations = append(violations, fmt.Sprintf("table %s is neither an atespace table nor a global table; give it an atespace column or add it to globalTables in TestAtespaceTablesShardable", table))
+		}
+	}
+	fks, err := foreignKeys(t.Context(), pool)
+	if err != nil {
+		t.Fatal(err)
+	}
+	for _, fk := range fks {
+		if slices.Contains(exempt, fk.Name) {
+			continue
+		}
+		if slices.Contains(atespaceTables, fk.From) != slices.Contains(atespaceTables, fk.To) {
+			violations = append(violations, fmt.Sprintf("foreign key %s from %s to %s crosses between an atespace's tables and the global ones, so the atespace could not move to its own database; remove it or add it to the exemptions in TestAtespaceTablesShardable", fk.Name, fk.From, fk.To))
+		}
+	}
+	return violations
+}

docs/dev/postgresql-schema-evolution.md

@@ -35,6 +35,8 @@ To preserve the option to partition the `actors` table by `atespace` or by `name
 - A foreign key that references one of these tables by columns that omit `atespace`, or `actors` by anything other than `(atespace, name)`.
 - A query on one of these tables that does not filter on `atespace`, or a query on `actors` that does not also filter on `name`. A statement whose result spans every partition by definition, such as a global list, must be listed in `TestActorsTablePartitionable`. Anything else that reads more than one partition needs an exemption there, agreed with the community.
 
+To preserve the option to move each atespace's tables (those above, plus `atespaces`) to a database of their own, do not add a foreign key between them and the global tables (`workers`, `worker_assignments`, `worker_outbox`, `worker_outbox_trim`, `leases`). `TestAtespaceTablesShardable` fails on such a key, and on a new table until it is classified as one or the other: a table with an `atespace` column is an atespace's, and a global one must be listed in the test.
+
 ## Expand and contract
 
 Use an expand and contract sequence for a schema replacement or removal: