From d1b79ded63a8ebea51ab83cc303bcc448f00233e Mon Sep 17 00:00:00 2001 From: Nikolas Garza <90273783+nmgarza5@users.noreply.github.com> Date: Fri, 2 Oct 2026 19:58:45 +0000 Subject: [PATCH] feat(tenants): snapshot a migrated template schema per shard (#15278) --- backend/alembic/check_tenant_snapshot.py | 47 ++ backend/alembic/run_multitenant_migrations.py | 65 ++- ...5e0d2a7c9f41_add_tenant_schema_snapshot.py | 42 ++ backend/ee/onyx/db/tenant_snapshot.py | 486 ++++++++++++++++++ backend/onyx/db/engine/tenant_utils.py | 22 +- backend/onyx/db/models.py | 21 + backend/scripts/reencrypt_secrets.py | 23 +- backend/shared_configs/configs.py | 4 + .../db/test_tenant_snapshot_clone.py | 182 +++++++ 9 files changed, 866 insertions(+), 26 deletions(-) create mode 100755 backend/alembic/check_tenant_snapshot.py create mode 100644 backend/alembic_tenants/versions/5e0d2a7c9f41_add_tenant_schema_snapshot.py create mode 100644 backend/ee/onyx/db/tenant_snapshot.py create mode 100644 backend/tests/external_dependency_unit/db/test_tenant_snapshot_clone.py diff --git a/backend/alembic/check_tenant_snapshot.py b/backend/alembic/check_tenant_snapshot.py new file mode 100755 index 0000000000..140d5b268e --- /dev/null +++ b/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()) diff --git a/backend/alembic/run_multitenant_migrations.py b/backend/alembic/run_multitenant_migrations.py index c69d16d0c2..668349883a 100755 --- a/backend/alembic/run_multitenant_migrations.py +++ b/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__": diff --git a/backend/alembic_tenants/versions/5e0d2a7c9f41_add_tenant_schema_snapshot.py b/backend/alembic_tenants/versions/5e0d2a7c9f41_add_tenant_schema_snapshot.py new file mode 100644 index 0000000000..8d887e0c07 --- /dev/null +++ b/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") diff --git a/backend/ee/onyx/db/tenant_snapshot.py b/backend/ee/onyx/db/tenant_snapshot.py new file mode 100644 index 0000000000..fd345665f0 --- /dev/null +++ b/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'(? 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)) diff --git a/backend/onyx/db/engine/tenant_utils.py b/backend/onyx/db/engine/tenant_utils.py index 9aadbb286e..e8866fabc3 100644 --- a/backend/onyx/db/engine/tenant_utils.py +++ b/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 + ) diff --git a/backend/onyx/db/models.py b/backend/onyx/db/models.py index 167cdb9072..bb2ef8aef4 100644 --- a/backend/onyx/db/models.py +++ b/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" diff --git a/backend/scripts/reencrypt_secrets.py b/backend/scripts/reencrypt_secrets.py index ce58a8fd93..575a49b54c 100644 --- a/backend/scripts/reencrypt_secrets.py +++ b/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) diff --git a/backend/shared_configs/configs.py b/backend/shared_configs/configs.py index 62356f483d..51f1f6a36a 100644 --- a/backend/shared_configs/configs.py +++ b/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 = ( diff --git a/backend/tests/external_dependency_unit/db/test_tenant_snapshot_clone.py b/backend/tests/external_dependency_unit/db/test_tenant_snapshot_clone.py new file mode 100644 index 0000000000..dd5112b495 --- /dev/null +++ b/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()