diff --git a/tests/scripts/test_bootstrap_shared_schema.py b/tests/scripts/test_bootstrap_shared_schema.py new file mode 100644 index 00000000..41b7bc29 --- /dev/null +++ b/tests/scripts/test_bootstrap_shared_schema.py @@ -0,0 +1,129 @@ +"""Integration test for the shared-tenant schema bootstrap. + +Runs ``scripts/bootstrap_shared_schema`` against a throwaway database and asserts +the partitioned schema it produces: every tenant-scoped table is +HASH(tenant_id)-partitioned with its full set of partitions, the primary key +leads with ``tenant_id``, and the tenant-scoped composite keys enforce (two +tenants can share a workspace name, a duplicate name within a tenant is rejected, +and an unknown ``tenant_id`` is rejected by the FK to ``tenants``). + +Heavier than a metadata-only shape check on purpose: it exercises the exact DDL +the migration track will run, against real Postgres + pgvector, and it does not +go through the app's ``create_all`` fixture (which does not create partitions). +""" + +import pytest +from sqlalchemy import Engine, create_engine, make_url, text +from sqlalchemy.exc import IntegrityError +from sqlalchemy.orm import Session + +from scripts.bootstrap_shared_schema import ( + PARTITION_COUNT, + bootstrap_shared_schema, + partitioned_tables, +) +from src.config import settings +from src.models import Tenant, Workspace + + +@pytest.fixture +def bootstrapped_engine(worker_id: str): + """A throwaway database with the shared-tenant schema bootstrapped into it.""" + db_name = f"test_bootstrap_shared_{worker_id}" + admin_url = make_url(settings.DB.CONNECTION_URI).set(database="postgres") + + def _drop_create(create: bool) -> None: + # Pass the URL object (not str) so the password is not masked to '***'. + admin = create_engine(admin_url, isolation_level="AUTOCOMMIT") + with admin.connect() as conn: + conn.execute(text(f'DROP DATABASE IF EXISTS "{db_name}" WITH (FORCE)')) + if create: + conn.execute(text(f'CREATE DATABASE "{db_name}"')) + admin.dispose() + + _drop_create(create=True) + engine = create_engine(make_url(settings.DB.CONNECTION_URI).set(database=db_name)) + # AUTOCOMMIT: the partition DDL accumulates too many locks for one transaction. + with engine.connect().execution_options(isolation_level="AUTOCOMMIT") as conn: + conn.execute(text("CREATE EXTENSION IF NOT EXISTS vector")) + bootstrap_shared_schema(conn) + try: + yield engine + finally: + engine.dispose() + _drop_create(create=False) + + +def test_every_tenant_table_is_partitioned(bootstrapped_engine: Engine) -> None: + """Each tenant-scoped table is a HASH partitioned table carrying all its + partitions — guards the drift trap where a partitioned parent with zero + partitions rejects every insert.""" + expected = partitioned_tables() + assert len(expected) == 9 + with bootstrapped_engine.connect() as conn: + for name in expected: + relkind = conn.execute( + text( + "SELECT relkind FROM pg_class" + + " WHERE relname = :n AND relnamespace = 'public'::regnamespace" + ), + {"n": name}, + ).scalar() + assert relkind == "p", f"{name} is not partitioned (relkind={relkind!r})" + n_partitions = conn.execute( + text( + "SELECT count(*) FROM pg_inherits i" + + " JOIN pg_class p ON p.oid = i.inhparent WHERE p.relname = :n" + ), + {"n": name}, + ).scalar() + assert n_partitions == PARTITION_COUNT, ( + f"{name} has {n_partitions} partitions, expected {PARTITION_COUNT}" + ) + + +def test_primary_key_leads_with_tenant_id(bootstrapped_engine: Engine) -> None: + with bootstrapped_engine.connect() as conn: + pk_cols = ( + conn.execute( + text( + "SELECT a.attname FROM pg_index i" + + " JOIN pg_attribute a" + + " ON a.attrelid = i.indrelid AND a.attnum = ANY(i.indkey)" + + " WHERE i.indrelid = 'public.workspaces'::regclass AND i.indisprimary" + + " ORDER BY array_position(i.indkey, a.attnum)" + ) + ) + .scalars() + .all() + ) + assert pk_cols == ["tenant_id", "id"] + + +def test_tenant_scoped_uniqueness_and_fk(bootstrapped_engine: Engine) -> None: + """The tenant-scoped composite keys enforce as designed.""" + with Session(bootstrapped_engine) as session: + session.add_all([Tenant(tenant_id="tenant_a"), Tenant(tenant_id="tenant_b")]) + session.commit() + + # Two different tenants may each own a workspace named "default": this is + # the whole point of UNIQUE(tenant_id, name) replacing UNIQUE(name). + session.add_all( + [ + Workspace(tenant_id="tenant_a", name="default"), + Workspace(tenant_id="tenant_b", name="default"), + ] + ) + session.commit() + + # Same (tenant, name) twice is rejected. + with Session(bootstrapped_engine) as session: + session.add(Workspace(tenant_id="tenant_a", name="default")) + with pytest.raises(IntegrityError): + session.commit() + + # An unknown tenant_id is rejected by the FK to tenants. + with Session(bootstrapped_engine) as session: + session.add(Workspace(tenant_id="ghost_tenant", name="scratch")) + with pytest.raises(IntegrityError): + session.commit()