diff --git a/scripts/bootstrap_shared_schema.py b/scripts/bootstrap_shared_schema.py index 9d67e948..2bc30a48 100644 --- a/scripts/bootstrap_shared_schema.py +++ b/scripts/bootstrap_shared_schema.py @@ -1,60 +1,67 @@ # honcho/scripts/bootstrap_shared_schema.py -"""Bootstrap the shared-tenant partitioned schema on a fresh shared database. +"""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). -""" +# region ai +# 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). +# endregion import os import sys -# Add the project root to the path (this script is run from the scripts directory). +# ai: 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 +# region ai # 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. +# endregion from src.models import Base # noqa: E402 # pyright: ignore +# region ai # 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. +# endregion PARTITION_COUNT = 128 def partitioned_tables() -> set[str]: - """Names of the tables declared with HASH(tenant_id) partitioning. + """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. - """ + # region ai + # Derived from the models, never hand-listed: a hand list drifts out of + # sync, and a partitioned parent with no partitions fails every insert. + # endregion return { table.name for table in Base.metadata.tables.values() @@ -79,31 +86,39 @@ def _create_hash_partitions(conn: Connection, table: Table) -> None: def bootstrap_shared_schema(conn: Connection) -> None: - """Create the shared partitioned schema on ``conn``. + """Create the shared partitioned schema on ``conn``.""" - ``conn`` must be in AUTOCOMMIT (see the module docstring on lock accumulation). - """ + # region ai + # ``conn`` must be in AUTOCOMMIT (see the module docstring on lock + # accumulation). + # endregion schema = Base.metadata.schema if schema and schema != "public": conn.execute(text(f'CREATE SCHEMA IF NOT EXISTS "{schema}"')) partitioned = partitioned_tables() + # region ai # Create tables in FK-dependency order (tenants first); a partitioned parent's # HASH partitions are created immediately after it. + # endregion for table in Base.metadata.sorted_tables: conn.execute(CreateTable(table)) if table.name in partitioned: _create_hash_partitions(conn, table) + # region ai # Indexes after every partition exists, so each partitioned index cascades to # all partitions. + # endregion for table in Base.metadata.sorted_tables: for index in table.indexes: conn.execute(CreateIndex(index)) if __name__ == "__main__": + # region ai # The app's async engine exposes a sync engine; CONNECTION_URI must point at # the fresh shared database. + # endregion from src.db import engine # noqa: E402 connection = engine.sync_engine.connect().execution_options( diff --git a/src/models.py b/src/models.py index 0255e89d..78535c5b 100644 --- a/src/models.py +++ b/src/models.py @@ -41,7 +41,7 @@ logger = getLogger(__name__) session_peers_table = Table( "session_peers", Base.metadata, - # tenant_id leads the all-natural-key PK and is the HASH partition key. + # ai: tenant_id leads the all-natural-key PK and is the HASH partition key. Column( "tenant_id", TEXT, @@ -108,27 +108,26 @@ session_peers_table = Table( @final class Tenant(Base): - """The tenant primitive: one row per tenant, the FK target for every - tenant-scoped table's ``tenant_id``. - - honcho's data plane and the control plane that owns the canonical tenant - registry live in separate databases, so a cross-database foreign key to the - real source of truth is impossible; this table is honcho's local mirror, - kept in sync from the control plane (backfilled for existing tenants; - written on provision or upserted on first authenticated request for new - ones). It is also the one-row-per-tenant home for facts that have nowhere - else to live: - - - ``app_name``: the tenant's original per-instance name. Keeping it lets a - tenant's external vector-store namespace stay stable when the tenant is - moved onto a shared backend, which avoids a full and very expensive - re-embed of its vectors. - - ``tier``: whether the tenant runs on a dedicated or a shared backend. - """ + """One row per tenant; the FK target for every tenant-scoped table's ``tenant_id``.""" + # region ai + # honcho's data plane and the control plane that owns the canonical tenant + # registry live in separate databases, so a cross-database FK to the real + # source of truth is impossible; this table is honcho's local mirror, kept in + # sync from the control plane (backfilled for existing tenants; written on + # provision, or upserted on first authenticated request, for new ones). It is + # also the one-row-per-tenant home for facts with nowhere else to live: + # - legacy_app_name: the tenant's original per-instance name. After the + # groudon shared-pool allocation, app_name is "shared" for every pooled + # tenant, so this column preserves the per-tenant value that keeps a + # tenant's external vector-store namespace stable across the move (avoids a + # full, expensive re-embed). Named "legacy_" so it never competes with the + # pool's app_name. + # - tier: whether the tenant runs on a dedicated or a shared backend. + # endregion __tablename__: str = "tenants" tenant_id: Mapped[str] = mapped_column(TEXT, primary_key=True) - app_name: Mapped[str | None] = mapped_column(TEXT, nullable=True, index=True) + legacy_app_name: Mapped[str | None] = mapped_column(TEXT, nullable=True, index=True) tier: Mapped[str] = mapped_column(TEXT, nullable=False, server_default="dedicated") created_at: Mapped[datetime.datetime] = mapped_column( DateTime(timezone=True), server_default=func.now() @@ -138,8 +137,10 @@ class Tenant(Base): @final class Workspace(Base): __tablename__: str = "workspaces" + # region ai # tenant_id is the HASH partition key and leads the composite PK, so it is # declared first. It FKs to the local tenants mirror (see Tenant). + # endregion tenant_id: Mapped[str] = mapped_column( TEXT, ForeignKey("tenants.tenant_id"), nullable=False ) @@ -167,9 +168,11 @@ class Workspace(Base): webhook_endpoints = relationship("WebhookEndpoint", back_populates="workspace") __table_args__ = ( + # region ai # Partitioned by HASH(tenant_id): Postgres requires the partition key in # the PK and in every UNIQUE. `name` is unique WITHIN a tenant, not # globally — many tenants share the SDK-default "default" workspace. + # endregion PrimaryKeyConstraint("tenant_id", "id"), UniqueConstraint("tenant_id", "name"), CheckConstraint("length(id) = 21", name="id_length"), @@ -196,8 +199,10 @@ class Peer(Base): created_at: Mapped[datetime.datetime] = mapped_column( DateTime(timezone=True), server_default=func.now(), index=True ) + # region ai # workspace_name's FK to workspaces is now the composite (below), since # workspaces.name is only unique within a tenant. + # endregion workspace_name: Mapped[str] = mapped_column(TEXT, nullable=False) configuration: Mapped[dict[str, Any]] = mapped_column( JSONB, default=dict, server_default=text("'{}'::jsonb") @@ -305,7 +310,7 @@ class Message(Base): __table_args__ = ( PrimaryKeyConstraint("tenant_id", "id"), - # (tenant_id, public_id) is the unique that message_embeddings' FK targets. + # ai: (tenant_id, public_id) is the unique that message_embeddings' FK targets. UniqueConstraint("tenant_id", "public_id"), CheckConstraint("length(public_id) = 21", name="public_id_length"), CheckConstraint("public_id ~ '^[A-Za-z0-9_-]+$'", name="public_id_format"), @@ -340,9 +345,11 @@ class Message(Base): "session_name", "seq_in_session", ), + # region ai # GIN can't lead with a scalar column without btree_gin; the table is # HASH(tenant_id)-partitioned, so this index is per-partition — queries # prune to one partition, then tenant_id filters the FTS candidates. + # endregion Index( "ix_messages_content_gin", text("to_tsvector('english', content)"), @@ -386,15 +393,19 @@ class MessageEmbedding(Base): __table_args__ = ( PrimaryKeyConstraint("tenant_id", "id"), + # region ai # message_id → messages.public_id is now composite: messages' unique is # (tenant_id, public_id) under partitioning. + # endregion ForeignKeyConstraint( ["tenant_id", "message_id"], ["messages.tenant_id", "messages.public_id"], ondelete="CASCADE", ), + # region ai # Composite FK to workspaces, for parity with the other tenant-scoped # tables (workspace_name is only unique within a tenant). + # endregion ForeignKeyConstraint( ["workspace_name", "tenant_id"], ["workspaces.name", "workspaces.tenant_id"], @@ -407,12 +418,16 @@ class MessageEmbedding(Base): ["peer_name", "workspace_name", "tenant_id"], ["peers.name", "peers.workspace_name", "peers.tenant_id"], ), + # region ai # 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. + # endregion Index("ix_message_embeddings_message_tenant", "message_id", "tenant_id"), + # region ai # HNSW is a single-column vector index (can't lead with tenant_id); it # becomes per-partition automatically under HASH(tenant_id). + # endregion Index( "ix_message_embeddings_embedding_hnsw", "embedding", @@ -420,9 +435,11 @@ class MessageEmbedding(Base): postgresql_with={"m": 16, "ef_construction": 64}, postgresql_ops={"embedding": "vector_cosine_ops"}, ), + # region ai # NOT tenant_id-leading on purpose: the reconciler scans this cross-tenant # (sync_state='pending' over all tenants), so a tenant_id prefix wouldn't # help. (Also drops the redundant single-column sync_state index.) + # endregion Index( "ix_message_embeddings_sync_state_last_sync_at", "sync_state", @@ -565,7 +582,8 @@ class Document(Base): ["session_name", "workspace_name", "tenant_id"], ["sessions.name", "sessions.workspace_name", "sessions.tenant_id"], ), - # Tenant-scoped collection lookups (replaces the single observer/observed indexes) + # Tenant-scoped collection lookups + # ai: replaces the single observer/observed indexes Index( "ix_documents_tenant_collection", "tenant_id", @@ -573,7 +591,7 @@ class Document(Base): "observed", "workspace_name", ), - # HNSW is a single-column vector index (per-partition under HASH(tenant_id)) + # ai: HNSW is a single-column vector index (per-partition under HASH(tenant_id)) Index( "ix_documents_embedding_hnsw", "embedding", @@ -589,8 +607,10 @@ class Document(Base): "source_ids", postgresql_using="gin", ), + # region ai # Reconciler scans this cross-tenant (sync_state='pending'), so NOT # tenant_id-leading. Also drops the redundant single-column sync_state index. + # endregion Index( "ix_documents_sync_state_last_sync_at", "sync_state", @@ -606,11 +626,13 @@ class QueueItem(Base): id: Mapped[int] = mapped_column( BigInteger, Identity(), primary_key=True, autoincrement=True ) + # region ai # Service table: NOT partitioned and drained-not-copied at migration, so it # keeps a sole-id PK. tenant_id is a plain attribution / fair-scheduling column # (no FK, no RLS); the FKs to the now-partitioned sessions / messages / # workspaces are dropped — the app manages queue lifecycle and already # tolerates missing referents. + # endregion tenant_id: Mapped[str | None] = mapped_column(TEXT, nullable=True, index=True) session_id: Mapped[str | None] = mapped_column(TEXT, nullable=True, index=True) work_unit_key: Mapped[str] = mapped_column(TEXT, nullable=False) @@ -665,7 +687,7 @@ class ActiveQueueSession(Base): id: Mapped[str] = mapped_column(TEXT, default=generate_nanoid, primary_key=True) - # Service table (unpartitioned): tenant_id is plain attribution, no FK / RLS. + # ai: Service table (unpartitioned): tenant_id is plain attribution, no FK / RLS. tenant_id: Mapped[str | None] = mapped_column(TEXT, nullable=True) work_unit_key: Mapped[str] = mapped_column(TEXT, unique=True) diff --git a/tests/scripts/test_bootstrap_shared_schema.py b/tests/scripts/test_bootstrap_shared_schema.py index 41b7bc29..60b8a921 100644 --- a/tests/scripts/test_bootstrap_shared_schema.py +++ b/tests/scripts/test_bootstrap_shared_schema.py @@ -1,16 +1,18 @@ -"""Integration test for the shared-tenant schema bootstrap. +"""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). -""" +# region ai +# 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). +# endregion import pytest from sqlalchemy import Engine, create_engine, make_url, text @@ -33,7 +35,7 @@ def bootstrapped_engine(worker_id: str): 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 '***'. + # ai: 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)')) @@ -43,7 +45,7 @@ def bootstrapped_engine(worker_id: str): _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. + # ai: 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) @@ -55,9 +57,12 @@ def bootstrapped_engine(worker_id: str): 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.""" + """Every tenant-scoped table is HASH-partitioned, carrying all its partitions.""" + + # region ai + # Guards the drift trap where a partitioned parent with zero partitions + # rejects every insert. + # endregion expected = partitioned_tables() assert len(expected) == 9 with bootstrapped_engine.connect() as conn: @@ -106,8 +111,10 @@ def test_tenant_scoped_uniqueness_and_fk(bootstrapped_engine: Engine) -> None: session.add_all([Tenant(tenant_id="tenant_a"), Tenant(tenant_id="tenant_b")]) session.commit() + # region ai # Two different tenants may each own a workspace named "default": this is # the whole point of UNIQUE(tenant_id, name) replacing UNIQUE(name). + # endregion session.add_all( [ Workspace(tenant_id="tenant_a", name="default"),