RomanNum3ral
/
honcho
mirror of https://github.com/plastic-labs/honcho.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

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 <noreply@anthropic.com>
This commit is contained in:
Phil 2026-09-02 13:20:50 -04:00
parent 70c4584e7c
commit 117a83de96
3 changed files with 132 additions and 107 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

104
migrations/versions/09dfcd158c6c_add_tenant_id_first_class_primitive_and_.py
View File

@ -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")

114
scripts/bootstrap_shared_schema.py Normal file
View File

@ -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.")

21
src/models.py
View File

@ -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