diff --git a/cmd/ateapi/internal/store/atepg/guardrails_test.go b/cmd/ateapi/internal/store/atepg/guardrails_test.go new file mode 100644 index 0000000000..5445de9169 --- /dev/null +++ b/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]) +} diff --git a/cmd/ateapi/internal/store/atepg/partition_test.go b/cmd/ateapi/internal/store/atepg/partition_test.go index ea7f765546..49f3d5ded5 100644 --- a/cmd/ateapi/internal/store/atepg/partition_test.go +++ b/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 -} diff --git a/cmd/ateapi/internal/store/atepg/sharding_test.go b/cmd/ateapi/internal/store/atepg/sharding_test.go new file mode 100644 index 0000000000..17c1f0af0c --- /dev/null +++ b/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 +} diff --git a/docs/dev/postgresql-schema-evolution.md b/docs/dev/postgresql-schema-evolution.md index 4b60280fc8..c26c2d00a8 100644 --- a/docs/dev/postgresql-schema-evolution.md +++ b/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: