feat(models): make tenant_id a first-class primitive for shared-tenant schema
Introduce a `tenants` registry table and a `tenant_id` column on every tenant-scoped table, with composite `(tenant_id, id)` primary keys, composite foreign keys (natural-key chains plus a `tenant_id` -> tenants FK), and `HASH(tenant_id)` partitioning (N=128) so many tenants can share one database with per-tenant isolation. `workspaces.name` uniqueness becomes `UNIQUE(tenant_id, name)` and indexes lead with `tenant_id`. Service tables (queue, active_queue_sessions) carry a plain `tenant_id` for attribution and stay unpartitioned; their FKs to the now-partitioned parents are dropped. Adds the bootstrap migration that creates the partitioned shared schema from the declarative models (written, not run). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Phil
parent
a026bebdef
commit
70c4584e7c
|
|
@ -0,0 +1,104 @@
|
|||
"""add tenant_id first class primitive and shared partitioned schema
|
||||
|
||||
Bootstraps the shared-tenant target schema: a ``tenants`` registry table, a
|
||||
``tenant_id`` column on every tenant-scoped table, composite ``(tenant_id, id)``
|
||||
primary keys, composite foreign keys, and ``HASH(tenant_id)`` partitioning on the
|
||||
data tables. The service tables (``queue``, ``active_queue_sessions``) carry a
|
||||
plain ``tenant_id`` and stay unpartitioned.
|
||||
|
||||
This CREATES the partitioned tables from scratch — partitioning cannot be
|
||||
introduced by ``ALTER TABLE`` — so it bootstraps a fresh shared database that
|
||||
per-tenant data is consolidated into, rather than transforming an existing
|
||||
single-tenant database in place.
|
||||
|
||||
The DDL is generated from the declarative models (``src.models.Base.metadata``)
|
||||
rather than hand-transcribed. The models are the single source of truth for the
|
||||
target schema; the compiler renders ``PARTITION BY`` from each table's
|
||||
``postgresql_partition_by`` option and applies the naming convention; and this
|
||||
keeps the bootstrap exactly in sync with the accepted data model. Only the
|
||||
per-table HASH partitions (which the model layer does not enumerate) are created
|
||||
explicitly. The target schema and ``search_path`` are established by
|
||||
``migrations/env.py`` before this runs.
|
||||
|
||||
Revision ID: 09dfcd158c6c
|
||||
Revises: e4eba9cfaa6f
|
||||
Create Date: 2026-09-02 11:38:27.736205
|
||||
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import op
|
||||
from sqlalchemy import Table
|
||||
from sqlalchemy.schema import CreateIndex, CreateTable
|
||||
|
||||
# Importing Base from src.models (rather than src.db) also registers every model
|
||||
# on Base.metadata as a side effect — that populated metadata is the whole schema
|
||||
# this migration builds.
|
||||
from src.models import Base # pyright: ignore[reportPrivateLocalImportUsage]
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = "09dfcd158c6c"
|
||||
down_revision: str | None = "e4eba9cfaa6f"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
# HASH(tenant_id) partition count for the tenant-scoped tables. Sized so the
|
||||
# largest table (a couple hundred GB fleet-wide) lands at roughly a couple GB per
|
||||
# partition — small enough to stay maintainable and well under a node's memory,
|
||||
# with headroom for growth. Contention is not the driver (write volume is tiny).
|
||||
PARTITION_COUNT = 128
|
||||
|
||||
# Tenant-scoped data tables partitioned by HASH(tenant_id). The service tables
|
||||
# (queue, active_queue_sessions) and the tenants registry are NOT partitioned.
|
||||
PARTITIONED_TABLES = frozenset(
|
||||
{
|
||||
"workspaces",
|
||||
"peers",
|
||||
"sessions",
|
||||
"messages",
|
||||
"message_embeddings",
|
||||
"collections",
|
||||
"documents",
|
||||
"session_peers",
|
||||
"webhook_endpoints",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _qualified(table: Table) -> str:
|
||||
return f'"{table.schema}"."{table.name}"' if table.schema else f'"{table.name}"'
|
||||
|
||||
|
||||
def _create_hash_partitions(table: Table) -> None:
|
||||
"""Create the N HASH partitions for a partitioned parent table."""
|
||||
schema_prefix = f'"{table.schema}".' if table.schema else ""
|
||||
for remainder in range(PARTITION_COUNT):
|
||||
op.execute(
|
||||
f'CREATE TABLE {schema_prefix}"{table.name}_p{remainder:03d}"'
|
||||
+ f" PARTITION OF {_qualified(table)}"
|
||||
+ f" FOR VALUES WITH (MODULUS {PARTITION_COUNT}, REMAINDER {remainder})"
|
||||
)
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
metadata = Base.metadata
|
||||
# Create tables in FK-dependency order (tenants first). A partitioned
|
||||
# parent's HASH partitions are created immediately after the parent.
|
||||
for table in metadata.sorted_tables:
|
||||
op.execute(CreateTable(table))
|
||||
if table.name in PARTITIONED_TABLES:
|
||||
_create_hash_partitions(table)
|
||||
# Create indexes once every partition exists, so each partitioned index
|
||||
# cascades onto all partitions.
|
||||
for table in metadata.sorted_tables:
|
||||
for index in table.indexes:
|
||||
op.execute(CreateIndex(index))
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
metadata = Base.metadata
|
||||
# Reverse dependency order; CASCADE drops each partitioned table's partitions
|
||||
# (and any dependent FKs) along with the parent.
|
||||
for table in reversed(metadata.sorted_tables):
|
||||
op.execute(f"DROP TABLE IF EXISTS {_qualified(table)} CASCADE")
|
||||
318
src/models.py
318
src/models.py
|
|
@ -1,6 +1,6 @@
|
|||
import datetime
|
||||
from logging import getLogger
|
||||
from typing import Any, final
|
||||
from typing import Any, final, override
|
||||
|
||||
from dotenv import load_dotenv
|
||||
from nanoid import generate as generate_nanoid
|
||||
|
|
@ -16,6 +16,7 @@ from sqlalchemy import (
|
|||
Identity,
|
||||
Index,
|
||||
Integer,
|
||||
PrimaryKeyConstraint,
|
||||
Table,
|
||||
UniqueConstraint,
|
||||
text,
|
||||
|
|
@ -23,7 +24,6 @@ from sqlalchemy import (
|
|||
from sqlalchemy.dialects.postgresql import JSONB, TEXT
|
||||
from sqlalchemy.orm import Mapped, MappedColumn, mapped_column, relationship
|
||||
from sqlalchemy.sql import func
|
||||
from typing_extensions import override
|
||||
|
||||
from src.config import settings
|
||||
from src.utils.types import DocumentLevel, TaskType, VectorSyncState
|
||||
|
|
@ -41,10 +41,17 @@ logger = getLogger(__name__)
|
|||
session_peers_table = Table(
|
||||
"session_peers",
|
||||
Base.metadata,
|
||||
# tenant_id leads the all-natural-key PK and is the HASH partition key.
|
||||
Column(
|
||||
"tenant_id",
|
||||
TEXT,
|
||||
ForeignKey("tenants.tenant_id"),
|
||||
primary_key=True,
|
||||
nullable=False,
|
||||
),
|
||||
Column(
|
||||
"workspace_name",
|
||||
TEXT,
|
||||
ForeignKey("workspaces.name"),
|
||||
primary_key=True,
|
||||
nullable=False,
|
||||
),
|
||||
|
|
@ -80,24 +87,64 @@ session_peers_table = Table(
|
|||
DateTime(timezone=True),
|
||||
nullable=True,
|
||||
),
|
||||
# Composite foreign key constraint for workspaces
|
||||
ForeignKeyConstraint(
|
||||
["workspace_name", "tenant_id"],
|
||||
["workspaces.name", "workspaces.tenant_id"],
|
||||
),
|
||||
# Composite foreign key constraint for sessions
|
||||
ForeignKeyConstraint(
|
||||
["session_name", "workspace_name"],
|
||||
["sessions.name", "sessions.workspace_name"],
|
||||
["session_name", "workspace_name", "tenant_id"],
|
||||
["sessions.name", "sessions.workspace_name", "sessions.tenant_id"],
|
||||
),
|
||||
# Composite foreign key constraint for peers
|
||||
ForeignKeyConstraint(
|
||||
["peer_name", "workspace_name"],
|
||||
["peers.name", "peers.workspace_name"],
|
||||
["peer_name", "workspace_name", "tenant_id"],
|
||||
["peers.name", "peers.workspace_name", "peers.tenant_id"],
|
||||
),
|
||||
postgresql_partition_by="HASH (tenant_id)",
|
||||
)
|
||||
|
||||
|
||||
@final
|
||||
class Tenant(Base):
|
||||
"""The tenant primitive: one row per tenant, the FK target for every
|
||||
tenant-scoped table's ``tenant_id``.
|
||||
|
||||
honcho's data plane and the control plane that owns the canonical tenant
|
||||
registry live in separate databases, so a cross-database foreign key to the
|
||||
real source of truth is impossible; this table is honcho's local mirror,
|
||||
kept in sync from the control plane (backfilled for existing tenants;
|
||||
written on provision or upserted on first authenticated request for new
|
||||
ones). It is also the one-row-per-tenant home for facts that have nowhere
|
||||
else to live:
|
||||
|
||||
- ``app_name``: the tenant's original per-instance name. Keeping it lets a
|
||||
tenant's external vector-store namespace stay stable when the tenant is
|
||||
moved onto a shared backend, which avoids a full and very expensive
|
||||
re-embed of its vectors.
|
||||
- ``tier``: whether the tenant runs on a dedicated or a shared backend.
|
||||
"""
|
||||
|
||||
__tablename__: str = "tenants"
|
||||
tenant_id: Mapped[str] = mapped_column(TEXT, primary_key=True)
|
||||
app_name: Mapped[str | None] = mapped_column(TEXT, nullable=True, index=True)
|
||||
tier: Mapped[str] = mapped_column(TEXT, nullable=False, server_default="dedicated")
|
||||
created_at: Mapped[datetime.datetime] = mapped_column(
|
||||
DateTime(timezone=True), server_default=func.now()
|
||||
)
|
||||
|
||||
|
||||
@final
|
||||
class Workspace(Base):
|
||||
__tablename__: str = "workspaces"
|
||||
id: Mapped[str] = mapped_column(TEXT, default=generate_nanoid, primary_key=True)
|
||||
name: Mapped[str] = mapped_column(TEXT, unique=True)
|
||||
# tenant_id is the HASH partition key and leads the composite PK, so it is
|
||||
# declared first. It FKs to the local tenants mirror (see Tenant).
|
||||
tenant_id: Mapped[str] = mapped_column(
|
||||
TEXT, ForeignKey("tenants.tenant_id"), nullable=False
|
||||
)
|
||||
id: Mapped[str] = mapped_column(TEXT, default=generate_nanoid)
|
||||
name: Mapped[str] = mapped_column(TEXT)
|
||||
created_at: Mapped[datetime.datetime] = mapped_column(
|
||||
DateTime(timezone=True), server_default=func.now(), index=True
|
||||
)
|
||||
|
|
@ -120,16 +167,25 @@ class Workspace(Base):
|
|||
webhook_endpoints = relationship("WebhookEndpoint", back_populates="workspace")
|
||||
|
||||
__table_args__ = (
|
||||
# Partitioned by HASH(tenant_id): Postgres requires the partition key in
|
||||
# the PK and in every UNIQUE. `name` is unique WITHIN a tenant, not
|
||||
# globally — many tenants share the SDK-default "default" workspace.
|
||||
PrimaryKeyConstraint("tenant_id", "id"),
|
||||
UniqueConstraint("tenant_id", "name"),
|
||||
CheckConstraint("length(id) = 21", name="id_length"),
|
||||
CheckConstraint("length(name) <= 512", name="name_length"),
|
||||
CheckConstraint("id ~ '^[A-Za-z0-9_-]+$'", name="id_format"),
|
||||
{"postgresql_partition_by": "HASH (tenant_id)"},
|
||||
)
|
||||
|
||||
|
||||
@final
|
||||
class Peer(Base):
|
||||
__tablename__: str = "peers"
|
||||
id: Mapped[str] = mapped_column(TEXT, default=generate_nanoid, primary_key=True)
|
||||
tenant_id: Mapped[str] = mapped_column(
|
||||
TEXT, ForeignKey("tenants.tenant_id"), nullable=False
|
||||
)
|
||||
id: Mapped[str] = mapped_column(TEXT, default=generate_nanoid)
|
||||
name: Mapped[str] = mapped_column(TEXT, nullable=False)
|
||||
h_metadata: Mapped[dict[str, Any]] = mapped_column(
|
||||
"metadata", JSONB, default=dict, server_default=text("'{}'::jsonb")
|
||||
|
|
@ -140,9 +196,9 @@ class Peer(Base):
|
|||
created_at: Mapped[datetime.datetime] = mapped_column(
|
||||
DateTime(timezone=True), server_default=func.now(), index=True
|
||||
)
|
||||
workspace_name: Mapped[str] = mapped_column(
|
||||
ForeignKey("workspaces.name"), nullable=False, index=True
|
||||
)
|
||||
# workspace_name's FK to workspaces is now the composite (below), since
|
||||
# workspaces.name is only unique within a tenant.
|
||||
workspace_name: Mapped[str] = mapped_column(TEXT, nullable=False)
|
||||
configuration: Mapped[dict[str, Any]] = mapped_column(
|
||||
JSONB, default=dict, server_default=text("'{}'::jsonb")
|
||||
)
|
||||
|
|
@ -153,20 +209,30 @@ class Peer(Base):
|
|||
)
|
||||
|
||||
__table_args__ = (
|
||||
UniqueConstraint("name", "workspace_name"),
|
||||
PrimaryKeyConstraint("tenant_id", "id"),
|
||||
UniqueConstraint("tenant_id", "name", "workspace_name"),
|
||||
ForeignKeyConstraint(
|
||||
["workspace_name", "tenant_id"],
|
||||
["workspaces.name", "workspaces.tenant_id"],
|
||||
),
|
||||
Index("ix_peers_tenant_workspace", "tenant_id", "workspace_name"),
|
||||
CheckConstraint("length(id) = 21", name="id_length"),
|
||||
CheckConstraint("length(name) <= 512", name="name_length"),
|
||||
CheckConstraint("id ~ '^[A-Za-z0-9_-]+$'", name="id_format"),
|
||||
{"postgresql_partition_by": "HASH (tenant_id)"},
|
||||
)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"Peer(id={self.id}, name={self.name}, workspace_name={self.workspace_name}, created_at={self.created_at}, h_metadata={self.h_metadata}, configuration={self.configuration})"
|
||||
return f"Peer(tenant_id={self.tenant_id}, id={self.id}, name={self.name}, workspace_name={self.workspace_name}, created_at={self.created_at}, h_metadata={self.h_metadata}, configuration={self.configuration})"
|
||||
|
||||
|
||||
@final
|
||||
class Session(Base):
|
||||
__tablename__: str = "sessions"
|
||||
id: Mapped[str] = mapped_column(TEXT, primary_key=True, default=generate_nanoid)
|
||||
tenant_id: Mapped[str] = mapped_column(
|
||||
TEXT, ForeignKey("tenants.tenant_id"), nullable=False
|
||||
)
|
||||
id: Mapped[str] = mapped_column(TEXT, default=generate_nanoid)
|
||||
name: Mapped[str] = mapped_column(TEXT)
|
||||
is_active: Mapped[bool] = mapped_column(default=True, server_default=text("true"))
|
||||
h_metadata: Mapped[dict[str, Any]] = mapped_column(
|
||||
|
|
@ -178,9 +244,7 @@ class Session(Base):
|
|||
created_at: Mapped[datetime.datetime] = mapped_column(
|
||||
DateTime(timezone=True), server_default=func.now(), index=True
|
||||
)
|
||||
workspace_name: Mapped[str] = mapped_column(
|
||||
ForeignKey("workspaces.name"), nullable=False, index=True
|
||||
)
|
||||
workspace_name: Mapped[str] = mapped_column(TEXT, nullable=False)
|
||||
configuration: Mapped[dict[str, Any]] = mapped_column(
|
||||
JSONB, default=dict, server_default=text("'{}'::jsonb")
|
||||
)
|
||||
|
|
@ -192,27 +256,31 @@ class Session(Base):
|
|||
messages = relationship("Message", back_populates="session")
|
||||
|
||||
__table_args__ = (
|
||||
UniqueConstraint("name", "workspace_name"),
|
||||
PrimaryKeyConstraint("tenant_id", "id"),
|
||||
UniqueConstraint("tenant_id", "name", "workspace_name"),
|
||||
ForeignKeyConstraint(
|
||||
["workspace_name", "tenant_id"],
|
||||
["workspaces.name", "workspaces.tenant_id"],
|
||||
),
|
||||
Index("ix_sessions_tenant_workspace", "tenant_id", "workspace_name"),
|
||||
CheckConstraint("length(name) <= 512", name="name_length"),
|
||||
CheckConstraint("length(id) = 21", name="id_length"),
|
||||
CheckConstraint("id ~ '^[A-Za-z0-9_-]+$'", name="id_format"),
|
||||
{"postgresql_partition_by": "HASH (tenant_id)"},
|
||||
)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"Session(id={self.id}, name={self.name}, workspace_name={self.workspace_name}, is_active={self.is_active}, created_at={self.created_at}, h_metadata={self.h_metadata})"
|
||||
return f"Session(tenant_id={self.tenant_id}, id={self.id}, name={self.name}, workspace_name={self.workspace_name}, is_active={self.is_active}, created_at={self.created_at}, h_metadata={self.h_metadata})"
|
||||
|
||||
|
||||
@final
|
||||
class Message(Base):
|
||||
__tablename__: str = "messages"
|
||||
id: Mapped[int] = mapped_column(
|
||||
BigInteger, Identity(), primary_key=True, autoincrement=True
|
||||
)
|
||||
public_id: Mapped[str] = mapped_column(
|
||||
TEXT,
|
||||
unique=True,
|
||||
default=generate_nanoid,
|
||||
tenant_id: Mapped[str] = mapped_column(
|
||||
TEXT, ForeignKey("tenants.tenant_id"), nullable=False
|
||||
)
|
||||
id: Mapped[int] = mapped_column(BigInteger, Identity(), autoincrement=True)
|
||||
public_id: Mapped[str] = mapped_column(TEXT, default=generate_nanoid)
|
||||
# NOTE: Messages in Honcho 2.0 could historically be stored outside of a session.
|
||||
# We have since assigned all of these messages to a default session.
|
||||
session_name: Mapped[str] = mapped_column(TEXT, nullable=False)
|
||||
|
|
@ -230,72 +298,78 @@ class Message(Base):
|
|||
DateTime(timezone=True), server_default=func.now(), index=True
|
||||
)
|
||||
# Note: Foreign key relationships established via composite ForeignKeyConstraint below
|
||||
peer_name: Mapped[str] = mapped_column(TEXT, index=True)
|
||||
workspace_name: Mapped[str] = mapped_column(TEXT, index=True)
|
||||
peer_name: Mapped[str] = mapped_column(TEXT)
|
||||
workspace_name: Mapped[str] = mapped_column(TEXT)
|
||||
|
||||
session = relationship("Session", back_populates="messages")
|
||||
|
||||
__table_args__ = (
|
||||
PrimaryKeyConstraint("tenant_id", "id"),
|
||||
# (tenant_id, public_id) is the unique that message_embeddings' FK targets.
|
||||
UniqueConstraint("tenant_id", "public_id"),
|
||||
CheckConstraint("length(public_id) = 21", name="public_id_length"),
|
||||
CheckConstraint("public_id ~ '^[A-Za-z0-9_-]+$'", name="public_id_format"),
|
||||
CheckConstraint("length(content) <= 65535", name="content_length"),
|
||||
# Composite foreign key constraint for sessions
|
||||
ForeignKeyConstraint(
|
||||
["session_name", "workspace_name"],
|
||||
["sessions.name", "sessions.workspace_name"],
|
||||
["session_name", "workspace_name", "tenant_id"],
|
||||
["sessions.name", "sessions.workspace_name", "sessions.tenant_id"],
|
||||
),
|
||||
# Composite foreign key constraint for peers
|
||||
ForeignKeyConstraint(
|
||||
["peer_name", "workspace_name"],
|
||||
["peers.name", "peers.workspace_name"],
|
||||
["peer_name", "workspace_name", "tenant_id"],
|
||||
["peers.name", "peers.workspace_name", "peers.tenant_id"],
|
||||
),
|
||||
Index(
|
||||
"ix_messages_session_lookup",
|
||||
"tenant_id",
|
||||
"session_name",
|
||||
"id",
|
||||
postgresql_include=["id", "created_at"],
|
||||
postgresql_include=["created_at"],
|
||||
),
|
||||
Index("ix_messages_tenant_peer", "tenant_id", "peer_name"),
|
||||
UniqueConstraint(
|
||||
"tenant_id",
|
||||
"workspace_name",
|
||||
"session_name",
|
||||
"seq_in_session",
|
||||
),
|
||||
# Full text search index on content column
|
||||
# GIN can't lead with a scalar column without btree_gin; the table is
|
||||
# HASH(tenant_id)-partitioned, so this index is per-partition — queries
|
||||
# prune to one partition, then tenant_id filters the FTS candidates.
|
||||
Index(
|
||||
"ix_messages_content_gin",
|
||||
text("to_tsvector('english', content)"),
|
||||
postgresql_using="gin",
|
||||
),
|
||||
{"postgresql_partition_by": "HASH (tenant_id)"},
|
||||
)
|
||||
|
||||
@override
|
||||
def __repr__(self) -> str:
|
||||
return f"Message(id={self.id}, session_name={self.session_name}, peer_name={self.peer_name}, content={self.content})"
|
||||
return f"Message(tenant_id={self.tenant_id}, id={self.id}, session_name={self.session_name}, peer_name={self.peer_name}, content={self.content})"
|
||||
|
||||
|
||||
@final
|
||||
class MessageEmbedding(Base):
|
||||
__tablename__: str = "message_embeddings"
|
||||
|
||||
id: Mapped[int] = mapped_column(
|
||||
BigInteger, Identity(), primary_key=True, autoincrement=True
|
||||
tenant_id: Mapped[str] = mapped_column(
|
||||
TEXT, ForeignKey("tenants.tenant_id"), nullable=False
|
||||
)
|
||||
id: Mapped[int] = mapped_column(BigInteger, Identity(), autoincrement=True)
|
||||
content: Mapped[str] = mapped_column(TEXT)
|
||||
embedding: MappedColumn[Any] = mapped_column(Vector(_VECTOR_DIM), nullable=True)
|
||||
message_id: Mapped[str] = mapped_column(
|
||||
ForeignKey("messages.public_id", ondelete="CASCADE"), nullable=False, index=True
|
||||
)
|
||||
workspace_name: Mapped[str] = mapped_column(
|
||||
ForeignKey("workspaces.name"), nullable=False, index=True
|
||||
)
|
||||
session_name: Mapped[str] = mapped_column(TEXT, nullable=False, index=True)
|
||||
peer_name: Mapped[str] = mapped_column(TEXT, nullable=False, index=True)
|
||||
message_id: Mapped[str] = mapped_column(TEXT, nullable=False)
|
||||
workspace_name: Mapped[str] = mapped_column(TEXT, nullable=False)
|
||||
session_name: Mapped[str] = mapped_column(TEXT, nullable=False)
|
||||
peer_name: Mapped[str] = mapped_column(TEXT, nullable=False)
|
||||
created_at: Mapped[datetime.datetime] = mapped_column(
|
||||
DateTime(timezone=True), server_default=func.now(), index=True
|
||||
)
|
||||
# Vector sync state tracking
|
||||
sync_state: Mapped[VectorSyncState] = mapped_column(
|
||||
TEXT, nullable=False, server_default="pending", index=True
|
||||
TEXT, nullable=False, server_default="pending"
|
||||
)
|
||||
last_sync_at: Mapped[datetime.datetime | None] = mapped_column(
|
||||
DateTime(timezone=True), nullable=True
|
||||
|
|
@ -305,16 +379,25 @@ class MessageEmbedding(Base):
|
|||
)
|
||||
|
||||
__table_args__ = (
|
||||
# Compound foreign key constraints
|
||||
PrimaryKeyConstraint("tenant_id", "id"),
|
||||
# message_id → messages.public_id is now composite: messages' unique is
|
||||
# (tenant_id, public_id) under partitioning.
|
||||
ForeignKeyConstraint(
|
||||
["session_name", "workspace_name"],
|
||||
["sessions.name", "sessions.workspace_name"],
|
||||
["tenant_id", "message_id"],
|
||||
["messages.tenant_id", "messages.public_id"],
|
||||
ondelete="CASCADE",
|
||||
),
|
||||
ForeignKeyConstraint(
|
||||
["peer_name", "workspace_name"],
|
||||
["peers.name", "peers.workspace_name"],
|
||||
["session_name", "workspace_name", "tenant_id"],
|
||||
["sessions.name", "sessions.workspace_name", "sessions.tenant_id"],
|
||||
),
|
||||
# HNSW index on embedding column for efficient similarity search
|
||||
ForeignKeyConstraint(
|
||||
["peer_name", "workspace_name", "tenant_id"],
|
||||
["peers.name", "peers.workspace_name", "peers.tenant_id"],
|
||||
),
|
||||
Index("ix_message_embeddings_tenant_message", "tenant_id", "message_id"),
|
||||
# HNSW is a single-column vector index (can't lead with tenant_id); it
|
||||
# becomes per-partition automatically under HASH(tenant_id).
|
||||
Index(
|
||||
"ix_message_embeddings_embedding_hnsw",
|
||||
"embedding",
|
||||
|
|
@ -322,12 +405,15 @@ class MessageEmbedding(Base):
|
|||
postgresql_with={"m": 16, "ef_construction": 64},
|
||||
postgresql_ops={"embedding": "vector_cosine_ops"},
|
||||
),
|
||||
# Composite index for efficient reconciliation queries
|
||||
# NOT tenant_id-leading on purpose: the reconciler scans this cross-tenant
|
||||
# (sync_state='pending' over all tenants), so a tenant_id prefix wouldn't
|
||||
# help. (Also drops the redundant single-column sync_state index.)
|
||||
Index(
|
||||
"ix_message_embeddings_sync_state_last_sync_at",
|
||||
"sync_state",
|
||||
"last_sync_at",
|
||||
),
|
||||
{"postgresql_partition_by": "HASH (tenant_id)"},
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -335,9 +421,12 @@ class MessageEmbedding(Base):
|
|||
class Collection(Base):
|
||||
__tablename__: str = "collections"
|
||||
|
||||
id: Mapped[str] = mapped_column(TEXT, default=generate_nanoid, primary_key=True)
|
||||
observer: Mapped[str] = mapped_column(TEXT, index=True)
|
||||
observed: Mapped[str] = mapped_column(TEXT, index=True)
|
||||
tenant_id: Mapped[str] = mapped_column(
|
||||
TEXT, ForeignKey("tenants.tenant_id"), nullable=False
|
||||
)
|
||||
id: Mapped[str] = mapped_column(TEXT, default=generate_nanoid)
|
||||
observer: Mapped[str] = mapped_column(TEXT)
|
||||
observed: Mapped[str] = mapped_column(TEXT)
|
||||
created_at: Mapped[datetime.datetime] = mapped_column(
|
||||
DateTime(timezone=True), server_default=func.now(), index=True
|
||||
)
|
||||
|
|
@ -350,35 +439,43 @@ class Collection(Base):
|
|||
documents = relationship(
|
||||
"Document", back_populates="collection", cascade="all, delete, delete-orphan"
|
||||
)
|
||||
workspace_name: Mapped[str] = mapped_column(
|
||||
ForeignKey("workspaces.name"), nullable=False, index=True
|
||||
)
|
||||
workspace_name: Mapped[str] = mapped_column(TEXT, nullable=False)
|
||||
|
||||
__table_args__ = (
|
||||
PrimaryKeyConstraint("tenant_id", "id"),
|
||||
UniqueConstraint(
|
||||
"tenant_id",
|
||||
"observer",
|
||||
"observed",
|
||||
"workspace_name",
|
||||
),
|
||||
CheckConstraint("length(id) = 21", name="id_length"),
|
||||
CheckConstraint("id ~ '^[A-Za-z0-9_-]+$'", name="id_format"),
|
||||
ForeignKeyConstraint(
|
||||
["workspace_name", "tenant_id"],
|
||||
["workspaces.name", "workspaces.tenant_id"],
|
||||
),
|
||||
# Composite foreign key constraint for observer peer
|
||||
ForeignKeyConstraint(
|
||||
["observer", "workspace_name"],
|
||||
["peers.name", "peers.workspace_name"],
|
||||
["observer", "workspace_name", "tenant_id"],
|
||||
["peers.name", "peers.workspace_name", "peers.tenant_id"],
|
||||
),
|
||||
# Composite foreign key constraint for observed peer
|
||||
ForeignKeyConstraint(
|
||||
["observed", "workspace_name"],
|
||||
["peers.name", "peers.workspace_name"],
|
||||
["observed", "workspace_name", "tenant_id"],
|
||||
["peers.name", "peers.workspace_name", "peers.tenant_id"],
|
||||
),
|
||||
{"postgresql_partition_by": "HASH (tenant_id)"},
|
||||
)
|
||||
|
||||
|
||||
@final
|
||||
class Document(Base):
|
||||
__tablename__: str = "documents"
|
||||
id: Mapped[str] = mapped_column(TEXT, default=generate_nanoid, primary_key=True)
|
||||
tenant_id: Mapped[str] = mapped_column(
|
||||
TEXT, ForeignKey("tenants.tenant_id"), nullable=False
|
||||
)
|
||||
id: Mapped[str] = mapped_column(TEXT, default=generate_nanoid)
|
||||
internal_metadata: Mapped[dict[str, Any]] = mapped_column(
|
||||
"internal_metadata", JSONB, default=dict, server_default=text("'{}'::jsonb")
|
||||
)
|
||||
|
|
@ -397,19 +494,17 @@ class Document(Base):
|
|||
DateTime(timezone=True), server_default=func.now(), index=True
|
||||
)
|
||||
|
||||
observer: Mapped[str] = mapped_column(TEXT, index=True)
|
||||
observed: Mapped[str] = mapped_column(TEXT, index=True)
|
||||
workspace_name: Mapped[str] = mapped_column(
|
||||
ForeignKey("workspaces.name"), nullable=False, index=True
|
||||
)
|
||||
session_name: Mapped[str | None] = mapped_column(TEXT, nullable=True, index=True)
|
||||
observer: Mapped[str] = mapped_column(TEXT)
|
||||
observed: Mapped[str] = mapped_column(TEXT)
|
||||
workspace_name: Mapped[str] = mapped_column(TEXT, nullable=False)
|
||||
session_name: Mapped[str | None] = mapped_column(TEXT, nullable=True)
|
||||
deleted_at: Mapped[datetime.datetime | None] = mapped_column(
|
||||
DateTime(timezone=True), nullable=True, index=True, default=None
|
||||
)
|
||||
|
||||
# Vector sync state tracking
|
||||
sync_state: Mapped[VectorSyncState] = mapped_column(
|
||||
TEXT, nullable=False, server_default="pending", index=True
|
||||
TEXT, nullable=False, server_default="pending"
|
||||
)
|
||||
last_sync_at: Mapped[datetime.datetime | None] = mapped_column(
|
||||
DateTime(timezone=True), nullable=True
|
||||
|
|
@ -421,34 +516,49 @@ class Document(Base):
|
|||
collection = relationship("Collection", back_populates="documents")
|
||||
|
||||
__table_args__ = (
|
||||
PrimaryKeyConstraint("tenant_id", "id"),
|
||||
CheckConstraint("length(id) = 21", name="id_length"),
|
||||
CheckConstraint("length(content) <= 65535", name="content_length"),
|
||||
CheckConstraint("id ~ '^[A-Za-z0-9_-]+$'", name="id_format"),
|
||||
# Composite foreign key constraint for workspaces
|
||||
ForeignKeyConstraint(
|
||||
["workspace_name", "tenant_id"],
|
||||
["workspaces.name", "workspaces.tenant_id"],
|
||||
),
|
||||
# Composite foreign key constraint for collections
|
||||
ForeignKeyConstraint(
|
||||
["observer", "observed", "workspace_name"],
|
||||
["observer", "observed", "workspace_name", "tenant_id"],
|
||||
[
|
||||
"collections.observer",
|
||||
"collections.observed",
|
||||
"collections.workspace_name",
|
||||
"collections.tenant_id",
|
||||
],
|
||||
),
|
||||
# Composite foreign key constraint for observer peer
|
||||
ForeignKeyConstraint(
|
||||
["observer", "workspace_name"],
|
||||
["peers.name", "peers.workspace_name"],
|
||||
["observer", "workspace_name", "tenant_id"],
|
||||
["peers.name", "peers.workspace_name", "peers.tenant_id"],
|
||||
),
|
||||
# Composite foreign key constraint for observed peer
|
||||
ForeignKeyConstraint(
|
||||
["observed", "workspace_name"],
|
||||
["peers.name", "peers.workspace_name"],
|
||||
["observed", "workspace_name", "tenant_id"],
|
||||
["peers.name", "peers.workspace_name", "peers.tenant_id"],
|
||||
),
|
||||
# Composite foreign key constraint for sessions
|
||||
ForeignKeyConstraint(
|
||||
["session_name", "workspace_name"],
|
||||
["sessions.name", "sessions.workspace_name"],
|
||||
["session_name", "workspace_name", "tenant_id"],
|
||||
["sessions.name", "sessions.workspace_name", "sessions.tenant_id"],
|
||||
),
|
||||
# HNSW index on embedding column
|
||||
# Tenant-scoped collection lookups (replaces the single observer/observed indexes)
|
||||
Index(
|
||||
"ix_documents_tenant_collection",
|
||||
"tenant_id",
|
||||
"observer",
|
||||
"observed",
|
||||
"workspace_name",
|
||||
),
|
||||
# HNSW is a single-column vector index (per-partition under HASH(tenant_id))
|
||||
Index(
|
||||
"ix_documents_embedding_hnsw",
|
||||
"embedding",
|
||||
|
|
@ -464,12 +574,14 @@ class Document(Base):
|
|||
"source_ids",
|
||||
postgresql_using="gin",
|
||||
),
|
||||
# Composite index for efficient reconciliation queries
|
||||
# Reconciler scans this cross-tenant (sync_state='pending'), so NOT
|
||||
# tenant_id-leading. Also drops the redundant single-column sync_state index.
|
||||
Index(
|
||||
"ix_documents_sync_state_last_sync_at",
|
||||
"sync_state",
|
||||
"last_sync_at",
|
||||
),
|
||||
{"postgresql_partition_by": "HASH (tenant_id)"},
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -479,9 +591,13 @@ class QueueItem(Base):
|
|||
id: Mapped[int] = mapped_column(
|
||||
BigInteger, Identity(), primary_key=True, autoincrement=True
|
||||
)
|
||||
session_id: Mapped[str | None] = mapped_column(
|
||||
ForeignKey("sessions.id"), nullable=True, index=True
|
||||
)
|
||||
# Service table: NOT partitioned and drained-not-copied at migration, so it
|
||||
# keeps a sole-id PK. tenant_id is a plain attribution / fair-scheduling column
|
||||
# (no FK, no RLS); the FKs to the now-partitioned sessions / messages /
|
||||
# workspaces are dropped — the app manages queue lifecycle and already
|
||||
# tolerates missing referents.
|
||||
tenant_id: Mapped[str | None] = mapped_column(TEXT, nullable=True, index=True)
|
||||
session_id: Mapped[str | None] = mapped_column(TEXT, nullable=True, index=True)
|
||||
work_unit_key: Mapped[str] = mapped_column(TEXT, nullable=False)
|
||||
|
||||
task_type: Mapped[TaskType] = mapped_column(TEXT, nullable=False)
|
||||
|
|
@ -493,12 +609,8 @@ class QueueItem(Base):
|
|||
created_at: Mapped[datetime.datetime] = mapped_column(
|
||||
DateTime(timezone=True), server_default=func.now(), index=True
|
||||
)
|
||||
workspace_name: Mapped[str | None] = mapped_column(
|
||||
ForeignKey("workspaces.name"), nullable=True, index=True
|
||||
)
|
||||
message_id: Mapped[int | None] = mapped_column(
|
||||
BigInteger, ForeignKey("messages.id"), nullable=True
|
||||
)
|
||||
workspace_name: Mapped[str | None] = mapped_column(TEXT, nullable=True, index=True)
|
||||
message_id: Mapped[int | None] = mapped_column(BigInteger, nullable=True)
|
||||
|
||||
__table_args__ = (
|
||||
Index(
|
||||
|
|
@ -538,6 +650,9 @@ class ActiveQueueSession(Base):
|
|||
|
||||
id: Mapped[str] = mapped_column(TEXT, default=generate_nanoid, primary_key=True)
|
||||
|
||||
# Service table (unpartitioned): tenant_id is plain attribution, no FK / RLS.
|
||||
tenant_id: Mapped[str | None] = mapped_column(TEXT, nullable=True)
|
||||
|
||||
work_unit_key: Mapped[str] = mapped_column(TEXT, unique=True)
|
||||
|
||||
last_updated: Mapped[datetime.datetime] = mapped_column(
|
||||
|
|
@ -548,10 +663,11 @@ class ActiveQueueSession(Base):
|
|||
@final
|
||||
class WebhookEndpoint(Base):
|
||||
__tablename__: str = "webhook_endpoints"
|
||||
id: Mapped[str] = mapped_column(TEXT, default=generate_nanoid, primary_key=True)
|
||||
workspace_name: Mapped[str] = mapped_column(
|
||||
ForeignKey("workspaces.name"), nullable=False, index=True
|
||||
tenant_id: Mapped[str] = mapped_column(
|
||||
TEXT, ForeignKey("tenants.tenant_id"), nullable=False
|
||||
)
|
||||
id: Mapped[str] = mapped_column(TEXT, default=generate_nanoid)
|
||||
workspace_name: Mapped[str] = mapped_column(TEXT, nullable=False)
|
||||
url: Mapped[str] = mapped_column(TEXT, nullable=False)
|
||||
created_at: Mapped[datetime.datetime] = mapped_column(
|
||||
DateTime(timezone=True), server_default=func.now()
|
||||
|
|
@ -559,10 +675,19 @@ class WebhookEndpoint(Base):
|
|||
|
||||
workspace = relationship("Workspace", back_populates="webhook_endpoints")
|
||||
|
||||
__table_args__ = (CheckConstraint("length(url) <= 2048", name="url_length"),)
|
||||
__table_args__ = (
|
||||
PrimaryKeyConstraint("tenant_id", "id"),
|
||||
ForeignKeyConstraint(
|
||||
["workspace_name", "tenant_id"],
|
||||
["workspaces.name", "workspaces.tenant_id"],
|
||||
),
|
||||
Index("ix_webhook_endpoints_tenant_workspace", "tenant_id", "workspace_name"),
|
||||
CheckConstraint("length(url) <= 2048", name="url_length"),
|
||||
{"postgresql_partition_by": "HASH (tenant_id)"},
|
||||
)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"WebhookEndpoint(id={self.id}, workspace_name={self.workspace_name}, url={self.url})"
|
||||
return f"WebhookEndpoint(tenant_id={self.tenant_id}, id={self.id}, workspace_name={self.workspace_name}, url={self.url})"
|
||||
|
||||
|
||||
@final
|
||||
|
|
@ -570,6 +695,7 @@ class SessionPeer(Base):
|
|||
__table__: Table = session_peers_table
|
||||
|
||||
# Type annotations for the columns
|
||||
tenant_id: Mapped[str]
|
||||
workspace_name: Mapped[str]
|
||||
session_name: Mapped[str]
|
||||
peer_name: Mapped[str]
|
||||
|
|
|
|||
Loading…
Reference in New Issue