Merge 6adacd981f into 997b4764b9
This commit is contained in:
commit
927763fc77
|
|
@ -0,0 +1,374 @@
|
|||
"""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.
|
||||
"""
|
||||
|
||||
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)
|
||||
350
src/models.py
350
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,
|
||||
# ai: 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,59 @@ 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):
|
||||
"""One row per tenant; the FK target for every tenant-scoped table's ``tenant_id``."""
|
||||
|
||||
# region ai
|
||||
# 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)
|
||||
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()
|
||||
)
|
||||
|
||||
|
||||
@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)
|
||||
# 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
|
||||
)
|
||||
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 +162,27 @@ 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"),
|
||||
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 +193,11 @@ 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
|
||||
)
|
||||
# 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")
|
||||
)
|
||||
|
|
@ -153,20 +208,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 +243,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 +255,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 +297,86 @@ 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"),
|
||||
# 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"),
|
||||
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_peer_lookup",
|
||||
"tenant_id",
|
||||
"workspace_name",
|
||||
"peer_name",
|
||||
"created_at",
|
||||
),
|
||||
UniqueConstraint(
|
||||
"tenant_id",
|
||||
"workspace_name",
|
||||
"session_name",
|
||||
"seq_in_session",
|
||||
),
|
||||
# Full text search index on content column
|
||||
# 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)"),
|
||||
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 +386,42 @@ class MessageEmbedding(Base):
|
|||
)
|
||||
|
||||
__table_args__ = (
|
||||
# Compound foreign key constraints
|
||||
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(
|
||||
["session_name", "workspace_name"],
|
||||
["sessions.name", "sessions.workspace_name"],
|
||||
["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"],
|
||||
),
|
||||
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"],
|
||||
),
|
||||
# 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",
|
||||
|
|
@ -322,12 +429,17 @@ class MessageEmbedding(Base):
|
|||
postgresql_with={"m": 16, "ef_construction": 64},
|
||||
postgresql_ops={"embedding": "vector_cosine_ops"},
|
||||
),
|
||||
# Composite index for efficient reconciliation queries
|
||||
# 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",
|
||||
"last_sync_at",
|
||||
),
|
||||
{"postgresql_partition_by": "HASH (tenant_id)"},
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -335,9 +447,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 +465,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 +520,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 +542,50 @@ 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
|
||||
# ai: replaces the single observer/observed indexes
|
||||
Index(
|
||||
"ix_documents_tenant_collection",
|
||||
"tenant_id",
|
||||
"observer",
|
||||
"observed",
|
||||
"workspace_name",
|
||||
),
|
||||
# ai: HNSW is a single-column vector index (per-partition under HASH(tenant_id))
|
||||
Index(
|
||||
"ix_documents_embedding_hnsw",
|
||||
"embedding",
|
||||
|
|
@ -464,12 +601,16 @@ class Document(Base):
|
|||
"source_ids",
|
||||
postgresql_using="gin",
|
||||
),
|
||||
# Composite index for efficient reconciliation queries
|
||||
# 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",
|
||||
"last_sync_at",
|
||||
),
|
||||
{"postgresql_partition_by": "HASH (tenant_id)"},
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -479,9 +620,14 @@ 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
|
||||
)
|
||||
# region ai
|
||||
# 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)
|
||||
work_unit_key: Mapped[str] = mapped_column(TEXT, nullable=False)
|
||||
|
||||
task_type: Mapped[TaskType] = mapped_column(TEXT, nullable=False)
|
||||
|
|
@ -493,12 +639,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(
|
||||
|
|
@ -529,7 +671,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
|
||||
|
|
@ -538,6 +680,9 @@ class ActiveQueueSession(Base):
|
|||
|
||||
id: Mapped[str] = mapped_column(TEXT, default=generate_nanoid, primary_key=True)
|
||||
|
||||
# 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)
|
||||
|
||||
last_updated: Mapped[datetime.datetime] = mapped_column(
|
||||
|
|
@ -548,10 +693,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 +705,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 +725,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]
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
Loading…
Reference in New Issue