diff --git a/migrations/versions/09dfcd158c6c_add_tenant_id_first_class_primitive_and_.py b/migrations/versions/09dfcd158c6c_add_tenant_id_first_class_primitive_and_.py deleted file mode 100644 index ac0a1ae2..00000000 --- a/migrations/versions/09dfcd158c6c_add_tenant_id_first_class_primitive_and_.py +++ /dev/null @@ -1,104 +0,0 @@ -"""add tenant_id first class primitive and shared partitioned schema - -Bootstraps the shared-tenant target schema: a ``tenants`` registry table, a -``tenant_id`` column on every tenant-scoped table, composite ``(tenant_id, id)`` -primary keys, composite foreign keys, and ``HASH(tenant_id)`` partitioning on the -data tables. The service tables (``queue``, ``active_queue_sessions``) carry a -plain ``tenant_id`` and stay unpartitioned. - -This CREATES the partitioned tables from scratch — partitioning cannot be -introduced by ``ALTER TABLE`` — so it bootstraps a fresh shared database that -per-tenant data is consolidated into, rather than transforming an existing -single-tenant database in place. - -The DDL is generated from the declarative models (``src.models.Base.metadata``) -rather than hand-transcribed. The models are the single source of truth for the -target schema; the compiler renders ``PARTITION BY`` from each table's -``postgresql_partition_by`` option and applies the naming convention; and this -keeps the bootstrap exactly in sync with the accepted data model. Only the -per-table HASH partitions (which the model layer does not enumerate) are created -explicitly. The target schema and ``search_path`` are established by -``migrations/env.py`` before this runs. - -Revision ID: 09dfcd158c6c -Revises: e4eba9cfaa6f -Create Date: 2026-09-02 11:38:27.736205 - -""" - -from collections.abc import Sequence - -from alembic import op -from sqlalchemy import Table -from sqlalchemy.schema import CreateIndex, CreateTable - -# Importing Base from src.models (rather than src.db) also registers every model -# on Base.metadata as a side effect — that populated metadata is the whole schema -# this migration builds. -from src.models import Base # pyright: ignore[reportPrivateLocalImportUsage] - -# revision identifiers, used by Alembic. -revision: str = "09dfcd158c6c" -down_revision: str | None = "e4eba9cfaa6f" -branch_labels: str | Sequence[str] | None = None -depends_on: str | Sequence[str] | None = None - -# HASH(tenant_id) partition count for the tenant-scoped tables. Sized so the -# largest table (a couple hundred GB fleet-wide) lands at roughly a couple GB per -# partition — small enough to stay maintainable and well under a node's memory, -# with headroom for growth. Contention is not the driver (write volume is tiny). -PARTITION_COUNT = 128 - -# Tenant-scoped data tables partitioned by HASH(tenant_id). The service tables -# (queue, active_queue_sessions) and the tenants registry are NOT partitioned. -PARTITIONED_TABLES = frozenset( - { - "workspaces", - "peers", - "sessions", - "messages", - "message_embeddings", - "collections", - "documents", - "session_peers", - "webhook_endpoints", - } -) - - -def _qualified(table: Table) -> str: - return f'"{table.schema}"."{table.name}"' if table.schema else f'"{table.name}"' - - -def _create_hash_partitions(table: Table) -> None: - """Create the N HASH partitions for a partitioned parent table.""" - schema_prefix = f'"{table.schema}".' if table.schema else "" - for remainder in range(PARTITION_COUNT): - op.execute( - f'CREATE TABLE {schema_prefix}"{table.name}_p{remainder:03d}"' - + f" PARTITION OF {_qualified(table)}" - + f" FOR VALUES WITH (MODULUS {PARTITION_COUNT}, REMAINDER {remainder})" - ) - - -def upgrade() -> None: - metadata = Base.metadata - # Create tables in FK-dependency order (tenants first). A partitioned - # parent's HASH partitions are created immediately after the parent. - for table in metadata.sorted_tables: - op.execute(CreateTable(table)) - if table.name in PARTITIONED_TABLES: - _create_hash_partitions(table) - # Create indexes once every partition exists, so each partitioned index - # cascades onto all partitions. - for table in metadata.sorted_tables: - for index in table.indexes: - op.execute(CreateIndex(index)) - - -def downgrade() -> None: - metadata = Base.metadata - # Reverse dependency order; CASCADE drops each partitioned table's partitions - # (and any dependent FKs) along with the parent. - for table in reversed(metadata.sorted_tables): - op.execute(f"DROP TABLE IF EXISTS {_qualified(table)} CASCADE") diff --git a/scripts/bootstrap_shared_schema.py b/scripts/bootstrap_shared_schema.py new file mode 100644 index 00000000..9d67e948 --- /dev/null +++ b/scripts/bootstrap_shared_schema.py @@ -0,0 +1,114 @@ +# honcho/scripts/bootstrap_shared_schema.py +"""Bootstrap the shared-tenant partitioned schema on a fresh shared database. + +This is deliberately NOT an Alembic migration. As a chained revision it would run +after the per-tenant migration history that already creates these tables +(non-partitioned): ``alembic upgrade head`` would collide (DuplicateTable) and, +worse, it would run on existing single-tenant instances via ``init_db()`` and +error there too. Instead this is a standalone bootstrap that the +migration/consolidation track runs explicitly against the fresh shared database +that per-tenant data is consolidated into. How that database is provisioned and +version-stamped is owned by the migration track. + +The DDL is generated from the declarative models (``src.models.Base.metadata``): +the models are the single source of truth, the compiler renders ``PARTITION BY`` +from each table's ``postgresql_partition_by`` option, and this stays in lockstep +with the data model. Only the per-table HASH partitions (which the model layer +does not enumerate) are created explicitly, and the set of partitioned tables is +derived from the models — never hand-listed — so a newly partitioned model can't +silently ship a parent with zero partitions (which would fail every insert). + +Statements run in AUTOCOMMIT: the partitioned parents times ``PARTITION_COUNT`` +partitions, plus their composite FKs (each child FK locks every partition of the +referenced parent), accumulate tens of thousands of locks, which overflows +``max_locks_per_transaction`` if run in a single transaction. Committing per +statement releases each partition's locks incrementally. + +Prerequisite: the ``vector`` extension must already be installed (for the +embedding columns and their HNSW indexes). +""" + +import os +import sys + +# Add the project root to the path (this script is run from the scripts directory). +project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) +sys.path.insert(0, project_root) + +from sqlalchemy import Connection, Table, text # noqa: E402 +from sqlalchemy.schema import CreateIndex, CreateTable # noqa: E402 + +# Importing Base from src.models (rather than src.db) also registers every model +# on Base.metadata as a side effect — that populated metadata is the schema this +# script builds. +from src.models import Base # noqa: E402 # pyright: ignore + +# HASH(tenant_id) partition count. Sized so each partition stays around a couple +# GB — well under a node's memory, with headroom for growth. Contention is not the +# driver (write volume is low); this is a maintenance/pruning choice. +PARTITION_COUNT = 128 + + +def partitioned_tables() -> set[str]: + """Names of the tables declared with HASH(tenant_id) partitioning. + + Derived from the models, never hand-listed: a hand list drifts out of sync, + and a partitioned parent with no partitions fails every insert. + """ + return { + table.name + for table in Base.metadata.tables.values() + if table.dialect_options["postgresql"].get("partition_by") + } + + +def _qualified(table: Table) -> str: + return f'"{table.schema}"."{table.name}"' if table.schema else f'"{table.name}"' + + +def _create_hash_partitions(conn: Connection, table: Table) -> None: + schema_prefix = f'"{table.schema}".' if table.schema else "" + for remainder in range(PARTITION_COUNT): + conn.execute( + text( + f'CREATE TABLE {schema_prefix}"{table.name}_p{remainder:03d}"' + + f" PARTITION OF {_qualified(table)}" + + f" FOR VALUES WITH (MODULUS {PARTITION_COUNT}, REMAINDER {remainder})" + ) + ) + + +def bootstrap_shared_schema(conn: Connection) -> None: + """Create the shared partitioned schema on ``conn``. + + ``conn`` must be in AUTOCOMMIT (see the module docstring on lock accumulation). + """ + schema = Base.metadata.schema + if schema and schema != "public": + conn.execute(text(f'CREATE SCHEMA IF NOT EXISTS "{schema}"')) + + partitioned = partitioned_tables() + # Create tables in FK-dependency order (tenants first); a partitioned parent's + # HASH partitions are created immediately after it. + for table in Base.metadata.sorted_tables: + conn.execute(CreateTable(table)) + if table.name in partitioned: + _create_hash_partitions(conn, table) + # Indexes after every partition exists, so each partitioned index cascades to + # all partitions. + for table in Base.metadata.sorted_tables: + for index in table.indexes: + conn.execute(CreateIndex(index)) + + +if __name__ == "__main__": + # The app's async engine exposes a sync engine; CONNECTION_URI must point at + # the fresh shared database. + from src.db import engine # noqa: E402 + + connection = engine.sync_engine.connect().execution_options( + isolation_level="AUTOCOMMIT" + ) + with connection as conn: + bootstrap_shared_schema(conn) + print("Shared-tenant schema bootstrapped.") diff --git a/src/models.py b/src/models.py index 75fd155b..0255e89d 100644 --- a/src/models.py +++ b/src/models.py @@ -327,7 +327,13 @@ class Message(Base): "id", postgresql_include=["created_at"], ), - Index("ix_messages_tenant_peer", "tenant_id", "peer_name"), + Index( + "ix_messages_peer_lookup", + "tenant_id", + "workspace_name", + "peer_name", + "created_at", + ), UniqueConstraint( "tenant_id", "workspace_name", @@ -387,6 +393,12 @@ class MessageEmbedding(Base): ["messages.tenant_id", "messages.public_id"], ondelete="CASCADE", ), + # Composite FK to workspaces, for parity with the other tenant-scoped + # tables (workspace_name is only unique within a tenant). + ForeignKeyConstraint( + ["workspace_name", "tenant_id"], + ["workspaces.name", "workspaces.tenant_id"], + ), ForeignKeyConstraint( ["session_name", "workspace_name", "tenant_id"], ["sessions.name", "sessions.workspace_name", "sessions.tenant_id"], @@ -395,7 +407,10 @@ class MessageEmbedding(Base): ["peer_name", "workspace_name", "tenant_id"], ["peers.name", "peers.workspace_name", "peers.tenant_id"], ), - Index("ix_message_embeddings_tenant_message", "tenant_id", "message_id"), + # message_id-leading: every lookup on message_id is cross-tenant (the + # reconciler / embed_now filter by message_id with no tenant_id in scope), + # so a tenant_id prefix would force a scan of all partitions. + Index("ix_message_embeddings_message_tenant", "message_id", "tenant_id"), # HNSW is a single-column vector index (can't lead with tenant_id); it # becomes per-partition automatically under HASH(tenant_id). Index( @@ -641,7 +656,7 @@ class QueueItem(Base): ) def __repr__(self) -> str: - return f"QueueItem(id={self.id}, session_id={self.session_id}, work_unit_key={self.work_unit_key}, task_type={self.task_type}, payload={self.payload}, processed={self.processed}, workspace_name={self.workspace_name}, message_id={self.message_id})" + return f"QueueItem(id={self.id}, tenant_id={self.tenant_id}, session_id={self.session_id}, work_unit_key={self.work_unit_key}, task_type={self.task_type}, payload={self.payload}, processed={self.processed}, workspace_name={self.workspace_name}, message_id={self.message_id})" @final