From 70c4584e7c76bd0ceb770f7f97b7e4b3a27b4708 Mon Sep 17 00:00:00 2001 From: Phil Date: Wed, 2 Sep 2026 12:36:23 -0400 Subject: [PATCH 1/7] feat(models): make tenant_id a first-class primitive for shared-tenant schema Introduce a `tenants` registry table and a `tenant_id` column on every tenant-scoped table, with composite `(tenant_id, id)` primary keys, composite foreign keys (natural-key chains plus a `tenant_id` -> tenants FK), and `HASH(tenant_id)` partitioning (N=128) so many tenants can share one database with per-tenant isolation. `workspaces.name` uniqueness becomes `UNIQUE(tenant_id, name)` and indexes lead with `tenant_id`. Service tables (queue, active_queue_sessions) carry a plain `tenant_id` for attribution and stay unpartitioned; their FKs to the now-partitioned parents are dropped. Adds the bootstrap migration that creates the partitioned shared schema from the declarative models (written, not run). Co-Authored-By: Claude Opus 4.8 --- ...dd_tenant_id_first_class_primitive_and_.py | 104 ++++++ src/models.py | 318 ++++++++++++------ 2 files changed, 326 insertions(+), 96 deletions(-) create mode 100644 migrations/versions/09dfcd158c6c_add_tenant_id_first_class_primitive_and_.py 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 new file mode 100644 index 00000000..ac0a1ae2 --- /dev/null +++ b/migrations/versions/09dfcd158c6c_add_tenant_id_first_class_primitive_and_.py @@ -0,0 +1,104 @@ +"""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/src/models.py b/src/models.py index 6433225a..75fd155b 100644 --- a/src/models.py +++ b/src/models.py @@ -1,6 +1,6 @@ import datetime from logging import getLogger -from typing import Any, final +from typing import Any, final, override from dotenv import load_dotenv from nanoid import generate as generate_nanoid @@ -16,6 +16,7 @@ from sqlalchemy import ( Identity, Index, Integer, + PrimaryKeyConstraint, Table, UniqueConstraint, text, @@ -23,7 +24,6 @@ from sqlalchemy import ( from sqlalchemy.dialects.postgresql import JSONB, TEXT from sqlalchemy.orm import Mapped, MappedColumn, mapped_column, relationship from sqlalchemy.sql import func -from typing_extensions import override from src.config import settings from src.utils.types import DocumentLevel, TaskType, VectorSyncState @@ -41,10 +41,17 @@ 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. + Column( + "tenant_id", + TEXT, + ForeignKey("tenants.tenant_id"), + primary_key=True, + nullable=False, + ), Column( "workspace_name", TEXT, - ForeignKey("workspaces.name"), primary_key=True, nullable=False, ), @@ -80,24 +87,64 @@ session_peers_table = Table( DateTime(timezone=True), nullable=True, ), + # Composite foreign key constraint for workspaces + ForeignKeyConstraint( + ["workspace_name", "tenant_id"], + ["workspaces.name", "workspaces.tenant_id"], + ), # Composite foreign key constraint for sessions ForeignKeyConstraint( - ["session_name", "workspace_name"], - ["sessions.name", "sessions.workspace_name"], + ["session_name", "workspace_name", "tenant_id"], + ["sessions.name", "sessions.workspace_name", "sessions.tenant_id"], ), # Composite foreign key constraint for peers ForeignKeyConstraint( - ["peer_name", "workspace_name"], - ["peers.name", "peers.workspace_name"], + ["peer_name", "workspace_name", "tenant_id"], + ["peers.name", "peers.workspace_name", "peers.tenant_id"], ), + postgresql_partition_by="HASH (tenant_id)", ) +@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. + """ + + __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) + 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() + ) + + @final class Workspace(Base): __tablename__: str = "workspaces" - id: Mapped[str] = mapped_column(TEXT, default=generate_nanoid, primary_key=True) - name: Mapped[str] = mapped_column(TEXT, unique=True) + # 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). + tenant_id: Mapped[str] = mapped_column( + TEXT, ForeignKey("tenants.tenant_id"), nullable=False + ) + id: Mapped[str] = mapped_column(TEXT, default=generate_nanoid) + name: Mapped[str] = mapped_column(TEXT) created_at: Mapped[datetime.datetime] = mapped_column( DateTime(timezone=True), server_default=func.now(), index=True ) @@ -120,16 +167,25 @@ class Workspace(Base): webhook_endpoints = relationship("WebhookEndpoint", back_populates="workspace") __table_args__ = ( + # 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. + PrimaryKeyConstraint("tenant_id", "id"), + UniqueConstraint("tenant_id", "name"), CheckConstraint("length(id) = 21", name="id_length"), CheckConstraint("length(name) <= 512", name="name_length"), CheckConstraint("id ~ '^[A-Za-z0-9_-]+$'", name="id_format"), + {"postgresql_partition_by": "HASH (tenant_id)"}, ) @final class Peer(Base): __tablename__: str = "peers" - id: Mapped[str] = mapped_column(TEXT, default=generate_nanoid, primary_key=True) + tenant_id: Mapped[str] = mapped_column( + TEXT, ForeignKey("tenants.tenant_id"), nullable=False + ) + id: Mapped[str] = mapped_column(TEXT, default=generate_nanoid) name: Mapped[str] = mapped_column(TEXT, nullable=False) h_metadata: Mapped[dict[str, Any]] = mapped_column( "metadata", JSONB, default=dict, server_default=text("'{}'::jsonb") @@ -140,9 +196,9 @@ class Peer(Base): created_at: Mapped[datetime.datetime] = mapped_column( DateTime(timezone=True), server_default=func.now(), index=True ) - workspace_name: Mapped[str] = mapped_column( - ForeignKey("workspaces.name"), nullable=False, index=True - ) + # workspace_name's FK to workspaces is now the composite (below), since + # workspaces.name is only unique within a tenant. + workspace_name: Mapped[str] = mapped_column(TEXT, nullable=False) configuration: Mapped[dict[str, Any]] = mapped_column( JSONB, default=dict, server_default=text("'{}'::jsonb") ) @@ -153,20 +209,30 @@ class Peer(Base): ) __table_args__ = ( - UniqueConstraint("name", "workspace_name"), + PrimaryKeyConstraint("tenant_id", "id"), + UniqueConstraint("tenant_id", "name", "workspace_name"), + ForeignKeyConstraint( + ["workspace_name", "tenant_id"], + ["workspaces.name", "workspaces.tenant_id"], + ), + Index("ix_peers_tenant_workspace", "tenant_id", "workspace_name"), CheckConstraint("length(id) = 21", name="id_length"), CheckConstraint("length(name) <= 512", name="name_length"), CheckConstraint("id ~ '^[A-Za-z0-9_-]+$'", name="id_format"), + {"postgresql_partition_by": "HASH (tenant_id)"}, ) def __repr__(self) -> str: - return f"Peer(id={self.id}, name={self.name}, workspace_name={self.workspace_name}, created_at={self.created_at}, h_metadata={self.h_metadata}, configuration={self.configuration})" + return f"Peer(tenant_id={self.tenant_id}, id={self.id}, name={self.name}, workspace_name={self.workspace_name}, created_at={self.created_at}, h_metadata={self.h_metadata}, configuration={self.configuration})" @final class Session(Base): __tablename__: str = "sessions" - id: Mapped[str] = mapped_column(TEXT, primary_key=True, default=generate_nanoid) + tenant_id: Mapped[str] = mapped_column( + TEXT, ForeignKey("tenants.tenant_id"), nullable=False + ) + id: Mapped[str] = mapped_column(TEXT, default=generate_nanoid) name: Mapped[str] = mapped_column(TEXT) is_active: Mapped[bool] = mapped_column(default=True, server_default=text("true")) h_metadata: Mapped[dict[str, Any]] = mapped_column( @@ -178,9 +244,7 @@ class Session(Base): created_at: Mapped[datetime.datetime] = mapped_column( DateTime(timezone=True), server_default=func.now(), index=True ) - workspace_name: Mapped[str] = mapped_column( - ForeignKey("workspaces.name"), nullable=False, index=True - ) + workspace_name: Mapped[str] = mapped_column(TEXT, nullable=False) configuration: Mapped[dict[str, Any]] = mapped_column( JSONB, default=dict, server_default=text("'{}'::jsonb") ) @@ -192,27 +256,31 @@ class Session(Base): messages = relationship("Message", back_populates="session") __table_args__ = ( - UniqueConstraint("name", "workspace_name"), + PrimaryKeyConstraint("tenant_id", "id"), + UniqueConstraint("tenant_id", "name", "workspace_name"), + ForeignKeyConstraint( + ["workspace_name", "tenant_id"], + ["workspaces.name", "workspaces.tenant_id"], + ), + Index("ix_sessions_tenant_workspace", "tenant_id", "workspace_name"), CheckConstraint("length(name) <= 512", name="name_length"), CheckConstraint("length(id) = 21", name="id_length"), CheckConstraint("id ~ '^[A-Za-z0-9_-]+$'", name="id_format"), + {"postgresql_partition_by": "HASH (tenant_id)"}, ) def __repr__(self) -> str: - return f"Session(id={self.id}, name={self.name}, workspace_name={self.workspace_name}, is_active={self.is_active}, created_at={self.created_at}, h_metadata={self.h_metadata})" + return f"Session(tenant_id={self.tenant_id}, id={self.id}, name={self.name}, workspace_name={self.workspace_name}, is_active={self.is_active}, created_at={self.created_at}, h_metadata={self.h_metadata})" @final class Message(Base): __tablename__: str = "messages" - id: Mapped[int] = mapped_column( - BigInteger, Identity(), primary_key=True, autoincrement=True - ) - public_id: Mapped[str] = mapped_column( - TEXT, - unique=True, - default=generate_nanoid, + tenant_id: Mapped[str] = mapped_column( + TEXT, ForeignKey("tenants.tenant_id"), nullable=False ) + id: Mapped[int] = mapped_column(BigInteger, Identity(), autoincrement=True) + public_id: Mapped[str] = mapped_column(TEXT, default=generate_nanoid) # NOTE: Messages in Honcho 2.0 could historically be stored outside of a session. # We have since assigned all of these messages to a default session. session_name: Mapped[str] = mapped_column(TEXT, nullable=False) @@ -230,72 +298,78 @@ class Message(Base): DateTime(timezone=True), server_default=func.now(), index=True ) # Note: Foreign key relationships established via composite ForeignKeyConstraint below - peer_name: Mapped[str] = mapped_column(TEXT, index=True) - workspace_name: Mapped[str] = mapped_column(TEXT, index=True) + peer_name: Mapped[str] = mapped_column(TEXT) + workspace_name: Mapped[str] = mapped_column(TEXT) session = relationship("Session", back_populates="messages") __table_args__ = ( + PrimaryKeyConstraint("tenant_id", "id"), + # (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"), CheckConstraint("length(content) <= 65535", name="content_length"), # Composite foreign key constraint for sessions ForeignKeyConstraint( - ["session_name", "workspace_name"], - ["sessions.name", "sessions.workspace_name"], + ["session_name", "workspace_name", "tenant_id"], + ["sessions.name", "sessions.workspace_name", "sessions.tenant_id"], ), # Composite foreign key constraint for peers ForeignKeyConstraint( - ["peer_name", "workspace_name"], - ["peers.name", "peers.workspace_name"], + ["peer_name", "workspace_name", "tenant_id"], + ["peers.name", "peers.workspace_name", "peers.tenant_id"], ), Index( "ix_messages_session_lookup", + "tenant_id", "session_name", "id", - postgresql_include=["id", "created_at"], + postgresql_include=["created_at"], ), + Index("ix_messages_tenant_peer", "tenant_id", "peer_name"), UniqueConstraint( + "tenant_id", "workspace_name", "session_name", "seq_in_session", ), - # Full text search index on content column + # 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. Index( "ix_messages_content_gin", text("to_tsvector('english', content)"), postgresql_using="gin", ), + {"postgresql_partition_by": "HASH (tenant_id)"}, ) @override def __repr__(self) -> str: - return f"Message(id={self.id}, session_name={self.session_name}, peer_name={self.peer_name}, content={self.content})" + return f"Message(tenant_id={self.tenant_id}, id={self.id}, session_name={self.session_name}, peer_name={self.peer_name}, content={self.content})" @final class MessageEmbedding(Base): __tablename__: str = "message_embeddings" - id: Mapped[int] = mapped_column( - BigInteger, Identity(), primary_key=True, autoincrement=True + tenant_id: Mapped[str] = mapped_column( + TEXT, ForeignKey("tenants.tenant_id"), nullable=False ) + id: Mapped[int] = mapped_column(BigInteger, Identity(), autoincrement=True) content: Mapped[str] = mapped_column(TEXT) embedding: MappedColumn[Any] = mapped_column(Vector(_VECTOR_DIM), nullable=True) - message_id: Mapped[str] = mapped_column( - ForeignKey("messages.public_id", ondelete="CASCADE"), nullable=False, index=True - ) - workspace_name: Mapped[str] = mapped_column( - ForeignKey("workspaces.name"), nullable=False, index=True - ) - session_name: Mapped[str] = mapped_column(TEXT, nullable=False, index=True) - peer_name: Mapped[str] = mapped_column(TEXT, nullable=False, index=True) + message_id: Mapped[str] = mapped_column(TEXT, nullable=False) + workspace_name: Mapped[str] = mapped_column(TEXT, nullable=False) + session_name: Mapped[str] = mapped_column(TEXT, nullable=False) + peer_name: Mapped[str] = mapped_column(TEXT, nullable=False) created_at: Mapped[datetime.datetime] = mapped_column( DateTime(timezone=True), server_default=func.now(), index=True ) # Vector sync state tracking sync_state: Mapped[VectorSyncState] = mapped_column( - TEXT, nullable=False, server_default="pending", index=True + TEXT, nullable=False, server_default="pending" ) last_sync_at: Mapped[datetime.datetime | None] = mapped_column( DateTime(timezone=True), nullable=True @@ -305,16 +379,25 @@ class MessageEmbedding(Base): ) __table_args__ = ( - # Compound foreign key constraints + PrimaryKeyConstraint("tenant_id", "id"), + # message_id → messages.public_id is now composite: messages' unique is + # (tenant_id, public_id) under partitioning. ForeignKeyConstraint( - ["session_name", "workspace_name"], - ["sessions.name", "sessions.workspace_name"], + ["tenant_id", "message_id"], + ["messages.tenant_id", "messages.public_id"], + ondelete="CASCADE", ), ForeignKeyConstraint( - ["peer_name", "workspace_name"], - ["peers.name", "peers.workspace_name"], + ["session_name", "workspace_name", "tenant_id"], + ["sessions.name", "sessions.workspace_name", "sessions.tenant_id"], ), - # HNSW index on embedding column for efficient similarity search + ForeignKeyConstraint( + ["peer_name", "workspace_name", "tenant_id"], + ["peers.name", "peers.workspace_name", "peers.tenant_id"], + ), + Index("ix_message_embeddings_tenant_message", "tenant_id", "message_id"), + # HNSW is a single-column vector index (can't lead with tenant_id); it + # becomes per-partition automatically under HASH(tenant_id). Index( "ix_message_embeddings_embedding_hnsw", "embedding", @@ -322,12 +405,15 @@ class MessageEmbedding(Base): postgresql_with={"m": 16, "ef_construction": 64}, postgresql_ops={"embedding": "vector_cosine_ops"}, ), - # Composite index for efficient reconciliation queries + # 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.) Index( "ix_message_embeddings_sync_state_last_sync_at", "sync_state", "last_sync_at", ), + {"postgresql_partition_by": "HASH (tenant_id)"}, ) @@ -335,9 +421,12 @@ class MessageEmbedding(Base): class Collection(Base): __tablename__: str = "collections" - id: Mapped[str] = mapped_column(TEXT, default=generate_nanoid, primary_key=True) - observer: Mapped[str] = mapped_column(TEXT, index=True) - observed: Mapped[str] = mapped_column(TEXT, index=True) + tenant_id: Mapped[str] = mapped_column( + TEXT, ForeignKey("tenants.tenant_id"), nullable=False + ) + id: Mapped[str] = mapped_column(TEXT, default=generate_nanoid) + observer: Mapped[str] = mapped_column(TEXT) + observed: Mapped[str] = mapped_column(TEXT) created_at: Mapped[datetime.datetime] = mapped_column( DateTime(timezone=True), server_default=func.now(), index=True ) @@ -350,35 +439,43 @@ class Collection(Base): documents = relationship( "Document", back_populates="collection", cascade="all, delete, delete-orphan" ) - workspace_name: Mapped[str] = mapped_column( - ForeignKey("workspaces.name"), nullable=False, index=True - ) + workspace_name: Mapped[str] = mapped_column(TEXT, nullable=False) __table_args__ = ( + PrimaryKeyConstraint("tenant_id", "id"), UniqueConstraint( + "tenant_id", "observer", "observed", "workspace_name", ), CheckConstraint("length(id) = 21", name="id_length"), CheckConstraint("id ~ '^[A-Za-z0-9_-]+$'", name="id_format"), + ForeignKeyConstraint( + ["workspace_name", "tenant_id"], + ["workspaces.name", "workspaces.tenant_id"], + ), # Composite foreign key constraint for observer peer ForeignKeyConstraint( - ["observer", "workspace_name"], - ["peers.name", "peers.workspace_name"], + ["observer", "workspace_name", "tenant_id"], + ["peers.name", "peers.workspace_name", "peers.tenant_id"], ), # Composite foreign key constraint for observed peer ForeignKeyConstraint( - ["observed", "workspace_name"], - ["peers.name", "peers.workspace_name"], + ["observed", "workspace_name", "tenant_id"], + ["peers.name", "peers.workspace_name", "peers.tenant_id"], ), + {"postgresql_partition_by": "HASH (tenant_id)"}, ) @final class Document(Base): __tablename__: str = "documents" - id: Mapped[str] = mapped_column(TEXT, default=generate_nanoid, primary_key=True) + tenant_id: Mapped[str] = mapped_column( + TEXT, ForeignKey("tenants.tenant_id"), nullable=False + ) + id: Mapped[str] = mapped_column(TEXT, default=generate_nanoid) internal_metadata: Mapped[dict[str, Any]] = mapped_column( "internal_metadata", JSONB, default=dict, server_default=text("'{}'::jsonb") ) @@ -397,19 +494,17 @@ class Document(Base): DateTime(timezone=True), server_default=func.now(), index=True ) - observer: Mapped[str] = mapped_column(TEXT, index=True) - observed: Mapped[str] = mapped_column(TEXT, index=True) - workspace_name: Mapped[str] = mapped_column( - ForeignKey("workspaces.name"), nullable=False, index=True - ) - session_name: Mapped[str | None] = mapped_column(TEXT, nullable=True, index=True) + observer: Mapped[str] = mapped_column(TEXT) + observed: Mapped[str] = mapped_column(TEXT) + workspace_name: Mapped[str] = mapped_column(TEXT, nullable=False) + session_name: Mapped[str | None] = mapped_column(TEXT, nullable=True) deleted_at: Mapped[datetime.datetime | None] = mapped_column( DateTime(timezone=True), nullable=True, index=True, default=None ) # Vector sync state tracking sync_state: Mapped[VectorSyncState] = mapped_column( - TEXT, nullable=False, server_default="pending", index=True + TEXT, nullable=False, server_default="pending" ) last_sync_at: Mapped[datetime.datetime | None] = mapped_column( DateTime(timezone=True), nullable=True @@ -421,34 +516,49 @@ class Document(Base): collection = relationship("Collection", back_populates="documents") __table_args__ = ( + PrimaryKeyConstraint("tenant_id", "id"), CheckConstraint("length(id) = 21", name="id_length"), CheckConstraint("length(content) <= 65535", name="content_length"), CheckConstraint("id ~ '^[A-Za-z0-9_-]+$'", name="id_format"), + # Composite foreign key constraint for workspaces + ForeignKeyConstraint( + ["workspace_name", "tenant_id"], + ["workspaces.name", "workspaces.tenant_id"], + ), # Composite foreign key constraint for collections ForeignKeyConstraint( - ["observer", "observed", "workspace_name"], + ["observer", "observed", "workspace_name", "tenant_id"], [ "collections.observer", "collections.observed", "collections.workspace_name", + "collections.tenant_id", ], ), # Composite foreign key constraint for observer peer ForeignKeyConstraint( - ["observer", "workspace_name"], - ["peers.name", "peers.workspace_name"], + ["observer", "workspace_name", "tenant_id"], + ["peers.name", "peers.workspace_name", "peers.tenant_id"], ), # Composite foreign key constraint for observed peer ForeignKeyConstraint( - ["observed", "workspace_name"], - ["peers.name", "peers.workspace_name"], + ["observed", "workspace_name", "tenant_id"], + ["peers.name", "peers.workspace_name", "peers.tenant_id"], ), # Composite foreign key constraint for sessions ForeignKeyConstraint( - ["session_name", "workspace_name"], - ["sessions.name", "sessions.workspace_name"], + ["session_name", "workspace_name", "tenant_id"], + ["sessions.name", "sessions.workspace_name", "sessions.tenant_id"], ), - # HNSW index on embedding column + # Tenant-scoped collection lookups (replaces the single observer/observed indexes) + Index( + "ix_documents_tenant_collection", + "tenant_id", + "observer", + "observed", + "workspace_name", + ), + # HNSW is a single-column vector index (per-partition under HASH(tenant_id)) Index( "ix_documents_embedding_hnsw", "embedding", @@ -464,12 +574,14 @@ class Document(Base): "source_ids", postgresql_using="gin", ), - # Composite index for efficient reconciliation queries + # Reconciler scans this cross-tenant (sync_state='pending'), so NOT + # tenant_id-leading. Also drops the redundant single-column sync_state index. Index( "ix_documents_sync_state_last_sync_at", "sync_state", "last_sync_at", ), + {"postgresql_partition_by": "HASH (tenant_id)"}, ) @@ -479,9 +591,13 @@ class QueueItem(Base): id: Mapped[int] = mapped_column( BigInteger, Identity(), primary_key=True, autoincrement=True ) - session_id: Mapped[str | None] = mapped_column( - ForeignKey("sessions.id"), nullable=True, index=True - ) + # 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. + 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) task_type: Mapped[TaskType] = mapped_column(TEXT, nullable=False) @@ -493,12 +609,8 @@ class QueueItem(Base): created_at: Mapped[datetime.datetime] = mapped_column( DateTime(timezone=True), server_default=func.now(), index=True ) - workspace_name: Mapped[str | None] = mapped_column( - ForeignKey("workspaces.name"), nullable=True, index=True - ) - message_id: Mapped[int | None] = mapped_column( - BigInteger, ForeignKey("messages.id"), nullable=True - ) + workspace_name: Mapped[str | None] = mapped_column(TEXT, nullable=True, index=True) + message_id: Mapped[int | None] = mapped_column(BigInteger, nullable=True) __table_args__ = ( Index( @@ -538,6 +650,9 @@ 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. + tenant_id: Mapped[str | None] = mapped_column(TEXT, nullable=True) + work_unit_key: Mapped[str] = mapped_column(TEXT, unique=True) last_updated: Mapped[datetime.datetime] = mapped_column( @@ -548,10 +663,11 @@ class ActiveQueueSession(Base): @final class WebhookEndpoint(Base): __tablename__: str = "webhook_endpoints" - id: Mapped[str] = mapped_column(TEXT, default=generate_nanoid, primary_key=True) - workspace_name: Mapped[str] = mapped_column( - ForeignKey("workspaces.name"), nullable=False, index=True + tenant_id: Mapped[str] = mapped_column( + TEXT, ForeignKey("tenants.tenant_id"), nullable=False ) + id: Mapped[str] = mapped_column(TEXT, default=generate_nanoid) + workspace_name: Mapped[str] = mapped_column(TEXT, nullable=False) url: Mapped[str] = mapped_column(TEXT, nullable=False) created_at: Mapped[datetime.datetime] = mapped_column( DateTime(timezone=True), server_default=func.now() @@ -559,10 +675,19 @@ class WebhookEndpoint(Base): workspace = relationship("Workspace", back_populates="webhook_endpoints") - __table_args__ = (CheckConstraint("length(url) <= 2048", name="url_length"),) + __table_args__ = ( + PrimaryKeyConstraint("tenant_id", "id"), + ForeignKeyConstraint( + ["workspace_name", "tenant_id"], + ["workspaces.name", "workspaces.tenant_id"], + ), + Index("ix_webhook_endpoints_tenant_workspace", "tenant_id", "workspace_name"), + CheckConstraint("length(url) <= 2048", name="url_length"), + {"postgresql_partition_by": "HASH (tenant_id)"}, + ) def __repr__(self) -> str: - return f"WebhookEndpoint(id={self.id}, workspace_name={self.workspace_name}, url={self.url})" + return f"WebhookEndpoint(tenant_id={self.tenant_id}, id={self.id}, workspace_name={self.workspace_name}, url={self.url})" @final @@ -570,6 +695,7 @@ class SessionPeer(Base): __table__: Table = session_peers_table # Type annotations for the columns + tenant_id: Mapped[str] workspace_name: Mapped[str] session_name: Mapped[str] peer_name: Mapped[str] From 117a83de9639ca1f9c294854069d909457446736 Mon Sep 17 00:00:00 2001 From: Phil Date: Wed, 2 Sep 2026 13:20:50 -0400 Subject: [PATCH 2/7] review: fix migration execution + index/FK consistency for tenant_id schema - Move the shared-schema bootstrap out of the Alembic chain into a standalone script (scripts/bootstrap_shared_schema.py). As a chained revision it collided with the per-tenant history on `alembic upgrade head` (DuplicateTable) and would have run on existing single-tenant instances via init_db(). It is now run explicitly against the fresh shared DB; the migration track owns provisioning. - Run the bootstrap in AUTOCOMMIT so 1000+ partition creations plus composite FKs don't exhaust max_locks_per_transaction in a single transaction. - Derive the partitioned-table set from the models instead of hand-listing it (a hand list drifts and yields a partitioned parent with zero partitions). - message_embeddings: index (message_id, tenant_id) since message_id lookups are cross-tenant; restore the composite workspaces FK for parity with peer tables. - messages: ix_messages_peer_lookup = (tenant_id, workspace_name, peer_name, created_at) to serve the peer-history query without a sort. - QueueItem.__repr__: include tenant_id. Co-Authored-By: Claude Opus 4.8 --- ...dd_tenant_id_first_class_primitive_and_.py | 104 ---------------- scripts/bootstrap_shared_schema.py | 114 ++++++++++++++++++ src/models.py | 21 +++- 3 files changed, 132 insertions(+), 107 deletions(-) delete mode 100644 migrations/versions/09dfcd158c6c_add_tenant_id_first_class_primitive_and_.py create mode 100644 scripts/bootstrap_shared_schema.py 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 From ec16f4ad25d308a63f764d82e0143b6939852f14 Mon Sep 17 00:00:00 2001 From: Phil Date: Wed, 2 Sep 2026 13:42:45 -0400 Subject: [PATCH 3/7] test: add shared-tenant schema bootstrap integration test 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 partition set, 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). Runs against real Postgres + pgvector, bypassing the app's create_all fixture (which does not create partitions). Co-Authored-By: Claude Opus 4.8 --- tests/scripts/test_bootstrap_shared_schema.py | 129 ++++++++++++++++++ 1 file changed, 129 insertions(+) create mode 100644 tests/scripts/test_bootstrap_shared_schema.py diff --git a/tests/scripts/test_bootstrap_shared_schema.py b/tests/scripts/test_bootstrap_shared_schema.py new file mode 100644 index 00000000..41b7bc29 --- /dev/null +++ b/tests/scripts/test_bootstrap_shared_schema.py @@ -0,0 +1,129 @@ +"""Integration test for the shared-tenant schema bootstrap. + +Runs ``scripts/bootstrap_shared_schema`` against a throwaway database and asserts +the partitioned schema it produces: every tenant-scoped table is +HASH(tenant_id)-partitioned with its full set of partitions, the primary key +leads with ``tenant_id``, and the tenant-scoped composite keys enforce (two +tenants can share a workspace name, a duplicate name within a tenant is rejected, +and an unknown ``tenant_id`` is rejected by the FK to ``tenants``). + +Heavier than a metadata-only shape check on purpose: it exercises the exact DDL +the migration track will run, against real Postgres + pgvector, and it does not +go through the app's ``create_all`` fixture (which does not create partitions). +""" + +import pytest +from sqlalchemy import Engine, create_engine, make_url, text +from sqlalchemy.exc import IntegrityError +from sqlalchemy.orm import Session + +from scripts.bootstrap_shared_schema import ( + PARTITION_COUNT, + bootstrap_shared_schema, + partitioned_tables, +) +from src.config import settings +from src.models import Tenant, Workspace + + +@pytest.fixture +def bootstrapped_engine(worker_id: str): + """A throwaway database with the shared-tenant schema bootstrapped into it.""" + db_name = f"test_bootstrap_shared_{worker_id}" + admin_url = make_url(settings.DB.CONNECTION_URI).set(database="postgres") + + def _drop_create(create: bool) -> None: + # Pass the URL object (not str) so the password is not masked to '***'. + admin = create_engine(admin_url, isolation_level="AUTOCOMMIT") + with admin.connect() as conn: + conn.execute(text(f'DROP DATABASE IF EXISTS "{db_name}" WITH (FORCE)')) + if create: + conn.execute(text(f'CREATE DATABASE "{db_name}"')) + admin.dispose() + + _drop_create(create=True) + engine = create_engine(make_url(settings.DB.CONNECTION_URI).set(database=db_name)) + # AUTOCOMMIT: the partition DDL accumulates too many locks for one transaction. + with engine.connect().execution_options(isolation_level="AUTOCOMMIT") as conn: + conn.execute(text("CREATE EXTENSION IF NOT EXISTS vector")) + bootstrap_shared_schema(conn) + try: + yield engine + finally: + engine.dispose() + _drop_create(create=False) + + +def test_every_tenant_table_is_partitioned(bootstrapped_engine: Engine) -> None: + """Each tenant-scoped table is a HASH partitioned table carrying all its + partitions — guards the drift trap where a partitioned parent with zero + partitions rejects every insert.""" + expected = partitioned_tables() + assert len(expected) == 9 + with bootstrapped_engine.connect() as conn: + for name in expected: + relkind = conn.execute( + text( + "SELECT relkind FROM pg_class" + + " WHERE relname = :n AND relnamespace = 'public'::regnamespace" + ), + {"n": name}, + ).scalar() + assert relkind == "p", f"{name} is not partitioned (relkind={relkind!r})" + n_partitions = conn.execute( + text( + "SELECT count(*) FROM pg_inherits i" + + " JOIN pg_class p ON p.oid = i.inhparent WHERE p.relname = :n" + ), + {"n": name}, + ).scalar() + assert n_partitions == PARTITION_COUNT, ( + f"{name} has {n_partitions} partitions, expected {PARTITION_COUNT}" + ) + + +def test_primary_key_leads_with_tenant_id(bootstrapped_engine: Engine) -> None: + with bootstrapped_engine.connect() as conn: + pk_cols = ( + conn.execute( + text( + "SELECT a.attname FROM pg_index i" + + " JOIN pg_attribute a" + + " ON a.attrelid = i.indrelid AND a.attnum = ANY(i.indkey)" + + " WHERE i.indrelid = 'public.workspaces'::regclass AND i.indisprimary" + + " ORDER BY array_position(i.indkey, a.attnum)" + ) + ) + .scalars() + .all() + ) + assert pk_cols == ["tenant_id", "id"] + + +def test_tenant_scoped_uniqueness_and_fk(bootstrapped_engine: Engine) -> None: + """The tenant-scoped composite keys enforce as designed.""" + with Session(bootstrapped_engine) as session: + session.add_all([Tenant(tenant_id="tenant_a"), Tenant(tenant_id="tenant_b")]) + session.commit() + + # Two different tenants may each own a workspace named "default": this is + # the whole point of UNIQUE(tenant_id, name) replacing UNIQUE(name). + session.add_all( + [ + Workspace(tenant_id="tenant_a", name="default"), + Workspace(tenant_id="tenant_b", name="default"), + ] + ) + session.commit() + + # Same (tenant, name) twice is rejected. + with Session(bootstrapped_engine) as session: + session.add(Workspace(tenant_id="tenant_a", name="default")) + with pytest.raises(IntegrityError): + session.commit() + + # An unknown tenant_id is rejected by the FK to tenants. + with Session(bootstrapped_engine) as session: + session.add(Workspace(tenant_id="ghost_tenant", name="scratch")) + with pytest.raises(IntegrityError): + session.commit() From c630993680819dca11fced15b25e41f6bc6aa7f7 Mon Sep 17 00:00:00 2001 From: Phil Date: Wed, 2 Sep 2026 14:11:01 -0400 Subject: [PATCH 4/7] refactor: rename tenants.app_name to legacy_app_name; apply agent-comment markers - Rename Tenant.app_name -> legacy_app_name so the incumbent per-tenant name never collides with a shared pool's app_name (which is a single shared value for all pooled tenants). It's the value that keeps a tenant's external vector-store namespace stable when the tenant moves onto a shared backend. - Apply the agent-comment-marker convention (# ai: / # region ai) across the tenant_id schema, the bootstrap script, and its test: the why (receipts, anti-prior gotchas) is foldable-marked; short docstrings carry the what. Comments/docstrings + one column rename only; the bootstrap integration test stays green (no behavioral change). Co-Authored-By: Claude Opus 4.8 --- scripts/bootstrap_shared_schema.py | 85 +++++++++++-------- src/models.py | 68 ++++++++++----- tests/scripts/test_bootstrap_shared_schema.py | 41 +++++---- 3 files changed, 119 insertions(+), 75 deletions(-) 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"), From b1d1537345452aa66954d55c9bca67237ef9e963 Mon Sep 17 00:00:00 2001 From: Phil Date: Wed, 2 Sep 2026 15:33:06 -0400 Subject: [PATCH 5/7] docs(models,scripts): keep comments implementation-generic for public readers Rewrite the tenant-table and schema-bootstrap comments to the generic, standalone-honcho rationale an external reader needs, dropping deployment-specific detail. Comment-only; no behavior change. Co-Authored-By: Claude Opus 4.8 --- scripts/bootstrap_shared_schema.py | 12 ++++++------ src/models.py | 23 ++++++++++------------- 2 files changed, 16 insertions(+), 19 deletions(-) diff --git a/scripts/bootstrap_shared_schema.py b/scripts/bootstrap_shared_schema.py index 2bc30a48..d38fbe3f 100644 --- a/scripts/bootstrap_shared_schema.py +++ b/scripts/bootstrap_shared_schema.py @@ -7,9 +7,8 @@ # (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. +# run explicitly against a freshly provisioned shared database. How that +# database is provisioned and version-stamped is out of scope for this script. # # The DDL is generated from the declarative models (``src.models.Base.metadata``): # the models are the single source of truth, the compiler renders ``PARTITION @@ -48,9 +47,10 @@ from sqlalchemy.schema import CreateIndex, CreateTable # noqa: E402 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. +# HASH(tenant_id) partition count. This is a maintenance/pruning choice — keeping +# each partition small enough to vacuum and prune independently — not a contention +# one, since the hash key already spreads writes across partitions regardless of +# count. # endregion PARTITION_COUNT = 128 diff --git a/src/models.py b/src/models.py index 78535c5b..31171759 100644 --- a/src/models.py +++ b/src/models.py @@ -111,19 +111,16 @@ class Tenant(Base): """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. + # This table can be a local mirror rather than the source of truth: when an + # external layer owns the tenant registry, rows are populated out-of-band + # (backfilled for existing tenants; written on provision, or upserted on first + # authenticated request) rather than via an in-database FK. It is also the + # one-row-per-tenant home for facts with nowhere else to live: + # - legacy_app_name: the tenant's original app name, preserved so its + # external vector-store namespace stays stable if the active app name + # changes (avoids a full re-embed). Named "legacy_" so it never competes + # with the current app_name. + # - tier: which backend deployment class serves this tenant. # endregion __tablename__: str = "tenants" tenant_id: Mapped[str] = mapped_column(TEXT, primary_key=True) From 418e59ca2e300366ac4cb5f2b79a96e35fceb32d Mon Sep 17 00:00:00 2001 From: Phil Date: Thu, 3 Sep 2026 11:54:23 -0400 Subject: [PATCH 6/7] feat(migrations): add tenant_id primitive migration (OSS prosumer-safe) Add the alembic migration that makes tenant_id a first-class primitive: a new tenants table + tenant_id with composite PKs / uniques / FKs / indexes on every tenant-scoped table, matching the declarative models MINUS physical partitioning. Self-host/prosumer safe: transforms the existing single-tenant schema in place and backfills every row to a default tenant. A top-of-upgrade guard (tenants-exists) makes it a no-op on the shared/prod schema, which the internal bootstrap builds and alembic-stamps past. Rename tenants.legacy_app_name -> vector_correlation_id (impl-agnostic: the durable external vector-store namespace key, not legacy). Remove scripts/bootstrap_shared_schema.py + its test from OSS -- the prod-only shared-schema standup moves to the internal migration runbook. Strip internal / migration-transient comments from models.py per the OSS-safe pass. Verified: pytest tests/alembic -k e5fe7f8bcf62 passes on the full revision chain. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../e5fe7f8bcf62_add_tenant_id_primitive.py | 387 ++++++++++++++++++ scripts/bootstrap_shared_schema.py | 129 ------ src/models.py | 26 +- tests/alembic/revisions/__init__.py | 14 +- ...st_e5fe7f8bcf62_add_tenant_id_primitive.py | 127 ++++++ tests/scripts/test_bootstrap_shared_schema.py | 136 ------ 6 files changed, 533 insertions(+), 286 deletions(-) create mode 100644 migrations/versions/e5fe7f8bcf62_add_tenant_id_primitive.py delete mode 100644 scripts/bootstrap_shared_schema.py create mode 100644 tests/alembic/revisions/test_e5fe7f8bcf62_add_tenant_id_primitive.py delete mode 100644 tests/scripts/test_bootstrap_shared_schema.py diff --git a/migrations/versions/e5fe7f8bcf62_add_tenant_id_primitive.py b/migrations/versions/e5fe7f8bcf62_add_tenant_id_primitive.py new file mode 100644 index 00000000..2b58e2d1 --- /dev/null +++ b/migrations/versions/e5fe7f8bcf62_add_tenant_id_primitive.py @@ -0,0 +1,387 @@ +"""add tenant id primitive + +Revision ID: e5fe7f8bcf62 +Revises: e4eba9cfaa6f +Create Date: 2026-09-03 + +Adds ``tenant_id`` as a first-class primitive to the data model: a new +``tenants`` table and a ``tenant_id`` column (plus tenant-scoped composite +PKs / uniques / FKs / indexes) on every tenant-scoped table, matching the +declarative models MINUS physical partitioning. + +Scope: this is the OSS / self-host (prosumer) migration. It transforms an +existing single-tenant, non-partitioned schema in place and backfills every +row to a single default tenant. It intentionally does NOT create HASH +partitions — ``postgresql_partition_by`` in the models is a create-time hint +that only the (internal, prod-only) shared-schema bootstrap honours; +partitioning is transparent at query time, so the declarative models run +correctly against these plain tables. + +Prod is NOT migrated by this script: the shared partitioned schema is built by +the bootstrap and prod is ``alembic stamp``-ed past this revision. The guard at +the top of ``upgrade()`` also makes this a no-op wherever the tenant schema +already exists (prod, or a re-run), so it is idempotent. +""" + +from collections.abc import Sequence +from typing import Any + +import sqlalchemy as sa +from alembic import op + +from migrations.utils import ( + column_exists, + get_schema, + index_exists, + table_exists, +) + +# revision identifiers, used by Alembic. +revision: str = "e5fe7f8bcf62" +down_revision: str | None = "e4eba9cfaa6f" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + +schema = get_schema() + +# The single tenant every existing row is backfilled to on a self-host install. +DEFAULT_TENANT_ID = "default" + +# Tables that gain tenant_id + a composite PK + composite FKs (session_peers has +# an all-natural composite PK and is handled with the rest). +TENANT_SCOPED: tuple[str, ...] = ( + "workspaces", + "peers", + "sessions", + "messages", + "message_embeddings", + "collections", + "documents", + "webhook_endpoints", + "session_peers", +) + +# New primary keys (table -> ordered PK columns). Everything is (tenant_id, id) +# except the association table, whose PK is all-natural. +NEW_PKS: dict[str, list[str]] = { + "workspaces": ["tenant_id", "id"], + "peers": ["tenant_id", "id"], + "sessions": ["tenant_id", "id"], + "messages": ["tenant_id", "id"], + "message_embeddings": ["tenant_id", "id"], + "collections": ["tenant_id", "id"], + "documents": ["tenant_id", "id"], + "webhook_endpoints": ["tenant_id", "id"], + "session_peers": ["tenant_id", "workspace_name", "session_name", "peer_name"], +} + +# New unique constraints (table -> list of (name, cols)). +NEW_UNIQUES: dict[str, list[tuple[str, list[str]]]] = { + "workspaces": [("uq_workspaces_tenant_id_name", ["tenant_id", "name"])], + "peers": [ + ("uq_peers_tenant_id_name_workspace_name", ["tenant_id", "name", "workspace_name"]) + ], + "sessions": [ + ("uq_sessions_tenant_id_name_workspace_name", ["tenant_id", "name", "workspace_name"]) + ], + "messages": [ + ("uq_messages_tenant_id_public_id", ["tenant_id", "public_id"]), + ( + "uq_messages_tenant_id_ws_session_seq", + ["tenant_id", "workspace_name", "session_name", "seq_in_session"], + ), + ], + "collections": [ + ( + "uq_collections_tenant_id_observer_observed_ws", + ["tenant_id", "observer", "observed", "workspace_name"], + ) + ], +} + +# New foreign keys: (name, source_table, [local_cols], ref_table, [ref_cols], ondelete) +# Every tenant-scoped table also FKs its tenant_id -> tenants.tenant_id (added in a +# loop below); the composite FKs below carry the tenant_id into the natural-key refs. +NEW_FKS: list[tuple[str, str, list[str], str, list[str], str | None]] = [ + # peers + ("fk_peers_ws_tenant_workspaces", "peers", ["workspace_name", "tenant_id"], "workspaces", ["name", "tenant_id"], None), + # sessions + ("fk_sessions_ws_tenant_workspaces", "sessions", ["workspace_name", "tenant_id"], "workspaces", ["name", "tenant_id"], None), + # messages + ("fk_messages_session_ws_tenant_sessions", "messages", ["session_name", "workspace_name", "tenant_id"], "sessions", ["name", "workspace_name", "tenant_id"], None), + ("fk_messages_peer_ws_tenant_peers", "messages", ["peer_name", "workspace_name", "tenant_id"], "peers", ["name", "workspace_name", "tenant_id"], None), + # message_embeddings + ("fk_msg_emb_tenant_message_messages", "message_embeddings", ["tenant_id", "message_id"], "messages", ["tenant_id", "public_id"], "CASCADE"), + ("fk_msg_emb_ws_tenant_workspaces", "message_embeddings", ["workspace_name", "tenant_id"], "workspaces", ["name", "tenant_id"], None), + ("fk_msg_emb_session_ws_tenant_sessions", "message_embeddings", ["session_name", "workspace_name", "tenant_id"], "sessions", ["name", "workspace_name", "tenant_id"], None), + ("fk_msg_emb_peer_ws_tenant_peers", "message_embeddings", ["peer_name", "workspace_name", "tenant_id"], "peers", ["name", "workspace_name", "tenant_id"], None), + # collections + ("fk_collections_ws_tenant_workspaces", "collections", ["workspace_name", "tenant_id"], "workspaces", ["name", "tenant_id"], None), + ("fk_collections_observer_ws_tenant_peers", "collections", ["observer", "workspace_name", "tenant_id"], "peers", ["name", "workspace_name", "tenant_id"], None), + ("fk_collections_observed_ws_tenant_peers", "collections", ["observed", "workspace_name", "tenant_id"], "peers", ["name", "workspace_name", "tenant_id"], None), + # documents + ("fk_documents_ws_tenant_workspaces", "documents", ["workspace_name", "tenant_id"], "workspaces", ["name", "tenant_id"], None), + ("fk_documents_collection_tenant_collections", "documents", ["observer", "observed", "workspace_name", "tenant_id"], "collections", ["observer", "observed", "workspace_name", "tenant_id"], None), + ("fk_documents_observer_ws_tenant_peers", "documents", ["observer", "workspace_name", "tenant_id"], "peers", ["name", "workspace_name", "tenant_id"], None), + ("fk_documents_observed_ws_tenant_peers", "documents", ["observed", "workspace_name", "tenant_id"], "peers", ["name", "workspace_name", "tenant_id"], None), + ("fk_documents_session_ws_tenant_sessions", "documents", ["session_name", "workspace_name", "tenant_id"], "sessions", ["name", "workspace_name", "tenant_id"], None), + # webhook_endpoints + ("fk_webhook_ws_tenant_workspaces", "webhook_endpoints", ["workspace_name", "tenant_id"], "workspaces", ["name", "tenant_id"], None), + # session_peers + ("fk_session_peers_ws_tenant_workspaces", "session_peers", ["workspace_name", "tenant_id"], "workspaces", ["name", "tenant_id"], None), + ("fk_session_peers_session_ws_tenant_sessions", "session_peers", ["session_name", "workspace_name", "tenant_id"], "sessions", ["name", "workspace_name", "tenant_id"], None), + ("fk_session_peers_peer_ws_tenant_peers", "session_peers", ["peer_name", "workspace_name", "tenant_id"], "peers", ["name", "workspace_name", "tenant_id"], None), +] + + +def _inspector() -> sa.Inspector: + return sa.inspect(op.get_bind()) + + +def _drop_all_fks(tables: Sequence[str]) -> None: + """Drop every FK on the given tables (dynamic — old names are unknown).""" + insp = _inspector() + for tname in tables: + if not table_exists(tname): + continue + for fk in insp.get_foreign_keys(tname, schema=schema): + name = fk.get("name") + if name: + op.drop_constraint(name, tname, type_="foreignkey", schema=schema) + + +def _drop_pk(table: str) -> None: + insp = _inspector() + pk = insp.get_pk_constraint(table, schema=schema) + name = pk.get("name") if pk else None + if name: + op.drop_constraint(name, table, type_="primary", schema=schema) + + +def _drop_all_uniques(table: str) -> None: + """Drop unique constraints and pure unique indexes (all become tenant-scoped).""" + insp = _inspector() + for uq in insp.get_unique_constraints(table, schema=schema): + name = uq.get("name") + if name: + op.drop_constraint(name, table, type_="unique", schema=schema) + for idx in insp.get_indexes(table, schema=schema): + idx_name = idx.get("name") + if idx.get("unique") and idx_name: + op.drop_index(idx_name, table_name=table, schema=schema) + + +def upgrade() -> None: + # region ai + # Disable-on-prod + idempotency guard. Prod's shared schema is built by the + # bootstrap (which creates `tenants`) and prod is `alembic stamp`-ed past this + # revision, so it never runs the body; this also makes the migration a no-op on + # any DB that already has the tenant schema (re-runs). Self-host / prosumer DBs + # have no `tenants` table at this point, so they run the full transform. + # endregion + if table_exists("tenants"): + return + + # 1. tenants table + its index + the single self-host tenant. + op.create_table( + "tenants", + sa.Column("tenant_id", sa.TEXT(), nullable=False), + sa.Column("vector_correlation_id", sa.TEXT(), nullable=True), + sa.Column("tier", sa.TEXT(), server_default="dedicated", nullable=False), + sa.Column( + "created_at", + sa.DateTime(timezone=True), + server_default=sa.text("now()"), + nullable=False, + ), + sa.PrimaryKeyConstraint("tenant_id", name="pk_tenants"), + schema=schema, + ) + op.create_index( + "ix_tenants_vector_correlation_id", + "tenants", + ["vector_correlation_id"], + schema=schema, + ) + op.execute( + sa.text( + f'INSERT INTO "{schema}"."tenants" (tenant_id, tier) ' + + "VALUES (:tid, 'dedicated') ON CONFLICT DO NOTHING" + ).bindparams(tid=DEFAULT_TENANT_ID) + ) + + # 2. Add tenant_id everywhere + backfill + NOT NULL on the tenant-scoped set. + for tname in TENANT_SCOPED: + if not column_exists(tname, "tenant_id"): + op.add_column( + tname, sa.Column("tenant_id", sa.TEXT(), nullable=True), schema=schema + ) + op.execute( + sa.text( + f'UPDATE "{schema}"."{tname}" SET tenant_id = :tid ' + + "WHERE tenant_id IS NULL" + ).bindparams(tid=DEFAULT_TENANT_ID) + ) + op.alter_column(tname, "tenant_id", nullable=False, schema=schema) + + # 3. Drop ALL old FKs first (they block the PK/unique changes below), incl. + # queue.session_id -> sessions.id from an earlier revision. + _drop_all_fks((*TENANT_SCOPED, "queue")) + + # 4. Rebuild primary keys -> tenant-leading composite. + for tname, cols in NEW_PKS.items(): + _drop_pk(tname) + op.create_primary_key(f"pk_{tname}", tname, cols, schema=schema) + + # 5. Drop old uniques, add the new tenant-scoped ones. + for tname in TENANT_SCOPED: + _drop_all_uniques(tname) + for tname, uniques in NEW_UNIQUES.items(): + for uname, cols in uniques: + op.create_unique_constraint(uname, tname, cols, schema=schema) + + # 6. tenant_id -> tenants.tenant_id on every tenant-scoped table, then the + # composite natural-key FKs (their targets — the new uniques — now exist). + for tname in TENANT_SCOPED: + op.create_foreign_key( + f"fk_{tname}_tenant_id_tenants", + tname, + "tenants", + ["tenant_id"], + ["tenant_id"], + source_schema=schema, + referent_schema=schema, + ) + for name, src, local_cols, ref, ref_cols, ondelete in NEW_FKS: + op.create_foreign_key( + name, + src, + ref, + local_cols, + ref_cols, + ondelete=ondelete, + source_schema=schema, + referent_schema=schema, + ) + + # 7. Service tables: tenant_id is plain attribution — nullable, no FK, no RLS. + for tname in ("queue", "active_queue_sessions"): + if not column_exists(tname, "tenant_id"): + op.add_column( + tname, sa.Column("tenant_id", sa.TEXT(), nullable=True), schema=schema + ) + if not index_exists("queue", "ix_queue_tenant_id"): + op.create_index("ix_queue_tenant_id", "queue", ["tenant_id"], schema=schema) + + # 8. New tenant-scoped / vector / fts indexes (idempotent). + _create_new_indexes() + + +def _create_new_indexes() -> None: + def _mk(name: str, table: str, cols: list[Any], **kw: Any) -> None: + if not index_exists(table, name): + op.create_index(name, table, cols, schema=schema, **kw) + + _mk("ix_peers_tenant_workspace", "peers", ["tenant_id", "workspace_name"]) + _mk("ix_sessions_tenant_workspace", "sessions", ["tenant_id", "workspace_name"]) + _mk( + "ix_messages_session_lookup", + "messages", + ["tenant_id", "session_name", "id"], + postgresql_include=["created_at"], + ) + _mk( + "ix_messages_peer_lookup", + "messages", + ["tenant_id", "workspace_name", "peer_name", "created_at"], + ) + _mk( + "ix_messages_content_gin", + "messages", + [sa.text("to_tsvector('english', content)")], + postgresql_using="gin", + ) + _mk( + "ix_message_embeddings_message_tenant", + "message_embeddings", + ["message_id", "tenant_id"], + ) + _mk( + "ix_message_embeddings_embedding_hnsw", + "message_embeddings", + ["embedding"], + postgresql_using="hnsw", + postgresql_with={"m": 16, "ef_construction": 64}, + postgresql_ops={"embedding": "vector_cosine_ops"}, + ) + _mk( + "ix_message_embeddings_sync_state_last_sync_at", + "message_embeddings", + ["sync_state", "last_sync_at"], + ) + _mk( + "ix_documents_tenant_collection", + "documents", + ["tenant_id", "observer", "observed", "workspace_name"], + ) + _mk( + "ix_documents_embedding_hnsw", + "documents", + ["embedding"], + postgresql_using="hnsw", + postgresql_with={"m": 16, "ef_construction": 64}, + postgresql_ops={"embedding": "vector_cosine_ops"}, + ) + _mk("ix_documents_source_ids_gin", "documents", ["source_ids"], postgresql_using="gin") + _mk( + "ix_documents_sync_state_last_sync_at", + "documents", + ["sync_state", "last_sync_at"], + ) + _mk( + "ix_webhook_endpoints_tenant_workspace", + "webhook_endpoints", + ["tenant_id", "workspace_name"], + ) + + +def downgrade() -> None: + # Reverse the transform: drop new FKs/uniques/indexes, restore sole-id PKs, + # drop tenant_id, drop tenants. Dropping the tenant_id column cascades away any + # residual constraint that references it. + if not table_exists("tenants"): + return + + for name, src, *_ in NEW_FKS: + insp = _inspector() + if any(fk.get("name") == name for fk in insp.get_foreign_keys(src, schema=schema)): + op.drop_constraint(name, src, type_="foreignkey", schema=schema) + for tname in TENANT_SCOPED: + insp = _inspector() + fk_name = f"fk_{tname}_tenant_id_tenants" + if any(fk.get("name") == fk_name for fk in insp.get_foreign_keys(tname, schema=schema)): + op.drop_constraint(fk_name, tname, type_="foreignkey", schema=schema) + + for tname, uniques in NEW_UNIQUES.items(): + for uname, _cols in uniques: + insp = _inspector() + if any( + uq.get("name") == uname + for uq in insp.get_unique_constraints(tname, schema=schema) + ): + op.drop_constraint(uname, tname, type_="unique", schema=schema) + + # Restore sole-id / natural PKs, then drop tenant_id (with any index on it). + old_pks = { + "session_peers": ["workspace_name", "session_name", "peer_name"], + } + for tname in TENANT_SCOPED: + _drop_pk(tname) + cols = old_pks.get(tname, ["id"]) + op.create_primary_key(f"pk_{tname}", tname, cols, schema=schema) + + for tname in (*TENANT_SCOPED, "queue", "active_queue_sessions"): + if column_exists(tname, "tenant_id"): + op.drop_column(tname, "tenant_id", schema=schema) + + op.drop_table("tenants", schema=schema) diff --git a/scripts/bootstrap_shared_schema.py b/scripts/bootstrap_shared_schema.py deleted file mode 100644 index d38fbe3f..00000000 --- a/scripts/bootstrap_shared_schema.py +++ /dev/null @@ -1,129 +0,0 @@ -# honcho/scripts/bootstrap_shared_schema.py -"""Bootstrap the shared-tenant partitioned schema on a fresh shared database.""" - -# 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 -# run explicitly against a freshly provisioned shared database. How that -# database is provisioned and version-stamped is out of scope for this script. -# -# 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 - -# 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. This is a maintenance/pruning choice — keeping -# each partition small enough to vacuum and prune independently — not a contention -# one, since the hash key already spreads writes across partitions regardless of -# count. -# endregion -PARTITION_COUNT = 128 - - -def partitioned_tables() -> set[str]: - """Names of the tables declared with HASH(tenant_id) partitioning.""" - - # 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() - 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``.""" - - # 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( - 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 31171759..b9d86294 100644 --- a/src/models.py +++ b/src/models.py @@ -111,20 +111,17 @@ class Tenant(Base): """One row per tenant; the FK target for every tenant-scoped table's ``tenant_id``.""" # region ai - # This table can be a local mirror rather than the source of truth: when an - # external layer owns the tenant registry, rows are populated out-of-band - # (backfilled for existing tenants; written on provision, or upserted on first - # authenticated request) rather than via an in-database FK. It is also the - # one-row-per-tenant home for facts with nowhere else to live: - # - legacy_app_name: the tenant's original app name, preserved so its - # external vector-store namespace stays stable if the active app name - # changes (avoids a full re-embed). Named "legacy_" so it never competes - # with the current app_name. + # Home for per-tenant facts with nowhere else to live: + # - vector_correlation_id: the tenant's external vector-store namespace key, + # preserved so the namespace stays stable if the tenant's app name changes + # (existing vectors resolve without a full, paid re-embed). # - tier: which backend deployment class serves this tenant. # endregion __tablename__: str = "tenants" tenant_id: Mapped[str] = mapped_column(TEXT, primary_key=True) - legacy_app_name: Mapped[str | None] = mapped_column(TEXT, nullable=True, index=True) + vector_correlation_id: 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() @@ -624,11 +621,10 @@ class QueueItem(Base): 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. + # Service table (unpartitioned), so it keeps a sole-id PK. tenant_id is a plain + # attribution / fair-scheduling column — no FK, no RLS — and the queue carries no + # FKs to sessions / messages / workspaces: 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) diff --git a/tests/alembic/revisions/__init__.py b/tests/alembic/revisions/__init__.py index 77c15483..e7847741 100644 --- a/tests/alembic/revisions/__init__.py +++ b/tests/alembic/revisions/__init__.py @@ -4,16 +4,16 @@ from . import ( test_05486ce795d5_make_session_name_required_on_messages, test_066e87ca5b07_align_schema_with_declarative_models, test_08894082221a_replace_collection_name_with_observer_, - test_7c0d9a4e3b1f_add_unique_index_for_pending_dreams, - test_20f89a421aff_rename_metamessage_type_to_label, - test_29ade7350c19_remove_document_level_valid_constraint, - test_66e63cf2cf77_add_indexes_to_documents_table, - test_76ffba56fe8c_add_error_field_to_queueitem, - test_88b0fb10906f_add_webhooks_table, test_110bdf470272_rename_deriver_disabled_to_deriver_, test_119a52b73c60_support_external_embeddings, + test_20f89a421aff_rename_metamessage_type_to_label, + test_29ade7350c19_remove_document_level_valid_constraint, test_556a16564f50_add_user_id_and_app_id_to_tables, test_564ba40505c5_add_session_name_column_to_documents, + test_66e63cf2cf77_add_indexes_to_documents_table, + test_76ffba56fe8c_add_error_field_to_queueitem, + test_7c0d9a4e3b1f_add_unique_index_for_pending_dreams, + test_88b0fb10906f_add_webhooks_table, test_917195d9b5e9_add_messageembedding_table, test_a1b2c3d4e5f6_initial_schema, test_b765d82110bd_change_metamessages_to_user_level_with_, @@ -23,6 +23,7 @@ from . import ( test_c3828084f472_add_indexes_for_messages_and_, test_d429de0e5338_adopt_peer_paradigm, test_e4eba9cfaa6f_make_document_session_name_nullable, + test_e5fe7f8bcf62_add_tenant_id_primitive, test_e9b705f9adf9_add_server_defaults_to_timestamp_, test_ec8f94139b02_codify_workspace_name_and_message_id_in_, test_f1a2b3c4d5e6_add_reasoning_tree_columns, @@ -51,6 +52,7 @@ __all__ = [ "test_c3828084f472_add_indexes_for_messages_and_", "test_d429de0e5338_adopt_peer_paradigm", "test_e4eba9cfaa6f_make_document_session_name_nullable", + "test_e5fe7f8bcf62_add_tenant_id_primitive", "test_e9b705f9adf9_add_server_defaults_to_timestamp_", "test_ec8f94139b02_codify_workspace_name_and_message_id_in_", "test_f1a2b3c4d5e6_add_reasoning_tree_columns", diff --git a/tests/alembic/revisions/test_e5fe7f8bcf62_add_tenant_id_primitive.py b/tests/alembic/revisions/test_e5fe7f8bcf62_add_tenant_id_primitive.py new file mode 100644 index 00000000..2445ccca --- /dev/null +++ b/tests/alembic/revisions/test_e5fe7f8bcf62_add_tenant_id_primitive.py @@ -0,0 +1,127 @@ +"""Hooks for revision e5fe7f8bcf62 (add_tenant_id_primitive).""" + +from __future__ import annotations + +from nanoid import generate as generate_nanoid +from sqlalchemy import text + +from tests.alembic.registry import register_after_upgrade, register_before_upgrade +from tests.alembic.verifier import MigrationVerifier + +WORKSPACE_NAME = "w1" +PEER_NAME = "p1" +SESSION_NAME = "s1" + +# Tenant-scoped tables that must end up with a NOT NULL tenant_id. +_TENANT_SCOPED = ( + "workspaces", + "peers", + "sessions", + "messages", + "message_embeddings", + "collections", + "documents", + "webhook_endpoints", + "session_peers", +) + + +@register_before_upgrade("e5fe7f8bcf62") +def prepare_add_tenant_id_primitive(verifier: MigrationVerifier) -> None: + """Seed old-schema rows before upgrading to e5fe7f8bcf62.""" + # Pre-state: no tenants table, no tenant_id yet. + verifier.assert_table_exists("tenants", exists=False) + verifier.assert_column_exists("workspaces", "tenant_id", exists=False) + verifier.assert_column_exists("messages", "tenant_id", exists=False) + + conn = verifier.conn + schema = verifier.schema + + conn.execute( + text( + f'INSERT INTO "{schema}"."workspaces" ("id", "name") VALUES (:id, :name)' + ), + {"id": generate_nanoid(), "name": WORKSPACE_NAME}, + ) + conn.execute( + text( + f'INSERT INTO "{schema}"."peers" ("id", "name", "workspace_name") ' + + "VALUES (:id, :name, :ws)" + ), + {"id": generate_nanoid(), "name": PEER_NAME, "ws": WORKSPACE_NAME}, + ) + conn.execute( + text( + f'INSERT INTO "{schema}"."sessions" ("id", "name", "workspace_name") ' + + "VALUES (:id, :name, :ws)" + ), + {"id": generate_nanoid(), "name": SESSION_NAME, "ws": WORKSPACE_NAME}, + ) + conn.execute( + text( + f'INSERT INTO "{schema}"."messages" ' + + '("public_id", "workspace_name", "session_name", "peer_name", "content", "seq_in_session") ' + + "VALUES (:pid, :ws, :sn, :pn, :content, :seq)" + ), + { + "pid": generate_nanoid(), + "ws": WORKSPACE_NAME, + "sn": SESSION_NAME, + "pn": PEER_NAME, + "content": "hello", + "seq": 0, + }, + ) + + +@register_after_upgrade("e5fe7f8bcf62") +def verify_add_tenant_id_primitive(verifier: MigrationVerifier) -> None: + """Assert the tenant_id primitive landed.""" + schema = verifier.schema + conn = verifier.conn + + # tenants table + the renamed correlation column + the default tenant row. + verifier.assert_table_exists("tenants") + verifier.assert_column_exists("tenants", "vector_correlation_id") + verifier.assert_column_exists("tenants", "legacy_app_name", exists=False) + default_count = conn.execute( + text(f"SELECT COUNT(*) FROM \"{schema}\".\"tenants\" WHERE tenant_id = 'default'") + ).scalar() + assert default_count == 1, f"expected the default tenant, found {default_count}" + + # tenant_id is NOT NULL on every tenant-scoped table, and existing rows backfilled. + for table in _TENANT_SCOPED: + verifier.assert_column_exists(table, "tenant_id", nullable=False) + for table in ("workspaces", "peers", "sessions", "messages"): + verifier.assert_no_nulls(table, "tenant_id") + backfilled = conn.execute( + text( + f'SELECT COUNT(*) FROM "{schema}"."{table}" ' + + "WHERE tenant_id <> 'default'" + ) + ).scalar() + assert backfilled == 0, f"{table} has rows not backfilled to the default tenant" + + # Composite PK (tenant_id, id) on workspaces. + pk = verifier.get_inspector().get_pk_constraint("workspaces", schema=schema) + assert pk["constrained_columns"] == [ + "tenant_id", + "id", + ], f"workspaces PK is {pk['constrained_columns']}" + + # Tenant-scoped uniqueness replaced the global one. + verifier.assert_constraint_exists( + "workspaces", "uq_workspaces_tenant_id_name", "unique" + ) + + # tenant_id FK to tenants + a representative composite FK. + verifier.assert_constraint_exists( + "peers", "fk_peers_tenant_id_tenants", "foreign_key" + ) + verifier.assert_constraint_exists( + "peers", "fk_peers_ws_tenant_workspaces", "foreign_key" + ) + + # Service tables carry tenant_id but leave it nullable (attribution only). + verifier.assert_column_exists("queue", "tenant_id", nullable=True) + verifier.assert_column_exists("active_queue_sessions", "tenant_id", nullable=True) diff --git a/tests/scripts/test_bootstrap_shared_schema.py b/tests/scripts/test_bootstrap_shared_schema.py deleted file mode 100644 index 60b8a921..00000000 --- a/tests/scripts/test_bootstrap_shared_schema.py +++ /dev/null @@ -1,136 +0,0 @@ -"""Integration test for the shared-tenant schema bootstrap.""" - -# 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 -from sqlalchemy.exc import IntegrityError -from sqlalchemy.orm import Session - -from scripts.bootstrap_shared_schema import ( - PARTITION_COUNT, - bootstrap_shared_schema, - partitioned_tables, -) -from src.config import settings -from src.models import Tenant, Workspace - - -@pytest.fixture -def bootstrapped_engine(worker_id: str): - """A throwaway database with the shared-tenant schema bootstrapped into it.""" - db_name = f"test_bootstrap_shared_{worker_id}" - admin_url = make_url(settings.DB.CONNECTION_URI).set(database="postgres") - - def _drop_create(create: bool) -> None: - # 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)')) - if create: - conn.execute(text(f'CREATE DATABASE "{db_name}"')) - admin.dispose() - - _drop_create(create=True) - engine = create_engine(make_url(settings.DB.CONNECTION_URI).set(database=db_name)) - # 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) - try: - yield engine - finally: - engine.dispose() - _drop_create(create=False) - - -def test_every_tenant_table_is_partitioned(bootstrapped_engine: Engine) -> None: - """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: - for name in expected: - relkind = conn.execute( - text( - "SELECT relkind FROM pg_class" - + " WHERE relname = :n AND relnamespace = 'public'::regnamespace" - ), - {"n": name}, - ).scalar() - assert relkind == "p", f"{name} is not partitioned (relkind={relkind!r})" - n_partitions = conn.execute( - text( - "SELECT count(*) FROM pg_inherits i" - + " JOIN pg_class p ON p.oid = i.inhparent WHERE p.relname = :n" - ), - {"n": name}, - ).scalar() - assert n_partitions == PARTITION_COUNT, ( - f"{name} has {n_partitions} partitions, expected {PARTITION_COUNT}" - ) - - -def test_primary_key_leads_with_tenant_id(bootstrapped_engine: Engine) -> None: - with bootstrapped_engine.connect() as conn: - pk_cols = ( - conn.execute( - text( - "SELECT a.attname FROM pg_index i" - + " JOIN pg_attribute a" - + " ON a.attrelid = i.indrelid AND a.attnum = ANY(i.indkey)" - + " WHERE i.indrelid = 'public.workspaces'::regclass AND i.indisprimary" - + " ORDER BY array_position(i.indkey, a.attnum)" - ) - ) - .scalars() - .all() - ) - assert pk_cols == ["tenant_id", "id"] - - -def test_tenant_scoped_uniqueness_and_fk(bootstrapped_engine: Engine) -> None: - """The tenant-scoped composite keys enforce as designed.""" - with Session(bootstrapped_engine) as session: - session.add_all([Tenant(tenant_id="tenant_a"), Tenant(tenant_id="tenant_b")]) - session.commit() - - # 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"), - Workspace(tenant_id="tenant_b", name="default"), - ] - ) - session.commit() - - # Same (tenant, name) twice is rejected. - with Session(bootstrapped_engine) as session: - session.add(Workspace(tenant_id="tenant_a", name="default")) - with pytest.raises(IntegrityError): - session.commit() - - # An unknown tenant_id is rejected by the FK to tenants. - with Session(bootstrapped_engine) as session: - session.add(Workspace(tenant_id="ghost_tenant", name="scratch")) - with pytest.raises(IntegrityError): - session.commit() From 6adacd981f40a1872b6429291cb333494eadc5e9 Mon Sep 17 00:00:00 2001 From: Phil Date: Thu, 3 Sep 2026 11:59:47 -0400 Subject: [PATCH 7/7] Update e5fe7f8bcf62_add_tenant_id_primitive.py remove non-OSS facing docstring; leaving the #ai comment for now --- .../e5fe7f8bcf62_add_tenant_id_primitive.py | 13 ------------- 1 file changed, 13 deletions(-) diff --git a/migrations/versions/e5fe7f8bcf62_add_tenant_id_primitive.py b/migrations/versions/e5fe7f8bcf62_add_tenant_id_primitive.py index 2b58e2d1..e621b7ea 100644 --- a/migrations/versions/e5fe7f8bcf62_add_tenant_id_primitive.py +++ b/migrations/versions/e5fe7f8bcf62_add_tenant_id_primitive.py @@ -8,19 +8,6 @@ Adds ``tenant_id`` as a first-class primitive to the data model: a new ``tenants`` table and a ``tenant_id`` column (plus tenant-scoped composite PKs / uniques / FKs / indexes) on every tenant-scoped table, matching the declarative models MINUS physical partitioning. - -Scope: this is the OSS / self-host (prosumer) migration. It transforms an -existing single-tenant, non-partitioned schema in place and backfills every -row to a single default tenant. It intentionally does NOT create HASH -partitions — ``postgresql_partition_by`` in the models is a create-time hint -that only the (internal, prod-only) shared-schema bootstrap honours; -partitioning is transparent at query time, so the declarative models run -correctly against these plain tables. - -Prod is NOT migrated by this script: the shared partitioned schema is built by -the bootstrap and prod is ``alembic stamp``-ed past this revision. The guard at -the top of ``upgrade()`` also makes this a no-op wherever the tenant schema -already exists (prod, or a re-run), so it is idempotent. """ from collections.abc import Sequence