feat(tenants): snapshot a migrated template schema per shard (#15278)

backend/alembic/check_tenant_snapshot.py

@@ -0,0 +1,47 @@
+#!/usr/bin/env python3
+"""Deploy gate for tenant schema snapshots, meant to run after the migration job:
+clone each shard's snapshot for the code's head, compare it with a freshly migrated
+schema, and fail on any difference so no snapshot ships that the chain would not build."""
+
+import sys
+
+from ee.onyx.db.tenant_snapshot import check_snapshot_parity
+from ee.onyx.db.tenant_snapshot import get_head_revision
+from ee.onyx.db.tenant_snapshot import get_snapshot
+from onyx.db.engine.shard_registry import get_shard_specs
+from onyx.db.engine.sql_engine import SqlEngine
+
+
+def main() -> int:
+    head_rev = get_head_revision()
+    if head_rev is None:
+        print("Could not determine head revision.", file=sys.stderr)
+        return 1
+
+    failed = False
+    with SqlEngine.scoped_engine(pool_size=5, max_overflow=2):
+        for shard_name in sorted(get_shard_specs()):
+            dump = get_snapshot(shard_name, head_rev)
+            if dump is None:
+                print(f"{shard_name}: no snapshot for head {head_rev}", file=sys.stderr)
+                failed = True
+                continue
+            differences = check_snapshot_parity(shard_name, dump)
+            if differences:
+                print(f"{shard_name}: snapshot differs from a migrated schema:")
+                print("\n".join(differences))
+                print(
+                    "Drop the template schema on that shard and rerun the migration "
+                    "job with --snapshot-template to rebuild it from the chain."
+                )
+                failed = True
+            else:
+                print(
+                    f"{shard_name}: snapshot for {head_rev} matches a migrated schema"
+                )
+
+    return 1 if failed else 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

backend/alembic/run_multitenant_migrations.py

@@ -24,11 +24,13 @@ import time
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from typing import NamedTuple
 
-from alembic.config import Config
-from alembic.script import ScriptDirectory
 
 from onyx.db.engine.sql_engine import SqlEngine
 from onyx.db.engine.tenant_utils import get_schemas_needing_migration
+from ee.onyx.db.tenant_snapshot import ensure_template_schema
+from ee.onyx.db.tenant_snapshot import get_head_revision
+from ee.onyx.db.tenant_snapshot import store_template_snapshots
+from onyx.db.engine.shard_registry import get_shard_specs
 from onyx.db.engine.tenant_utils import get_tenant_ids_by_shard
 from shared_configs.configs import TENANT_ID_PREFIX
 
@@ -40,6 +42,7 @@ from shared_configs.configs import TENANT_ID_PREFIX
 class Args(NamedTuple):
     jobs: int
     batch_size: int
+    snapshot_template: bool
 
 
 class Batch(NamedTuple):
@@ -104,13 +107,6 @@ def run_alembic_for_batch(batch: Batch) -> BatchResult:
     return BatchResult(schemas, False, retry.stdout or "", elapsed)
 
 
-def get_head_revision() -> str | None:
-    """Get the head revision from the alembic script directory."""
-    alembic_cfg = Config("alembic.ini")
-    script = ScriptDirectory.from_config(alembic_cfg)
-    return script.get_current_head()
-
-
 def run_migrations_parallel(
     schemas_by_shard: dict[str, list[str]],
     max_workers: int,
@@ -252,12 +248,24 @@ def parse_args() -> Args:
         metavar="N",
         help="Schemas per alembic process (default: 50)",
     )
+    parser.add_argument(
+        "--snapshot-template",
+        action="store_true",
+        help=(
+            "Migrate the template schema on every shard with the tenants and "
+            "store its dump afterwards (cloud provisioning clones it)"
+        ),
+    )
     args = parser.parse_args()
     if args.jobs < 1:
         parser.error("--jobs must be >= 1")
     if args.batch_size < 1:
         parser.error("--batch-size must be >= 1")
-    return Args(jobs=args.jobs, batch_size=args.batch_size)
+    return Args(
+        jobs=args.jobs,
+        batch_size=args.batch_size,
+        snapshot_template=args.snapshot_template,
+    )
 
 
 def main() -> int:
@@ -270,6 +278,10 @@ def main() -> int:
 
     schemas_by_shard: dict[str, list[str]] = {}
     with SqlEngine.scoped_engine(pool_size=5, max_overflow=2):
+        if args.snapshot_template:
+            for shard_name in sorted(get_shard_specs()):
+                ensure_template_schema(shard_name)
+
         # The prefix filter drops `public`, which enumeration reports as the sole
         # "tenant" outside multi-tenant mode. That is what makes the hint below fire.
         tenants_by_shard = {
@@ -295,22 +307,27 @@ def main() -> int:
                 schemas_by_shard[shard_name] = needing
 
     total_to_migrate = sum(len(s) for s in schemas_by_shard.values())
-    if not total_to_migrate:
+    if total_to_migrate:
+        print(
+            f"{total_to_migrate}/{total_tenants} tenants need migration (head: {head_rev})."
+        )
+        success = run_migrations_parallel(
+            schemas_by_shard,
+            max_workers=args.jobs,
+            batch_size=args.batch_size,
+        )
+        print(
+            f"\n{'All migrations successful' if success else 'Some migrations failed'}"
+        )
+        if not success:
+            return 1
+    else:
         print(f"All {total_tenants} tenants are already at head revision ({head_rev}).")
-        return 0
 
-    print(
-        f"{total_to_migrate}/{total_tenants} tenants need migration (head: {head_rev})."
-    )
-
-    success = run_migrations_parallel(
-        schemas_by_shard,
-        max_workers=args.jobs,
-        batch_size=args.batch_size,
-    )
-
-    print(f"\n{'All migrations successful' if success else 'Some migrations failed'}")
-    return 0 if success else 1
+    if args.snapshot_template:
+        with SqlEngine.scoped_engine(pool_size=5, max_overflow=2):
+            store_template_snapshots(head_rev)
+    return 0
 
 
 if __name__ == "__main__":

backend/alembic_tenants/versions/5e0d2a7c9f41_add_tenant_schema_snapshot.py

@@ -0,0 +1,42 @@
+"""add tenant schema snapshot
+
+One row per shard and head revision: the SQL dump of that shard's template schema,
+which apply_snapshot renders into a new tenant schema instead of replaying the chain.
+
+Revision ID: 5e0d2a7c9f41
+Revises: a754e4f72e60
+Create Date: 2026-09-29
+
+"""
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "5e0d2a7c9f41"
+down_revision = "a754e4f72e60"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    op.create_table(
+        "tenant_schema_snapshot",
+        sa.Column("id", sa.Integer(), nullable=False),
+        sa.Column("shard_name", sa.String(), nullable=False),
+        sa.Column("alembic_revision", sa.String(), nullable=False),
+        sa.Column("dump", sa.Text(), nullable=False),
+        sa.Column(
+            "created_at",
+            sa.DateTime(timezone=True),
+            server_default=sa.func.now(),
+            nullable=False,
+        ),
+        sa.PrimaryKeyConstraint("id"),
+        sa.UniqueConstraint("shard_name", "alembic_revision"),
+        schema="public",
+    )
+
+
+def downgrade() -> None:
+    op.drop_table("tenant_schema_snapshot", schema="public")

backend/ee/onyx/db/tenant_snapshot.py

@@ -0,0 +1,486 @@
+"""Snapshot of a migrated tenant schema, cloned into new tenants.
+
+The rollout job keeps one template schema per shard at head, dumps it after the
+migration run, and stores the dump keyed by revision. Rendering the dump with a
+new tenant's name and applying it yields what an empty schema becomes after the
+whole migration chain, baseline rows included.
+"""
+
+import os
+import re
+import shutil
+import subprocess
+import uuid
+from collections.abc import Iterator
+from contextlib import contextmanager
+from difflib import unified_diff
+from pathlib import Path
+from urllib.parse import unquote_plus
+
+from alembic.config import Config
+from alembic.script import ScriptDirectory
+from sqlalchemy import (
+    DateTime,
+    LargeBinary,
+    Text,
+    Uuid,
+    case,
+    cast,
+    column,
+    delete,
+    func,
+    inspect,
+    literal,
+    select,
+    table,
+    tuple_,
+)
+from sqlalchemy.dialects.postgresql import aggregate_order_by
+from sqlalchemy.engine import Connection, Engine
+from sqlalchemy.orm import Session
+from sqlalchemy.schema import CreateSchema, DropSchema
+from sqlalchemy.sql.elements import ColumnElement
+from sqlalchemy.sql.selectable import TableClause
+
+from onyx.configs.app_configs import AWS_REGION_NAME, DB_READONLY_USER, USE_IAM_AUTH
+from onyx.configs.constants import SSL_CERT_FILE
+from onyx.db.engine.iam_auth import get_iam_auth_token
+from onyx.db.engine.pg_ssl import pg_ssl_psycopg2_connect_args
+from onyx.db.engine.shard_registry import (
+    get_engine_for_shard,
+    get_shard_spec,
+    get_shard_specs,
+)
+from onyx.db.engine.sql_engine import get_catalog_session
+from onyx.db.engine.tenant_utils import validate_tenant_id
+from onyx.db.models import TenantSchemaSnapshot
+from onyx.utils.logger import setup_logger
+from shared_configs.configs import TENANT_TEMPLATE_SCHEMA
+
+logger = setup_logger()
+
+
+def _schema_name_pattern(schema: str) -> re.Pattern[str]:
+    """The name as an identifier, quoted or bare, never as part of a longer word."""
+    return re.compile(rf'(?<!\w)"?{re.escape(schema)}"?(?!\w)')
+
+
+# The image being replaced keeps its own snapshot until the rollout completes.
+_SNAPSHOTS_KEPT_PER_SHARD = 2
+# The dump names the template everywhere, including inside trigger function bodies,
+# so rendering is one identifier swap. pg_dump leaves the plain template name bare
+# but a real tenant name needs quoting, so every occurrence is rendered quoted.
+_TEMPLATE_NAME = _schema_name_pattern(TENANT_TEMPLATE_SCHEMA)
+# Dropped from the dump: the schema line (the caller owns the target schema),
+# psql meta-commands, and session settings, which vary by pg_dump version and
+# have no place inside a tenant build.
+_DRIVER_UNSAFE_LINES = re.compile(
+    r"^(CREATE SCHEMA .*;|\\(un)?restrict .*|SET \w+ = .*;|SELECT pg_catalog\.set_config\(.*\);)$",
+    re.MULTILINE,
+)
+# libpq reads these instead of a URL, so the password never appears in argv.
+_LIBPQ_SSL_ENV = {
+    "sslmode": "PGSSLMODE",
+    "sslrootcert": "PGSSLROOTCERT",
+    "sslcert": "PGSSLCERT",
+    "sslkey": "PGSSLKEY",
+}
+# Values that legitimately differ between a clone and a fresh migration: row
+# timestamps, encrypted blobs (a random salt per write) and generated ids.
+_UNCOMPARED_COLUMN_TYPES = (DateTime, LargeBinary, Uuid)
+_DIFF_LINES_REPORTED = 60
+# Seeded rows with one JSON field set to the migration's run date, so clone and
+# fresh schema differ by day. Only that field is dropped before comparing.
+# Table to (key column, key value, JSON column, dated field).
+_RUN_DATE_SEEDED_FIELDS = {
+    "key_value_store": ("key", "kg_config", "value", "KG_COVERAGE_START")
+}
+# Postgres deparses a varchar list in a CHECK or partial index either as an array
+# of casts or as a cast of an array, flipping form on every re-parse. Same
+# constraint, so both spellings compare as one.
+_ARRAY_OF_CASTS = re.compile(
+    r"ARRAY\[((?:\('[^']*'::character varying\)::text(?:, )?)+)\]"
+)
+_CAST_OF_ARRAY = re.compile(
+    r"\(ARRAY\[((?:'[^']*'::character varying(?:, )?)+)\]\)::text\[\]"
+)
+_QUOTED_LITERAL = re.compile(r"'[^']*'")
+# Where alembic.ini lives: the alembic subprocess and the head lookup run from here.
+_BACKEND_DIR = Path(__file__).resolve().parents[3]
+
+
+def get_head_revision() -> str | None:
+    config = Config(str(_BACKEND_DIR / "alembic.ini"))
+    # The ini names the scripts folder relative to the working directory.
+    config.set_main_option("script_location", str(_BACKEND_DIR / "alembic"))
+    return ScriptDirectory.from_config(config).get_current_head()
+
+
+def scratch_schema_name() -> str:
+    """A short-lived schema the parity check builds and drops. The name passes the
+    tenant validator but no scheduler treats it as a workspace."""
+    return f"{TENANT_TEMPLATE_SCHEMA}_{uuid.uuid4().hex}"
+
+
+@contextmanager
+def template_session(shard_name: str) -> Iterator[Session]:
+    """Session on one shard's template. Tenant routing cannot reach it, since
+    every shard holds a template under the same name."""
+    with (
+        get_engine_for_shard(shard_name)
+        .connect()
+        .execution_options(
+            schema_translate_map={None: TENANT_TEMPLATE_SCHEMA}
+        ) as connection
+    ):
+        session = Session(bind=connection, expire_on_commit=False)
+        try:
+            yield session
+        finally:
+            session.close()
+
+
+def ensure_template_schema(shard_name: str) -> None:
+    """Created before enumeration so the template is migrated in this run."""
+    with get_engine_for_shard(shard_name).begin() as connection:
+        connection.execute(CreateSchema(TENANT_TEMPLATE_SCHEMA, if_not_exists=True))
+
+
+def dump_schema(shard_name: str, schema: str, schema_only: bool = False) -> str:
+    """pg_dump of one schema as plain SQL a driver can execute in one go.
+
+    Rows come out as INSERTs rather than COPY blocks for that reason."""
+    if not validate_tenant_id(schema):
+        raise ValueError(f"Refusing to dump schema {schema!r}")
+    if shutil.which("pg_dump") is None:
+        raise RuntimeError("pg_dump is not installed, run from an image that has it")
+    command = [
+        "pg_dump",
+        "--schema",
+        schema,
+        "--no-owner",
+        "--no-privileges",
+        "--no-comments",
+        "--no-tablespaces",
+        "--no-security-labels",
+        "--inserts",
+    ]
+    if schema_only:
+        command.append("--schema-only")
+    result = subprocess.run(
+        command,
+        env={**os.environ, **_libpq_env(shard_name)},
+        capture_output=True,
+        text=True,
+    )
+    if result.returncode != 0:
+        raise RuntimeError(f"pg_dump of {schema} failed: {result.stderr.strip()}")
+    return _DRIVER_UNSAFE_LINES.sub("", result.stdout)
+
+
+def _libpq_env(shard_name: str) -> dict[str, str]:
+    """The shard's connection as libpq variables, with the same auth and TLS
+    settings the engine uses for that shard."""
+    spec = get_shard_spec(shard_name)
+    # The spec carries the URL-encoded form the engine embeds in its URL.
+    env = {
+        "PGHOST": spec.host,
+        "PGPORT": spec.port,
+        "PGUSER": spec.user,
+        "PGDATABASE": spec.db,
+        "PGPASSWORD": unquote_plus(spec.password),
+    }
+    if USE_IAM_AUTH:
+        # Same token and TLS the engine's IAM connect handler applies.
+        env["PGPASSWORD"] = get_iam_auth_token(
+            spec.host, spec.port, spec.user, AWS_REGION_NAME
+        )
+        env["PGSSLMODE"] = "require"
+        env["PGSSLROOTCERT"] = SSL_CERT_FILE
+        return env
+    ssl_args = pg_ssl_psycopg2_connect_args()
+    if "sslpassword" in ssl_args:
+        # libpq has no variable for it and argv would expose it.
+        raise RuntimeError("pg_dump cannot use a passphrase-protected client key")
+    for arg, variable in _LIBPQ_SSL_ENV.items():
+        value = ssl_args.get(arg)
+        if value:
+            env[variable] = value
+    return env
+
+
+def store_template_snapshots(alembic_revision: str) -> None:
+    """Dump every shard's template after a migration run. A fresh dump each run
+    keeps re-encrypted rows current and marks the running head as the newest."""
+    for shard_name in sorted(get_shard_specs()):
+        _require_template_at(shard_name, alembic_revision)
+        dump = dump_schema(shard_name, TENANT_TEMPLATE_SCHEMA)
+        store_snapshot(shard_name, alembic_revision, dump)
+        logger.info(
+            "Stored template snapshot for shard %s at %s (%d bytes)",
+            shard_name,
+            alembic_revision,
+            len(dump),
+        )
+
+
+def _require_template_at(shard_name: str, alembic_revision: str) -> None:
+    version_table = table(
+        "alembic_version", column("version_num"), schema=TENANT_TEMPLATE_SCHEMA
+    )
+    with get_engine_for_shard(shard_name).connect() as connection:
+        stamped = connection.scalar(select(version_table.c.version_num))
+    if stamped != alembic_revision:
+        raise RuntimeError(
+            f"Template on shard {shard_name} is at {stamped}, not {alembic_revision}"
+        )
+
+
+def store_snapshot(shard_name: str, alembic_revision: str, dump: str) -> None:
+    """Upsert the shard's snapshot for this revision as the newest, and drop all
+    but the newest two."""
+    with get_catalog_session() as db_session:
+        existing = db_session.scalar(
+            select(TenantSchemaSnapshot).where(
+                TenantSchemaSnapshot.shard_name == shard_name,
+                TenantSchemaSnapshot.alembic_revision == alembic_revision,
+            )
+        )
+        if existing is None:
+            db_session.add(
+                TenantSchemaSnapshot(
+                    shard_name=shard_name,
+                    alembic_revision=alembic_revision,
+                    dump=dump,
+                )
+            )
+        else:
+            existing.dump = dump
+            # A rollback re-stores an older revision, which must then outlive
+            # the one it replaced.
+            existing.created_at = func.now()
+        db_session.flush()
+
+        keep = db_session.scalars(
+            select(TenantSchemaSnapshot.id)
+            .where(TenantSchemaSnapshot.shard_name == shard_name)
+            .order_by(TenantSchemaSnapshot.created_at.desc())
+            .limit(_SNAPSHOTS_KEPT_PER_SHARD)
+        ).all()
+        db_session.execute(
+            delete(TenantSchemaSnapshot).where(
+                TenantSchemaSnapshot.shard_name == shard_name,
+                TenantSchemaSnapshot.id.not_in(keep),
+            )
+        )
+        db_session.commit()
+
+
+def get_snapshot(shard_name: str, alembic_revision: str) -> str | None:
+    with get_catalog_session() as db_session:
+        return db_session.scalar(
+            select(TenantSchemaSnapshot.dump).where(
+                TenantSchemaSnapshot.shard_name == shard_name,
+                TenantSchemaSnapshot.alembic_revision == alembic_revision,
+            )
+        )
+
+
+def render_snapshot(dump: str, tenant_id: str) -> str:
+    if not validate_tenant_id(tenant_id) or tenant_id == TENANT_TEMPLATE_SCHEMA:
+        raise ValueError(f"Refusing to render a snapshot for {tenant_id!r}")
+    return _TEMPLATE_NAME.sub(f'"{tenant_id}"', dump)
+
+
+def apply_snapshot(engine: Engine, dump: str, tenant_id: str) -> None:
+    """Build the tenant schema from the dump in one transaction, stamped at the
+    snapshot's revision by the version row the dump carries. Migration grants
+    schema usage to the read-only role, which a dump lacks, so it is granted here."""
+    rendered = render_snapshot(dump, tenant_id)
+    with engine.connect() as connection:
+        with connection.begin():
+            connection.execute(CreateSchema(tenant_id, if_not_exists=True))
+            _execute_verbatim(connection, rendered)
+            _grant_readonly_usage(connection, tenant_id)
+        # Defensive: a session setting the strip regex let through would
+        # otherwise stay on this pooled connection.
+        connection.invalidate()
+
+
+def _execute_verbatim(connection: Connection, statements: str) -> None:
+    """Straight to the driver cursor with no parameters: the dump and the grant
+    contain literal % signs that parameter interpolation would try to expand."""
+    cursor = connection.connection.cursor()
+    try:
+        cursor.execute(statements)
+    finally:
+        cursor.close()
+
+
+def _grant_readonly_usage(connection: Connection, tenant_id: str) -> None:
+    # tenant_id was validated in render_snapshot. The role name is trusted
+    # config, interpolated the way the migration that grants this does.
+    _execute_verbatim(
+        connection,
+        f"""
+        DO $$
+        BEGIN
+            IF EXISTS (SELECT FROM pg_catalog.pg_roles WHERE rolname = '{DB_READONLY_USER}') THEN
+                EXECUTE format('GRANT USAGE ON SCHEMA %I TO %I', '{tenant_id}', '{DB_READONLY_USER}');
+            END IF;
+        END
+        $$;
+        """,
+    )
+
+
+def check_snapshot_parity(shard_name: str, dump: str) -> list[str]:
+    """Differences between a clone of the dump and a freshly migrated schema.
+
+    Empty means the snapshot is safe to clone. Both scratch schemas are dropped."""
+    engine = get_engine_for_shard(shard_name)
+    cloned = scratch_schema_name()
+    migrated = scratch_schema_name()
+    with _dropped_afterwards(engine, cloned, migrated):
+        apply_snapshot(engine, dump, cloned)
+        _migrate_empty_schema(shard_name, migrated)
+        return compare_schemas(shard_name, cloned, migrated)
+
+
+def compare_schemas(shard_name: str, left: str, right: str) -> list[str]:
+    """Structure must match line for line after name, comment and array-cast
+    normalisation. Rows must match in count and, outside the uncompared column
+    types, in content."""
+    differences = _structure_differences(shard_name, left, right)
+    with get_engine_for_shard(shard_name).connect() as connection:
+        differences.extend(_row_differences(connection, left, right))
+    return differences
+
+
+def _structure_differences(shard_name: str, left: str, right: str) -> list[str]:
+    left_lines = _normalised_structure(dump_schema(shard_name, left, True), left)
+    right_lines = _normalised_structure(dump_schema(shard_name, right, True), right)
+    diff = list(
+        unified_diff(left_lines, right_lines, fromfile=left, tofile=right, lineterm="")
+    )
+    if not diff:
+        return []
+    return ["structure differs:"] + diff[:_DIFF_LINES_REPORTED]
+
+
+def _normalised_structure(dump: str, schema: str) -> list[str]:
+    body = _schema_name_pattern(schema).sub("SCHEMA", dump)
+    return [
+        _canonical_text_arrays(line)
+        for line in body.splitlines()
+        if line.strip() and not line.startswith("--")
+    ]
+
+
+def _canonical_text_arrays(line: str) -> str:
+    def literals_only(match: re.Match[str]) -> str:
+        return "ARRAY[" + ", ".join(_QUOTED_LITERAL.findall(match.group(1))) + "]"
+
+    return _CAST_OF_ARRAY.sub(literals_only, _ARRAY_OF_CASTS.sub(literals_only, line))
+
+
+def _row_differences(connection: Connection, left: str, right: str) -> list[str]:
+    differences: list[str] = []
+    for table_name in _tables(connection, left):
+        columns = _compared_columns(connection, left, table_name)
+        left_count, left_digest = _row_digest(connection, left, table_name, columns)
+        right_count, right_digest = _row_digest(connection, right, table_name, columns)
+        if left_count != right_count:
+            differences.append(
+                f"{table_name}: {left_count} rows in {left}, {right_count} in {right}"
+            )
+        elif left_digest != right_digest:
+            differences.append(f"{table_name}: row contents differ")
+    return differences
+
+
+def _tables(connection: Connection, schema: str) -> list[str]:
+    return sorted(inspect(connection).get_table_names(schema=schema))
+
+
+def _compared_columns(
+    connection: Connection, schema: str, table_name: str
+) -> list[str]:
+    """Columns whose values a clone and a fresh migration must agree on."""
+    return [
+        reflected["name"]
+        for reflected in inspect(connection).get_columns(table_name, schema=schema)
+        if "computed" not in reflected
+        and not isinstance(reflected["type"], _UNCOMPARED_COLUMN_TYPES)
+    ]
+
+
+def _row_digest(
+    connection: Connection, schema: str, table_name: str, columns: list[str]
+) -> tuple[int, str]:
+    source = table(table_name, *(column(name) for name in columns), schema=schema)
+    if not columns:
+        count = connection.scalar(select(func.count()).select_from(source))
+        return int(count or 0), ""
+    # A row constructor keeps NULL positions, so a NULL moving between columns
+    # still differs.
+    row_text = cast(
+        tuple_(*(_compared_value(source, name) for name in columns)),
+        Text,
+    ).label("row_text")
+    rows = select(row_text).select_from(source).subquery()
+    digest = func.md5(
+        func.coalesce(
+            func.string_agg(
+                rows.c.row_text, aggregate_order_by(literal("|"), rows.c.row_text)
+            ),
+            "",
+        )
+    )
+    row = connection.execute(select(func.count(), digest).select_from(rows)).one()
+    return int(row[0]), str(row[1])
+
+
+def _compared_value(source: TableClause, name: str) -> ColumnElement:
+    seeded = _RUN_DATE_SEEDED_FIELDS.get(source.name)
+    if seeded is None or name != seeded[2]:
+        return source.c[name]
+    key_column, key_value, _, dated_field = seeded
+    return case(
+        (source.c[key_column] == key_value, source.c[name].op("-")(dated_field)),
+        else_=source.c[name],
+    )
+
+
+def _migrate_empty_schema(shard_name: str, schema: str) -> None:
+    """Run the migration chain against a new schema, the way the rollout job does."""
+    result = subprocess.run(
+        [
+            "alembic",
+            "-x",
+            f"schemas={schema}",
+            "-x",
+            f"shard={shard_name}",
+            "upgrade",
+            "head",
+        ],
+        cwd=_BACKEND_DIR,
+        env=os.environ,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT,
+        text=True,
+    )
+    if result.returncode != 0:
+        raise RuntimeError(
+            f"Migrating {schema} for the parity check failed:\n{result.stdout}"
+        )
+
+
+@contextmanager
+def _dropped_afterwards(engine: Engine, *schemas: str) -> Iterator[None]:
+    try:
+        yield
+    finally:
+        with engine.begin() as connection:
+            for schema in schemas:
+                connection.execute(DropSchema(schema, cascade=True, if_exists=True))

backend/onyx/db/engine/tenant_utils.py

@@ -13,6 +13,7 @@ from shared_configs.configs import (
     MULTI_TENANT,
     POSTGRES_DEFAULT_SCHEMA,
     TENANT_ID_PREFIX,
+    TENANT_TEMPLATE_SCHEMA,
 )
 
 # Regex pattern for valid tenant IDs:
@@ -25,6 +26,7 @@ TENANT_ID_PATTERN = re.compile(
     r"[a-f0-9]{8}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{12}"  # UUID
     r"|i-[a-f0-9]+"  # AWS instance ID
     r"|dev"  # staff dev tenant
+    r"|template(_[a-f0-9]{32})?"  # rollout template, and its parity scratch copies
     r")$"
 )
 
@@ -179,4 +181,22 @@ def get_all_tenant_ids() -> list[str]:
         return [POSTGRES_DEFAULT_SCHEMA]
 
     # Deduped: a tenant mid-copy exists on two shards but is still one tenant.
-    return sorted(set(chain.from_iterable(get_tenant_ids_by_shard().values())))
+    # The template and the parity scratch schemas are nobody's workspace, so
+    # no scheduled work runs against them.
+    return sorted(
+        {
+            tenant_id
+            for tenant_id in chain.from_iterable(get_tenant_ids_by_shard().values())
+            if not tenant_id.startswith(TENANT_TEMPLATE_SCHEMA)
+        }
+    )
+
+
+def get_template_shards() -> list[str]:
+    """Shards holding a rollout template. Every shard has its own copy under one
+    name, so maintenance that must reach them, such as key rotation, goes by shard."""
+    return sorted(
+        shard_name
+        for shard_name, tenant_ids in get_tenant_ids_by_shard().items()
+        if TENANT_TEMPLATE_SCHEMA in tenant_ids
+    )

backend/onyx/db/models.py

@@ -5669,6 +5669,27 @@ class AvailableTenant(PublicBase):
     shard_name: Mapped[str | None] = mapped_column(String, nullable=True)
 
 
+class TenantSchemaSnapshot(PublicBase):
+    """SQL dump of a shard's template schema at one head revision, written by the
+    rollout job. apply_snapshot builds a tenant schema from it instead of replaying
+    the migration chain. The newest two per shard are kept so the image being
+    replaced keeps its own snapshot until the rollout completes."""
+
+    __tablename__ = "tenant_schema_snapshot"
+    __table_args__ = (
+        UniqueConstraint("shard_name", "alembic_revision"),
+        {"schema": "public"},
+    )
+
+    id: Mapped[int] = mapped_column(Integer, primary_key=True)
+    shard_name: Mapped[str] = mapped_column(String, nullable=False)
+    alembic_revision: Mapped[str] = mapped_column(String, nullable=False)
+    dump: Mapped[str] = mapped_column(Text, nullable=False)
+    created_at: Mapped[datetime.datetime] = mapped_column(
+        DateTime(timezone=True), server_default=func.now(), nullable=False
+    )
+
+
 # This is a mapping from tenant IDs to anonymous user paths
 class TenantAnonymousUserPath(PublicBase):
     __tablename__ = "tenant_anonymous_user_path"

backend/scripts/reencrypt_secrets.py

@@ -24,11 +24,15 @@ import sys
 parent_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
 sys.path.append(parent_dir)
 
+from ee.onyx.db.tenant_snapshot import template_session  # noqa: E402
 from onyx.db.engine.sql_engine import (  # noqa: E402
     SqlEngine,
     get_session_with_tenant,
 )
-from onyx.db.engine.tenant_utils import get_all_tenant_ids  # noqa: E402
+from onyx.db.engine.tenant_utils import (  # noqa: E402
+    get_all_tenant_ids,
+    get_template_shards,
+)
 from onyx.db.rotate_encryption_key import rotate_encryption_key  # noqa: E402
 from onyx.utils.variable_functionality import global_version  # noqa: E402
 from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA  # noqa: E402
@@ -38,7 +42,18 @@ def _run_for_tenant(tenant_id: str, old_key: str | None, dry_run: bool = False)
     print(f"Re-encrypting secrets for tenant: {tenant_id}")
     with get_session_with_tenant(tenant_id=tenant_id) as db_session:
         results = rotate_encryption_key(db_session, old_key=old_key, dry_run=dry_run)
+    _report(results, dry_run)
 
+
+def _run_for_template(shard_name: str, old_key: str | None, dry_run: bool) -> None:
+    # The template is nobody's workspace but its rows are cloned into new ones.
+    print(f"Re-encrypting secrets for the tenant template on shard: {shard_name}")
+    with template_session(shard_name) as db_session:
+        results = rotate_encryption_key(db_session, old_key=old_key, dry_run=dry_run)
+    _report(results, dry_run)
+
+
+def _report(results: dict[str, int], dry_run: bool) -> None:
     if results:
         for col, count in results.items():
             print(
@@ -95,6 +110,12 @@ def main() -> None:
             except Exception as e:
                 print(f"  ERROR for tenant {tid}: {e}")
                 failed_tenants.append(tid)
+        for shard_name in get_template_shards():
+            try:
+                _run_for_template(shard_name, old_key, dry_run=args.dry_run)
+            except Exception as e:
+                print(f"  ERROR for the template on shard {shard_name}: {e}")
+                failed_tenants.append(f"template on {shard_name}")
         if failed_tenants:
             print(f"FAILED tenants ({len(failed_tenants)}): {failed_tenants}")
             sys.exit(1)

backend/shared_configs/configs.py

@@ -213,6 +213,10 @@ async def async_return_default_schema(
 
 # Prefix used for all tenant ids
 TENANT_ID_PREFIX = "tenant_"
+# Schema the rollout job keeps at head and dumps into the snapshot that new tenants
+# are cloned from. Migrated and re-encrypted like a tenant, never scheduled work
+# or assigned a user.
+TENANT_TEMPLATE_SCHEMA = f"{TENANT_ID_PREFIX}template"
 
 DISALLOWED_SLACK_BOT_TENANT_IDS = os.environ.get("DISALLOWED_SLACK_BOT_TENANT_IDS")
 DISALLOWED_SLACK_BOT_TENANT_LIST = (

backend/tests/external_dependency_unit/db/test_tenant_snapshot_clone.py

@@ -0,0 +1,182 @@
+"""A tenant cloned from the template snapshot must equal a migrated schema:
+same structure, same baseline rows, stamped at head. The parity comparison
+must catch a structural or row difference, since it gates the deploy."""
+
+import os
+import uuid
+from collections.abc import Generator
+from typing import cast
+from unittest.mock import patch
+
+import pytest
+from sqlalchemy import Table, delete, func, select, text
+
+from ee.onyx.db import tenant_snapshot
+from onyx.db.engine import tenant_utils
+from onyx.db.engine.shard_registry import get_default_shard_name, get_engine_for_shard
+from onyx.db.engine.sql_engine import SqlEngine
+from onyx.db.engine.tenant_utils import get_template_shards
+from onyx.db.models import PublicBase, TenantSchemaSnapshot, Tool
+from shared_configs.configs import TENANT_TEMPLATE_SCHEMA
+
+
+@pytest.fixture(scope="module")
+def shard() -> Generator[str, None, None]:
+    SqlEngine.init_engine(pool_size=5, max_overflow=2)
+    shard_name = get_default_shard_name()
+    # The snapshot table comes from the catalog chain, which the single-tenant
+    # test database never runs.
+    PublicBase.metadata.create_all(
+        SqlEngine.get_engine(), tables=[cast(Table, TenantSchemaSnapshot.__table__)]
+    )
+    # The chain seeds cloud-only rows under this flag, which is what the
+    # rollout job's template gets.
+    os.environ["MULTI_TENANT"] = "true"
+    tenant_snapshot.ensure_template_schema(shard_name)
+    try:
+        tenant_snapshot._migrate_empty_schema(shard_name, TENANT_TEMPLATE_SCHEMA)
+        yield shard_name
+    finally:
+        os.environ.pop("MULTI_TENANT", None)
+        with tenant_snapshot._dropped_afterwards(
+            get_engine_for_shard(shard_name), TENANT_TEMPLATE_SCHEMA
+        ):
+            pass
+        SqlEngine.reset_engine()
+
+
+@pytest.fixture(scope="module")
+def dump(shard: str) -> str:
+    return tenant_snapshot.dump_schema(shard, TENANT_TEMPLATE_SCHEMA)
+
+
+@pytest.fixture
+def clone(shard: str, dump: str) -> Generator[str, None, None]:
+    tenant_id = tenant_snapshot.scratch_schema_name()
+    tenant_snapshot.apply_snapshot(get_engine_for_shard(shard), dump, tenant_id)
+    yield tenant_id
+    with tenant_snapshot._dropped_afterwards(get_engine_for_shard(shard), tenant_id):
+        pass
+
+
+def test_clone_matches_the_template(shard: str, clone: str) -> None:
+    assert tenant_snapshot.compare_schemas(shard, clone, TENANT_TEMPLATE_SCHEMA) == []
+
+    with get_engine_for_shard(shard).connect() as connection:
+        stamped = connection.scalar(
+            text(f'SELECT version_num FROM "{clone}".alembic_version')
+        )
+        tool_count = connection.scalar(text(f'SELECT count(*) FROM "{clone}".tool'))
+    assert stamped == tenant_snapshot.get_head_revision()
+    assert tool_count and tool_count > 0
+
+
+def test_parity_catches_a_missing_column(shard: str, clone: str) -> None:
+    with get_engine_for_shard(shard).begin() as connection:
+        connection.execute(
+            text(f'ALTER TABLE "{clone}".persona DROP COLUMN description')
+        )
+
+    differences = tenant_snapshot.compare_schemas(shard, clone, TENANT_TEMPLATE_SCHEMA)
+    assert differences and differences[0] == "structure differs:"
+
+
+def test_parity_ignores_the_migration_date_seed(shard: str, clone: str) -> None:
+    # The knowledge graph config is seeded with the migration's run date.
+    with get_engine_for_shard(shard).begin() as connection:
+        connection.execute(
+            text(
+                f'UPDATE "{clone}".key_value_store '
+                "SET value = jsonb_set(value, '{KG_COVERAGE_START}', '\"2000-01-01\"') "
+                "WHERE key = 'kg_config'"
+            )
+        )
+
+    assert tenant_snapshot.compare_schemas(shard, clone, TENANT_TEMPLATE_SCHEMA) == []
+
+    # Only the dated field is dropped: any other setting still has to match.
+    with get_engine_for_shard(shard).begin() as connection:
+        connection.execute(
+            text(
+                f'UPDATE "{clone}".key_value_store '
+                "SET value = jsonb_set(value, '{KG_ENABLED}', 'true') "
+                "WHERE key = 'kg_config'"
+            )
+        )
+    differences = tenant_snapshot.compare_schemas(shard, clone, TENANT_TEMPLATE_SCHEMA)
+    assert any(difference.startswith("key_value_store: ") for difference in differences)
+
+
+def test_parity_catches_a_missing_row(shard: str, clone: str) -> None:
+    with get_engine_for_shard(shard).begin() as connection:
+        connection.execute(
+            text(
+                f'DELETE FROM "{clone}".tool WHERE id = '
+                f'(SELECT min(id) FROM "{clone}".tool)'
+            )
+        )
+
+    differences = tenant_snapshot.compare_schemas(shard, clone, TENANT_TEMPLATE_SCHEMA)
+    assert any(difference.startswith("tool: ") for difference in differences)
+
+
+def test_deploy_gate_passes_for_the_template_snapshot(shard: str, dump: str) -> None:
+    # Migrates a scratch schema through the full chain, so this is the slow one.
+    assert tenant_snapshot.check_snapshot_parity(shard, dump) == []
+
+
+def test_rollout_stores_the_template_only_at_head(shard: str) -> None:
+    head = tenant_snapshot.get_head_revision()
+    assert head is not None
+    try:
+        tenant_snapshot.store_template_snapshots(head)
+        assert tenant_snapshot.get_snapshot(shard, head)
+        with pytest.raises(RuntimeError):
+            tenant_snapshot.store_template_snapshots("not-the-head")
+    finally:
+        _delete_snapshot_rows(shard, head)
+
+
+def test_template_is_rotated_by_shard(shard: str) -> None:
+    # Enumeration only looks past the default schema in multi-tenant mode.
+    with patch.object(tenant_utils, "MULTI_TENANT", True):
+        assert get_template_shards() == [shard]
+    # The session binds model tables to the template, which is what rotation reads.
+    with tenant_snapshot.template_session(shard) as db_session:
+        seeded_tools = db_session.scalar(select(func.count()).select_from(Tool))
+    assert seeded_tools and seeded_tools > 0
+
+
+def test_render_refuses_names_that_are_not_tenants(dump: str) -> None:
+    with pytest.raises(ValueError):
+        tenant_snapshot.render_snapshot(dump, "public")
+    with pytest.raises(ValueError):
+        tenant_snapshot.render_snapshot(dump, TENANT_TEMPLATE_SCHEMA)
+
+
+@pytest.mark.usefixtures("shard")
+def test_store_keeps_the_newest_two() -> None:
+    # Retention is per shard, so a made-up shard cannot evict real rows.
+    fake_shard = f"test-shard-{uuid.uuid4().hex[:8]}"
+    revisions = [f"test-{uuid.uuid4().hex[:8]}" for _ in range(3)]
+    try:
+        for revision in revisions:
+            tenant_snapshot.store_snapshot(fake_shard, revision, f"-- {revision}")
+        assert tenant_snapshot.get_snapshot(fake_shard, revisions[0]) is None
+        assert (
+            tenant_snapshot.get_snapshot(fake_shard, revisions[2])
+            == f"-- {revisions[2]}"
+        )
+    finally:
+        _delete_snapshot_rows(fake_shard)
+
+
+def _delete_snapshot_rows(shard: str, head: str | None = None) -> None:
+    statement = delete(TenantSchemaSnapshot).where(
+        TenantSchemaSnapshot.shard_name == shard
+    )
+    if head is not None:
+        statement = statement.where(TenantSchemaSnapshot.alembic_revision == head)
+    with tenant_snapshot.get_catalog_session() as db_session:
+        db_session.execute(statement)
+        db_session.commit()