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

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) <noreply@anthropic.com>
This commit is contained in:
Phil 2026-09-03 11:54:23 -04:00
parent b1d1537345
commit 418e59ca2e
6 changed files with 533 additions and 286 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

387
migrations/versions/e5fe7f8bcf62_add_tenant_id_primitive.py Normal file
View File

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

129
scripts/bootstrap_shared_schema.py
View File

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

26
src/models.py
View File

@ -111,20 +111,17 @@ class Tenant(Base):
"""One row per tenant; the FK target for every tenant-scoped table's ``tenant_id``.""" """One row per tenant; the FK target for every tenant-scoped table's ``tenant_id``."""
# region ai # region ai
# This table can be a local mirror rather than the source of truth: when an # Home for per-tenant facts with nowhere else to live:
# external layer owns the tenant registry, rows are populated out-of-band # - vector_correlation_id: the tenant's external vector-store namespace key,
# (backfilled for existing tenants; written on provision, or upserted on first # preserved so the namespace stays stable if the tenant's app name changes
# authenticated request) rather than via an in-database FK. It is also the # (existing vectors resolve without a full, paid re-embed).
# 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. # - tier: which backend deployment class serves this tenant.
# endregion # endregion
__tablename__: str = "tenants" __tablename__: str = "tenants"
tenant_id: Mapped[str] = mapped_column(TEXT, primary_key=True) 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") tier: Mapped[str] = mapped_column(TEXT, nullable=False, server_default="dedicated")
created_at: Mapped[datetime.datetime] = mapped_column( created_at: Mapped[datetime.datetime] = mapped_column(
DateTime(timezone=True), server_default=func.now() DateTime(timezone=True), server_default=func.now()
@ -624,11 +621,10 @@ class QueueItem(Base):
BigInteger, Identity(), primary_key=True, autoincrement=True BigInteger, Identity(), primary_key=True, autoincrement=True
) )
# region ai # region ai
# Service table: NOT partitioned and drained-not-copied at migration, so it # Service table (unpartitioned), so it keeps a sole-id PK. tenant_id is a plain
# keeps a sole-id PK. tenant_id is a plain attribution / fair-scheduling column # attribution / fair-scheduling column — no FK, no RLS — and the queue carries no
# (no FK, no RLS); the FKs to the now-partitioned sessions / messages / # FKs to sessions / messages / workspaces: the app manages queue lifecycle and
# workspaces are dropped — the app manages queue lifecycle and already # already tolerates missing referents.
# tolerates missing referents.
# endregion # endregion
tenant_id: Mapped[str | None] = mapped_column(TEXT, nullable=True, index=True) tenant_id: Mapped[str | None] = mapped_column(TEXT, nullable=True, index=True)
session_id: Mapped[str | None] = mapped_column(TEXT, nullable=True, index=True) session_id: Mapped[str | None] = mapped_column(TEXT, nullable=True, index=True)

14
tests/alembic/revisions/__init__.py
View File

@ -4,16 +4,16 @@ from . import (
test_05486ce795d5_make_session_name_required_on_messages, test_05486ce795d5_make_session_name_required_on_messages,
test_066e87ca5b07_align_schema_with_declarative_models, test_066e87ca5b07_align_schema_with_declarative_models,
test_08894082221a_replace_collection_name_with_observer_, 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_110bdf470272_rename_deriver_disabled_to_deriver_,
test_119a52b73c60_support_external_embeddings, 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_556a16564f50_add_user_id_and_app_id_to_tables,
test_564ba40505c5_add_session_name_column_to_documents, 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_917195d9b5e9_add_messageembedding_table,
test_a1b2c3d4e5f6_initial_schema, test_a1b2c3d4e5f6_initial_schema,
test_b765d82110bd_change_metamessages_to_user_level_with_, test_b765d82110bd_change_metamessages_to_user_level_with_,
@ -23,6 +23,7 @@ from . import (
test_c3828084f472_add_indexes_for_messages_and_, test_c3828084f472_add_indexes_for_messages_and_,
test_d429de0e5338_adopt_peer_paradigm, test_d429de0e5338_adopt_peer_paradigm,
test_e4eba9cfaa6f_make_document_session_name_nullable, test_e4eba9cfaa6f_make_document_session_name_nullable,
test_e5fe7f8bcf62_add_tenant_id_primitive,
test_e9b705f9adf9_add_server_defaults_to_timestamp_, test_e9b705f9adf9_add_server_defaults_to_timestamp_,
test_ec8f94139b02_codify_workspace_name_and_message_id_in_, test_ec8f94139b02_codify_workspace_name_and_message_id_in_,
test_f1a2b3c4d5e6_add_reasoning_tree_columns, test_f1a2b3c4d5e6_add_reasoning_tree_columns,
@ -51,6 +52,7 @@ __all__ = [
"test_c3828084f472_add_indexes_for_messages_and_", "test_c3828084f472_add_indexes_for_messages_and_",
"test_d429de0e5338_adopt_peer_paradigm", "test_d429de0e5338_adopt_peer_paradigm",
"test_e4eba9cfaa6f_make_document_session_name_nullable", "test_e4eba9cfaa6f_make_document_session_name_nullable",
"test_e5fe7f8bcf62_add_tenant_id_primitive",
"test_e9b705f9adf9_add_server_defaults_to_timestamp_", "test_e9b705f9adf9_add_server_defaults_to_timestamp_",
"test_ec8f94139b02_codify_workspace_name_and_message_id_in_", "test_ec8f94139b02_codify_workspace_name_and_message_id_in_",
"test_f1a2b3c4d5e6_add_reasoning_tree_columns", "test_f1a2b3c4d5e6_add_reasoning_tree_columns",

127
tests/alembic/revisions/test_e5fe7f8bcf62_add_tenant_id_primitive.py Normal file
View File

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

136
tests/scripts/test_bootstrap_shared_schema.py
View File

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