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 <noreply@anthropic.com>
This commit is contained in:
Phil 2026-09-02 14:11:01 -04:00
parent ec16f4ad25
commit c630993680
3 changed files with 119 additions and 75 deletions

View File

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

View File

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

View File

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