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

Standardize DB Constraint Conventions (#272) 

* fix: (db) Add standard naming conventions to SQLAlchemy Declarative Base

* chore: remove unnecessary relationships

* fix: (tests) handle old migrations being made before conventions were applied

* fix: wip migration to standardize naming for constraints

* fix (db): Migration wip)

* chore: rebase migration

* fix (db) Add migration tests

* chore: Code Review Comments

* fix: (db) fix remaining inconsistent index
This commit is contained in:
Vineeth Voruganti 2025-11-20 11:50:54 -05:00 committed by GitHub
parent 841e4bb808
commit a7520ce21d
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
14 changed files with 864 additions and 126 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

0
init.sql → database/init.sql
View File

2
migrations/env.py
View File

@ -143,7 +143,7 @@ def run_migrations_online() -> None:
connectable = engine_from_config(
configuration,
prefix="sqlalchemy.",
echo=True,
echo=False,
connect_args={
"prepare_threshold": None,
"options": "-c statement_timeout=300000", # 5 minutes in milliseconds

14
migrations/utils.py
View File

@ -68,6 +68,8 @@ def constraint_exists(
elif type == "primary":
constraint = inspector.get_pk_constraint(table_name, schema=schema)
return constraint["name"] == constraint_name
elif type == "foreignkey":
constraints = inspector.get_foreign_keys(table_name, schema=schema)
else:
raise ValueError(f"Invalid constraint type: {type}")
return any(constraint["name"] == constraint_name for constraint in constraints)
@ -151,9 +153,11 @@ def make_column_non_nullable_safe(table_name: str, column_name: str) -> None:
)
# Step 5: Drop the redundant CHECK constraint
op.drop_constraint(
constraint_name,
table_name,
type_="check",
schema=schema,
conn.execute(
sa.text(
f"""
ALTER TABLE {quoted_schema}.{quoted_table}
DROP CONSTRAINT {quoted_constraint}
"""
)
)

8
migrations/versions/05486ce795d5_make_session_name_required_on_messages.py
View File

@ -197,8 +197,12 @@ def upgrade() -> None:
# Step 9: Update the collections table name_length check constraint from 512 to 1025
print("Updating collections table name_length check constraint from 512 to 1025")
# Check for both old name and naming-convention-generated name
if constraint_exists("collections", "name_length", "check"):
op.drop_constraint("name_length", "collections", schema=schema)
elif constraint_exists("collections", "ck_collections_name_length", "check"):
op.drop_constraint("ck_collections_name_length", "collections", schema=schema)
op.create_check_constraint(
"name_length", "collections", "length(name) <= 1025", schema=schema
)
@ -215,8 +219,12 @@ def downgrade() -> None:
"Reverting collections table name_length check constraint from 1025 back to 512"
)
# Check for both old name and naming-convention-generated name
if constraint_exists("collections", "name_length", "check"):
op.drop_constraint("name_length", "collections", schema=schema)
elif constraint_exists("collections", "ck_collections_name_length", "check"):
op.drop_constraint("ck_collections_name_length", "collections", schema=schema)
op.create_check_constraint(
"name_length", "collections", "length(name) <= 512", schema=schema
)

53
migrations/versions/20f89a421aff_rename_metamessage_type_to_label.py
View File

@ -10,8 +10,9 @@ from collections.abc import Sequence
import sqlalchemy as sa
from alembic import op
from sqlalchemy import text
from migrations.utils import constraint_exists
from migrations.utils import constraint_exists, get_schema
# revision identifiers, used by Alembic.
revision: str = "20f89a421aff"
@ -21,6 +22,9 @@ depends_on: str | Sequence[str] | None = None
def upgrade() -> None:
schema = get_schema()
conn = op.get_bind()
# ### commands auto generated by Alembic - please adjust! ###
op.alter_column(
"metamessages",
@ -63,13 +67,38 @@ def upgrade() -> None:
)
# Rename check constraint
if constraint_exists("metamessages", "metamessage_type_length", "check"):
op.execute("ALTER TABLE metamessages DROP CONSTRAINT metamessage_type_length;")
op.create_check_constraint("label_length", "metamessages", "length(label) <= 512")
# Handle both naming convention and non-convention names for backward compatibility
if constraint_exists(
"metamessages", "ck_metamessages_metamessage_type_length", "check"
):
conn.execute(
text(
f"ALTER TABLE {schema}.metamessages DROP CONSTRAINT ck_metamessages_metamessage_type_length"
)
)
elif constraint_exists("metamessages", "metamessage_type_length", "check"):
conn.execute(
text(
f"ALTER TABLE {schema}.metamessages DROP CONSTRAINT metamessage_type_length"
)
)
conn.execute(
text(
f"""
ALTER TABLE {schema}.metamessages
ADD CONSTRAINT label_length
CHECK (length(label) <= 512)
"""
)
)
# ### end Alembic commands ###
def downgrade() -> None:
schema = get_schema()
conn = op.get_bind()
# ### commands auto generated by Alembic - please adjust! ###
op.alter_column(
"metamessages",
@ -112,10 +141,16 @@ def downgrade() -> None:
)
# Revert check constraint rename
op.execute("ALTER TABLE metamessages DROP CONSTRAINT label_length;")
op.create_check_constraint(
"metamessage_type_length",
"metamessages",
"length(metamessage_type) <= 512",
conn.execute(
text(f"ALTER TABLE {schema}.metamessages DROP CONSTRAINT label_length")
)
conn.execute(
text(
f"""
ALTER TABLE {schema}.metamessages
ADD CONSTRAINT metamessage_type_length
CHECK (length(metamessage_type) <= 512)
"""
)
)
# ### end Alembic commands ###

33
migrations/versions/29ade7350c19_remove_document_level_valid_constraint.py
View File

@ -10,6 +10,7 @@ from collections.abc import Sequence
import sqlalchemy as sa
from alembic import op
from sqlalchemy import text
from migrations.utils import constraint_exists, get_schema
@ -27,13 +28,18 @@ def upgrade() -> None:
connection = op.get_bind()
inspector = sa.inspect(connection)
# Drop CHECK constraint for level
# Drop CHECK constraint for level (check both naming convention and plain names)
# On old databases: "level_valid"
# On new databases with naming convention: "ck_documents_level_valid"
if constraint_exists("documents", "level_valid", "check", inspector):
op.drop_constraint(
"level_valid",
"documents",
type_="check",
schema=schema,
connection.execute(
text(f'ALTER TABLE {schema}.documents DROP CONSTRAINT "level_valid"')
)
elif constraint_exists("documents", "ck_documents_level_valid", "check", inspector):
connection.execute(
text(
f'ALTER TABLE {schema}.documents DROP CONSTRAINT "ck_documents_level_valid"'
)
)
@ -66,11 +72,14 @@ def downgrade() -> None:
if result.rowcount == 0:
break
# Recreate CHECK constraint for level
# Recreate CHECK constraint for level (use plain name to match original migration)
if not constraint_exists("documents", "level_valid", "check", inspector):
op.create_check_constraint(
"level_valid",
"documents",
"level IN ('explicit', 'deductive')",
schema=schema,
connection.execute(
text(
f"""
ALTER TABLE {schema}.documents
ADD CONSTRAINT level_valid
CHECK (level IN ('explicit', 'deductive'))
"""
)
)

12
migrations/versions/b765d82110bd_change_metamessages_to_user_level_with_.py
View File

@ -222,10 +222,14 @@ def upgrade() -> None:
)
if not constraint_exists:
op.create_check_constraint(
"message_requires_session",
"metamessages",
"(message_id IS NULL) OR (session_id IS NOT NULL)",
conn.execute(
text(
f"""
ALTER TABLE {schema}.metamessages
ADD CONSTRAINT message_requires_session
CHECK ((message_id IS NULL) OR (session_id IS NOT NULL))
"""
)
)
print("Created message_requires_session constraint")
else:

20
migrations/versions/b8183c5ffb48_codify_document_level_and_times_derived.py
View File

@ -100,11 +100,14 @@ def upgrade() -> None:
# Step 6: Add CHECK constraint for level
if not constraint_exists("documents", "level_valid", "check", inspector):
op.create_check_constraint(
"level_valid",
"documents",
"level IN ('explicit', 'deductive')",
schema=schema,
connection.execute(
text(
f"""
ALTER TABLE {schema}.documents
ADD CONSTRAINT level_valid
CHECK (level IN ('explicit', 'deductive'))
"""
)
)
@ -115,11 +118,8 @@ def downgrade() -> None:
# Step 1: Drop CHECK constraint for level
if constraint_exists("documents", "level_valid", "check", inspector):
op.drop_constraint(
"level_valid",
"documents",
type_="check",
schema=schema,
connection.execute(
text(f"ALTER TABLE {schema}.documents DROP CONSTRAINT level_valid")
)
# Step 2: Copy level and times_derived back to internal_metadata in batches (optional, for safety)

525
migrations/versions/baa22cad81e2_standardize_constraint_names.py Normal file
View File

@ -0,0 +1,525 @@
"""standardize_constraint_names
Revision ID: baa22cad81e2
Revises: 29ade7350c19
Create Date: 2025-11-15 01:16:40.937103
"""
from collections.abc import Sequence
import sqlalchemy as sa
from alembic import op
from sqlalchemy import text
from migrations.utils import constraint_exists, get_schema, index_exists
# revision identifiers, used by Alembic.
revision: str = "baa22cad81e2"
down_revision: str | None = "29ade7350c19"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
schema = get_schema()
def upgrade() -> None:
conn = op.get_bind()
inspector = sa.inspect(conn)
# ============================================================
# CONSTRAINT RENAMES (using raw SQL)
# ============================================================
# Unique Constraints
if constraint_exists(
"active_queue_sessions", "unique_work_unit_key", "unique", inspector
):
conn.execute(
text(
f'ALTER TABLE {schema}.active_queue_sessions RENAME CONSTRAINT "unique_work_unit_key" TO "uq_active_queue_sessions_work_unit_key"'
)
)
if constraint_exists(
"collections", "unique_observer_observed_collection", "unique", inspector
):
conn.execute(
text(
f'ALTER TABLE {schema}.collections RENAME CONSTRAINT "unique_observer_observed_collection" TO "uq_collections_observer_observed_workspace_name"'
)
)
if constraint_exists("peers", "unique_name_workspace_peer", "unique", inspector):
conn.execute(
text(
f'ALTER TABLE {schema}.peers RENAME CONSTRAINT "unique_name_workspace_peer" TO "uq_peers_name_workspace_name"'
)
)
if constraint_exists("sessions", "unique_session_name", "unique", inspector):
conn.execute(
text(
f'ALTER TABLE {schema}.sessions RENAME CONSTRAINT "unique_session_name" TO "uq_sessions_name_workspace_name"'
)
)
if constraint_exists("messages", "uq_messages_session_seq", "unique", inspector):
conn.execute(
text(
f'ALTER TABLE {schema}.messages RENAME CONSTRAINT "uq_messages_session_seq" TO "uq_messages_workspace_name_session_name_seq_in_session"'
)
)
if constraint_exists("workspaces", "uq_apps_name", "unique", inspector):
conn.execute(
text(
f'ALTER TABLE {schema}.workspaces RENAME CONSTRAINT "uq_apps_name" TO "uq_workspaces_name"'
)
)
# Foreign Keys
if constraint_exists(
"message_embeddings",
"message_embeddings_message_id_fkey",
"foreignkey",
inspector,
):
conn.execute(
text(
f'ALTER TABLE {schema}.message_embeddings RENAME CONSTRAINT "message_embeddings_message_id_fkey" TO "fk_message_embeddings_message_id_messages"'
)
)
# ============================================================
# INDEX RENAMES (using raw SQL)
# ============================================================
# Refresh inspector after constraint renames to ensure we see current state
inspector = sa.inspect(conn)
# Collections
if index_exists("collections", "idx_collections_observed", inspector):
conn.execute(
text(
f'ALTER INDEX {schema}."idx_collections_observed" RENAME TO "ix_collections_observed"'
)
)
if index_exists("collections", "idx_collections_observer", inspector):
conn.execute(
text(
f'ALTER INDEX {schema}."idx_collections_observer" RENAME TO "ix_collections_observer"'
)
)
# Documents
if index_exists("documents", "idx_documents_embedding_hnsw", inspector):
conn.execute(
text(
f'ALTER INDEX {schema}."idx_documents_embedding_hnsw" RENAME TO "ix_documents_embedding_hnsw"'
)
)
if index_exists("documents", "idx_documents_observed", inspector):
conn.execute(
text(
f'ALTER INDEX {schema}."idx_documents_observed" RENAME TO "ix_documents_observed"'
)
)
if index_exists("documents", "idx_documents_observer", inspector):
conn.execute(
text(
f'ALTER INDEX {schema}."idx_documents_observer" RENAME TO "ix_documents_observer"'
)
)
if index_exists("documents", "idx_documents_session_name", inspector):
conn.execute(
text(
f'ALTER INDEX {schema}."idx_documents_session_name" RENAME TO "ix_documents_session_name"'
)
)
# Message Embeddings
if index_exists(
"message_embeddings", "idx_message_embeddings_created_at", inspector
):
conn.execute(
text(
f'ALTER INDEX {schema}."idx_message_embeddings_created_at" RENAME TO "ix_message_embeddings_created_at"'
)
)
if index_exists(
"message_embeddings", "idx_message_embeddings_embedding_hnsw", inspector
):
conn.execute(
text(
f'ALTER INDEX {schema}."idx_message_embeddings_embedding_hnsw" RENAME TO "ix_message_embeddings_embedding_hnsw"'
)
)
if index_exists(
"message_embeddings", "idx_message_embeddings_message_id", inspector
):
conn.execute(
text(
f'ALTER INDEX {schema}."idx_message_embeddings_message_id" RENAME TO "ix_message_embeddings_message_id"'
)
)
if index_exists(
"message_embeddings", "idx_message_embeddings_peer_name", inspector
):
conn.execute(
text(
f'ALTER INDEX {schema}."idx_message_embeddings_peer_name" RENAME TO "ix_message_embeddings_peer_name"'
)
)
if index_exists(
"message_embeddings", "idx_message_embeddings_session_name", inspector
):
conn.execute(
text(
f'ALTER INDEX {schema}."idx_message_embeddings_session_name" RENAME TO "ix_message_embeddings_session_name"'
)
)
if index_exists(
"message_embeddings", "idx_message_embeddings_workspace_name", inspector
):
conn.execute(
text(
f'ALTER INDEX {schema}."idx_message_embeddings_workspace_name" RENAME TO "ix_message_embeddings_workspace_name"'
)
)
# Messages
if index_exists("messages", "idx_messages_content_gin", inspector):
conn.execute(
text(
f'ALTER INDEX {schema}."idx_messages_content_gin" RENAME TO "ix_messages_content_gin"'
)
)
if index_exists("messages", "idx_messages_session_lookup", inspector):
conn.execute(
text(
f'ALTER INDEX {schema}."idx_messages_session_lookup" RENAME TO "ix_messages_session_lookup"'
)
)
# Webhook Endpoints
if index_exists(
"webhook_endpoints", "idx_webhook_endpoints_workspace_lookup", inspector
):
conn.execute(
text(
f'ALTER INDEX {schema}."idx_webhook_endpoints_workspace_lookup" RENAME TO "ix_webhook_endpoints_workspace_name"'
)
)
# ============================================================
# INDEX DELETIONS (redundant indexes being removed)
# ============================================================
# Drop redundant indexes on primary keys and unique columns
if index_exists("messages", "ix_messages_id", inspector):
op.drop_index("ix_messages_id", table_name="messages", schema=schema)
if index_exists("messages", "ix_messages_public_id", inspector):
op.drop_index("ix_messages_public_id", table_name="messages", schema=schema)
if index_exists("peers", "ix_peers_name", inspector):
op.drop_index("ix_peers_name", table_name="peers", schema=schema)
if index_exists("peers", "idx_peers_workspace_lookup", inspector):
op.drop_index("idx_peers_workspace_lookup", table_name="peers", schema=schema)
if index_exists("workspaces", "ix_workspaces_name", inspector):
op.drop_index("ix_workspaces_name", table_name="workspaces", schema=schema)
# Drop redundant index on active_queue_sessions.work_unit_key (unique constraint auto-creates index)
if index_exists(
"active_queue_sessions", "ix_active_queue_sessions_work_unit_key", inspector
):
op.drop_index(
"ix_active_queue_sessions_work_unit_key",
table_name="active_queue_sessions",
schema=schema,
)
# Drop old queue indexes that will be replaced
if index_exists("queue", "ix_queue_workspace_name_processed", inspector):
op.drop_index(
"ix_queue_workspace_name_processed", table_name="queue", schema=schema
)
# ============================================================
# NEW INDEXES (being created, not renamed)
# ============================================================
# Queue - new indexes
if not index_exists("queue", "ix_queue_processed", inspector):
op.create_index(
"ix_queue_processed", "queue", ["processed"], unique=False, schema=schema
)
if not index_exists("queue", "ix_queue_work_unit_key_processed_id", inspector):
op.create_index(
"ix_queue_work_unit_key_processed_id",
"queue",
["work_unit_key", "processed", "id"],
unique=False,
schema=schema,
)
# Sessions - new index
if not index_exists("sessions", "ix_sessions_workspace_name", inspector):
op.create_index(
"ix_sessions_workspace_name",
"sessions",
["workspace_name"],
unique=False,
schema=schema,
)
# Messages - drop FK that's being removed
if constraint_exists(
"messages", "fk_messages_workspace_name_workspaces", "foreignkey", inspector
):
op.drop_constraint(
"fk_messages_workspace_name_workspaces",
"messages",
schema=schema,
type_="foreignkey",
)
# ============================================================
# FIX FK CASCADE (message_embeddings -> messages)
# ============================================================
# Refresh inspector to get latest state after potential renames
inspector = sa.inspect(conn)
# Drop and recreate FK with ON DELETE CASCADE
# This FK should exist at this point either from:
# 1. Being renamed from old name "message_embeddings_message_id_fkey", or
# 2. Already having the new name "fk_message_embeddings_message_id_messages" (from main branch)
if constraint_exists(
"message_embeddings",
"fk_message_embeddings_message_id_messages",
"foreignkey",
inspector,
):
conn.execute(
text(
f"ALTER TABLE {schema}.message_embeddings DROP CONSTRAINT fk_message_embeddings_message_id_messages"
)
)
conn.execute(
text(
f"""
ALTER TABLE {schema}.message_embeddings
ADD CONSTRAINT fk_message_embeddings_message_id_messages
FOREIGN KEY (message_id) REFERENCES {schema}.messages(public_id)
ON DELETE CASCADE
"""
)
)
def downgrade() -> None:
conn = op.get_bind()
inspector = sa.inspect(conn)
# ============================================================
# REVERSE: FIX FK CASCADE (message_embeddings -> messages)
# ============================================================
# Drop FK with CASCADE and recreate without CASCADE
if constraint_exists(
"message_embeddings",
"fk_message_embeddings_message_id_messages",
"foreignkey",
inspector,
):
conn.execute(
text(
f"ALTER TABLE {schema}.message_embeddings DROP CONSTRAINT fk_message_embeddings_message_id_messages"
)
)
conn.execute(
text(
f"""
ALTER TABLE {schema}.message_embeddings
ADD CONSTRAINT fk_message_embeddings_message_id_messages
FOREIGN KEY (message_id) REFERENCES {schema}.messages(public_id)
"""
)
)
# ============================================================
# REVERSE: NEW INDEXES AND FKS
# ============================================================
# Recreate FK
op.create_foreign_key(
"fk_messages_workspace_name_workspaces",
"messages",
"workspaces",
["workspace_name"],
["name"],
source_schema=schema,
referent_schema=schema,
)
# Drop new indexes
op.drop_index("ix_sessions_workspace_name", table_name="sessions", schema=schema)
op.drop_index("ix_queue_processed", table_name="queue", schema=schema)
# ============================================================
# REVERSE: INDEX DELETIONS (recreate them)
# ============================================================
# Recreate redundant index on active_queue_sessions.work_unit_key
op.create_index(
"ix_active_queue_sessions_work_unit_key",
"active_queue_sessions",
["work_unit_key"],
unique=False,
schema=schema,
)
op.create_index(
"ix_queue_workspace_name_processed",
"queue",
["workspace_name", "processed"],
unique=False,
schema=schema,
)
op.create_index(
"ix_workspaces_name", "workspaces", ["name"], unique=False, schema=schema
)
op.create_index(
"idx_peers_workspace_lookup",
"peers",
["workspace_name", "name"],
unique=False,
schema=schema,
)
op.create_index("ix_peers_name", "peers", ["name"], unique=False, schema=schema)
op.create_index(
"ix_messages_public_id", "messages", ["public_id"], unique=False, schema=schema
)
op.create_index("ix_messages_id", "messages", ["id"], unique=False, schema=schema)
# ============================================================
# REVERSE: INDEX RENAMES (using raw SQL)
# ============================================================
conn.execute(
text(
f'ALTER INDEX {schema}."ix_webhook_endpoints_workspace_name" RENAME TO "idx_webhook_endpoints_workspace_lookup"'
)
)
conn.execute(
text(
f'ALTER INDEX {schema}."ix_messages_session_lookup" RENAME TO "idx_messages_session_lookup"'
)
)
conn.execute(
text(
f'ALTER INDEX {schema}."ix_messages_content_gin" RENAME TO "idx_messages_content_gin"'
)
)
conn.execute(
text(
f'ALTER INDEX {schema}."ix_message_embeddings_workspace_name" RENAME TO "idx_message_embeddings_workspace_name"'
)
)
conn.execute(
text(
f'ALTER INDEX {schema}."ix_message_embeddings_session_name" RENAME TO "idx_message_embeddings_session_name"'
)
)
conn.execute(
text(
f'ALTER INDEX {schema}."ix_message_embeddings_peer_name" RENAME TO "idx_message_embeddings_peer_name"'
)
)
conn.execute(
text(
f'ALTER INDEX {schema}."ix_message_embeddings_message_id" RENAME TO "idx_message_embeddings_message_id"'
)
)
conn.execute(
text(
f'ALTER INDEX {schema}."ix_message_embeddings_embedding_hnsw" RENAME TO "idx_message_embeddings_embedding_hnsw"'
)
)
conn.execute(
text(
f'ALTER INDEX {schema}."ix_message_embeddings_created_at" RENAME TO "idx_message_embeddings_created_at"'
)
)
conn.execute(
text(
f'ALTER INDEX {schema}."ix_documents_session_name" RENAME TO "idx_documents_session_name"'
)
)
conn.execute(
text(
f'ALTER INDEX {schema}."ix_documents_observer" RENAME TO "idx_documents_observer"'
)
)
conn.execute(
text(
f'ALTER INDEX {schema}."ix_documents_observed" RENAME TO "idx_documents_observed"'
)
)
conn.execute(
text(
f'ALTER INDEX {schema}."ix_documents_embedding_hnsw" RENAME TO "idx_documents_embedding_hnsw"'
)
)
conn.execute(
text(
f'ALTER INDEX {schema}."ix_collections_observer" RENAME TO "idx_collections_observer"'
)
)
conn.execute(
text(
f'ALTER INDEX {schema}."ix_collections_observed" RENAME TO "idx_collections_observed"'
)
)
# ============================================================
# REVERSE: CONSTRAINT RENAMES
# ============================================================
# Foreign Keys
conn.execute(
text(
f'ALTER TABLE {schema}.message_embeddings RENAME CONSTRAINT "fk_message_embeddings_message_id_messages" TO "message_embeddings_message_id_fkey"'
)
)
# Unique Constraints
conn.execute(
text(
f'ALTER TABLE {schema}.workspaces RENAME CONSTRAINT "uq_workspaces_name" TO "uq_apps_name"'
)
)
conn.execute(
text(
f'ALTER TABLE {schema}.messages RENAME CONSTRAINT "uq_messages_workspace_name_session_name_seq_in_session" TO "uq_messages_session_seq"'
)
)
conn.execute(
text(
f'ALTER TABLE {schema}.sessions RENAME CONSTRAINT "uq_sessions_name_workspace_name" TO "unique_session_name"'
)
)
conn.execute(
text(
f'ALTER TABLE {schema}.peers RENAME CONSTRAINT "uq_peers_name_workspace_name" TO "unique_name_workspace_peer"'
)
)
conn.execute(
text(
f'ALTER TABLE {schema}.collections RENAME CONSTRAINT "uq_collections_observer_observed_workspace_name" TO "unique_observer_observed_collection"'
)
)
conn.execute(
text(
f'ALTER TABLE {schema}.active_queue_sessions RENAME CONSTRAINT "uq_active_queue_sessions_work_unit_key" TO "unique_work_unit_key"'
)
)

11
src/db.py
View File

@ -45,10 +45,19 @@ SessionLocal = async_sessionmaker(
bind=engine,
)
# Define your naming convention
convention = {
"ix": "ix_%(table_name)s_%(column_0_N_name)s", # Index - supports multi-column
"uq": "uq_%(table_name)s_%(column_0_N_name)s", # Unique constraint - supports multi-column
"ck": "ck_%(table_name)s_%(constraint_name)s", # Check constraint
"fk": "fk_%(table_name)s_%(column_0_N_name)s_%(referred_table_name)s", # Foreign key - supports composite keys
"pk": "pk_%(table_name)s", # Primary key
}
table_schema = settings.DB.SCHEMA
# Note: column_0_N_name expands to include all columns in multi-column constraints
# e.g., "workspace_id_tenant_id" for a composite constraint on both columns
meta = MetaData()
meta = MetaData(naming_convention=convention)
meta.schema = table_schema
Base = declarative_base(metadata=meta)

132
src/models.py
View File

@ -96,27 +96,31 @@ 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)
peers = relationship("Peer", back_populates="workspace")
webhook_endpoints = relationship("WebhookEndpoint", back_populates="workspace")
created_at: Mapped[datetime.datetime] = mapped_column(
DateTime(timezone=True), server_default=func.now()
DateTime(timezone=True), server_default=func.now(), index=True
)
h_metadata: Mapped[dict[str, Any]] = mapped_column(
"metadata", JSONB, default=dict, server_default=text("'{}'::jsonb")
)
internal_metadata: Mapped[dict[str, Any]] = mapped_column(
"internal_metadata", JSONB, default=dict
"internal_metadata", JSONB, default=dict, server_default=text("'{}'::jsonb")
)
configuration: Mapped[dict[str, Any]] = mapped_column(
JSONB, default=dict, server_default=text("'{}'::jsonb")
)
sessions = relationship(
"Session", back_populates="workspace", cascade="all, delete, delete-orphan"
)
peers = relationship(
"Peer", back_populates="workspace", cascade="all, delete, delete-orphan"
)
webhook_endpoints = relationship("WebhookEndpoint", back_populates="workspace")
__table_args__ = (
CheckConstraint("length(id) = 21", name="id_length"),
CheckConstraint("length(name) <= 512", name="name_length"),
CheckConstraint("id ~ '^[A-Za-z0-9_-]+$'", name="id_format"),
Index("ix_workspaces_created_at", "created_at"),
Index("ix_workspaces_name", "name"),
)
@ -132,10 +136,10 @@ class Peer(Base):
"internal_metadata", JSONB, default=dict, server_default=text("'{}'::jsonb")
)
created_at: Mapped[datetime.datetime] = mapped_column(
DateTime(timezone=True), server_default=func.now()
DateTime(timezone=True), server_default=func.now(), index=True
)
workspace_name: Mapped[str] = mapped_column(
ForeignKey("workspaces.name"), nullable=False
ForeignKey("workspaces.name"), nullable=False, index=True
)
configuration: Mapped[dict[str, Any]] = mapped_column(
JSONB, default=dict, server_default=text("'{}'::jsonb")
@ -147,14 +151,10 @@ class Peer(Base):
)
__table_args__ = (
UniqueConstraint("name", "workspace_name", name="unique_name_workspace_peer"),
UniqueConstraint("name", "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"),
Index("idx_peers_workspace_lookup", "workspace_name", "name"),
Index("ix_peers_created_at", "created_at"),
Index("ix_peers_name", "name"),
Index("ix_peers_workspace_name", "workspace_name"),
)
def __repr__(self) -> str:
@ -174,26 +174,26 @@ class Session(Base):
"internal_metadata", JSONB, default=dict, server_default=text("'{}'::jsonb")
)
created_at: Mapped[datetime.datetime] = mapped_column(
DateTime(timezone=True), server_default=func.now()
DateTime(timezone=True), server_default=func.now(), index=True
)
messages = relationship("Message", back_populates="session")
workspace_name: Mapped[str] = mapped_column(
ForeignKey("workspaces.name"), nullable=False
ForeignKey("workspaces.name"), nullable=False, index=True
)
configuration: Mapped[dict[str, Any]] = mapped_column(
JSONB, default=dict, server_default=text("'{}'::jsonb")
)
workspace = relationship("Workspace", back_populates="sessions")
peers = relationship(
"Peer", secondary=session_peers_table, back_populates="sessions"
)
messages = relationship("Message", back_populates="session")
__table_args__ = (
UniqueConstraint("name", "workspace_name", name="unique_session_name"),
UniqueConstraint("name", "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"),
Index("ix_sessions_created_at", "created_at"),
)
def __repr__(self) -> str:
@ -225,13 +225,13 @@ class Message(Base):
seq_in_session: Mapped[int] = mapped_column(BigInteger, nullable=False)
created_at: Mapped[datetime.datetime] = mapped_column(
DateTime(timezone=True), server_default=func.now()
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)
session = relationship("Session", back_populates="messages")
peer_name: Mapped[str] = mapped_column(TEXT)
workspace_name: Mapped[str] = mapped_column(
ForeignKey("workspaces.name"),
)
__table_args__ = (
CheckConstraint("length(public_id) = 21", name="public_id_length"),
@ -248,7 +248,7 @@ class Message(Base):
["peers.name", "peers.workspace_name"],
),
Index(
"idx_messages_session_lookup",
"ix_messages_session_lookup",
"session_name",
"id",
postgresql_include=["id", "created_at"],
@ -257,19 +257,13 @@ class Message(Base):
"workspace_name",
"session_name",
"seq_in_session",
name="uq_messages_session_seq",
),
# Full text search index on content column
Index(
"idx_messages_content_gin",
"ix_messages_content_gin",
text("to_tsvector('english', content)"),
postgresql_using="gin",
),
Index("ix_messages_created_at", "created_at"),
Index("ix_messages_id", "id"),
Index("ix_messages_peer_name", "peer_name"),
Index("ix_messages_public_id", "public_id"),
Index("ix_messages_workspace_name", "workspace_name"),
)
@override
@ -287,20 +281,17 @@ class MessageEmbedding(Base):
content: Mapped[str] = mapped_column(TEXT)
embedding: MappedColumn[Any] = mapped_column(Vector(1536))
message_id: Mapped[str] = mapped_column(
ForeignKey("messages.public_id"),
ForeignKey("messages.public_id", ondelete="CASCADE"), nullable=False, index=True
)
workspace_name: Mapped[str] = mapped_column(
ForeignKey("workspaces.name"),
ForeignKey("workspaces.name"), nullable=False, index=True
)
session_name: Mapped[str] = mapped_column(TEXT, nullable=False)
peer_name: Mapped[str] = mapped_column(TEXT)
session_name: Mapped[str] = mapped_column(TEXT, nullable=False, index=True)
peer_name: Mapped[str] = mapped_column(TEXT, nullable=False, index=True)
created_at: Mapped[datetime.datetime] = mapped_column(
DateTime(timezone=True), server_default=func.now()
DateTime(timezone=True), server_default=func.now(), index=True
)
# Relationship to Message
message = relationship("Message", backref="embeddings")
__table_args__ = (
# Compound foreign key constraints
ForeignKeyConstraint(
@ -313,17 +304,12 @@ class MessageEmbedding(Base):
),
# HNSW index on embedding column for efficient similarity search
Index(
"idx_message_embeddings_embedding_hnsw",
"ix_message_embeddings_embedding_hnsw",
"embedding",
postgresql_using="hnsw",
postgresql_with={"m": 16, "ef_construction": 64},
postgresql_ops={"embedding": "vector_cosine_ops"},
),
Index("idx_message_embeddings_created_at", "created_at"),
Index("idx_message_embeddings_message_id", "message_id"),
Index("idx_message_embeddings_peer_name", "peer_name"),
Index("idx_message_embeddings_session_name", "session_name"),
Index("idx_message_embeddings_workspace_name", "workspace_name"),
)
@ -332,10 +318,10 @@ class Collection(Base):
__tablename__: str = "collections"
id: Mapped[str] = mapped_column(TEXT, default=generate_nanoid, primary_key=True)
observer: Mapped[str] = mapped_column(TEXT)
observed: Mapped[str] = mapped_column(TEXT)
observer: Mapped[str] = mapped_column(TEXT, index=True)
observed: Mapped[str] = mapped_column(TEXT, index=True)
created_at: Mapped[datetime.datetime] = mapped_column(
DateTime(timezone=True), server_default=func.now()
DateTime(timezone=True), server_default=func.now(), index=True
)
h_metadata: Mapped[dict[str, Any]] = mapped_column(
"metadata", JSONB, default=dict, server_default=text("'{}'::jsonb")
@ -347,7 +333,7 @@ class Collection(Base):
"Document", back_populates="collection", cascade="all, delete, delete-orphan"
)
workspace_name: Mapped[str] = mapped_column(
ForeignKey("workspaces.name"),
ForeignKey("workspaces.name"), nullable=False, index=True
)
__table_args__ = (
@ -355,7 +341,6 @@ class Collection(Base):
"observer",
"observed",
"workspace_name",
name="unique_observer_observed_collection",
),
CheckConstraint("length(id) = 21", name="id_length"),
CheckConstraint("id ~ '^[A-Za-z0-9_-]+$'", name="id_format"),
@ -369,10 +354,6 @@ class Collection(Base):
["observed", "workspace_name"],
["peers.name", "peers.workspace_name"],
),
Index("idx_collections_observer", "observer"),
Index("idx_collections_observed", "observed"),
Index("ix_collections_created_at", "created_at"),
Index("ix_collections_workspace_name", "workspace_name"),
)
@ -392,13 +373,15 @@ class Document(Base):
)
embedding: MappedColumn[Any] = mapped_column(Vector(1536))
created_at: Mapped[datetime.datetime] = mapped_column(
DateTime(timezone=True), server_default=func.now()
DateTime(timezone=True), server_default=func.now(), index=True
)
observer: Mapped[str] = mapped_column(TEXT)
observed: Mapped[str] = mapped_column(TEXT)
workspace_name: Mapped[str] = mapped_column(ForeignKey("workspaces.name"))
session_name: Mapped[str] = mapped_column(TEXT)
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] = mapped_column(TEXT, index=True)
collection = relationship("Collection", back_populates="documents")
__table_args__ = (
@ -431,7 +414,7 @@ class Document(Base):
),
# HNSW index on embedding column
Index(
"idx_documents_embedding_hnsw",
"ix_documents_embedding_hnsw",
"embedding",
postgresql_using="hnsw", # HNSW index type
postgresql_with={"m": 16, "ef_construction": 64}, # HNSW parameters
@ -439,11 +422,6 @@ class Document(Base):
"embedding": "vector_cosine_ops"
}, # Cosine distance operator
),
Index("idx_documents_observer", "observer"),
Index("idx_documents_observed", "observed"),
Index("idx_documents_session_name", "session_name"),
Index("ix_documents_created_at", "created_at"),
Index("ix_documents_workspace_name", "workspace_name"),
)
@ -453,38 +431,33 @@ class QueueItem(Base):
id: Mapped[int] = mapped_column(
BigInteger, Identity(), primary_key=True, autoincrement=True
)
session_id: Mapped[str] = mapped_column(ForeignKey("sessions.id"), nullable=True)
session_id: Mapped[str | None] = mapped_column(
ForeignKey("sessions.id"), nullable=True, index=True
)
work_unit_key: Mapped[str] = mapped_column(TEXT, nullable=False)
task_type: Mapped[TaskType] = mapped_column(TEXT, nullable=False)
payload: Mapped[dict[str, Any]] = mapped_column(JSONB, nullable=False)
processed: Mapped[bool] = mapped_column(
Boolean, default=False, server_default=text("false")
Boolean, default=False, server_default=text("false"), index=True
)
error: Mapped[str | None] = mapped_column(TEXT, nullable=True)
created_at: Mapped[datetime.datetime] = mapped_column(
DateTime(timezone=True), server_default=func.now()
DateTime(timezone=True), server_default=func.now(), index=True
)
workspace_name: Mapped[str] = mapped_column(
ForeignKey("workspaces.name"), nullable=False
ForeignKey("workspaces.name"), nullable=False, index=True
)
message_id: Mapped[int | None] = mapped_column(
BigInteger, ForeignKey("messages.id"), nullable=True
)
__table_args__ = (
Index("ix_queue_created_at", "created_at"),
Index("ix_queue_session_id", "session_id"),
Index(
"ix_queue_workspace_name",
"workspace_name",
),
Index(
"ix_queue_message_id_not_null",
"message_id",
postgresql_where=text("message_id IS NOT NULL"),
),
Index("ix_queue_workspace_name_processed", "workspace_name", "processed"),
Index(
"ix_queue_work_unit_key_processed_id",
"work_unit_key",
@ -515,7 +488,7 @@ 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
ForeignKey("workspaces.name"), nullable=False, index=True
)
url: Mapped[str] = mapped_column(TEXT, nullable=False)
created_at: Mapped[datetime.datetime] = mapped_column(
@ -524,10 +497,7 @@ class WebhookEndpoint(Base):
workspace = relationship("Workspace", back_populates="webhook_endpoints")
__table_args__ = (
CheckConstraint("length(url) <= 2048", name="webhook_endpoint_url_length"),
Index("idx_webhook_endpoints_workspace_lookup", "workspace_name"),
)
__table_args__ = (CheckConstraint("length(url) <= 2048", name="url_length"),)
def __repr__(self) -> str:
return f"WebhookEndpoint(id={self.id}, workspace_name={self.workspace_name}, url={self.url})"

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

@ -15,6 +15,7 @@ from . import (
test_a1b2c3d4e5f6_initial_schema,
test_b765d82110bd_change_metamessages_to_user_level_with_,
test_b8183c5ffb48_codify_document_level_and_times_derived,
test_baa22cad81e2_standardize_constraint_names,
test_bb6fb3a7a643_add_message_seq_in_session_column,
test_c3828084f472_add_indexes_for_messages_and_,
test_d429de0e5338_adopt_peer_paradigm,
@ -37,6 +38,7 @@ __all__ = [
"test_a1b2c3d4e5f6_initial_schema",
"test_b765d82110bd_change_metamessages_to_user_level_with_",
"test_b8183c5ffb48_codify_document_level_and_times_derived",
"test_baa22cad81e2_standardize_constraint_names",
"test_bb6fb3a7a643_add_message_seq_in_session_column",
"test_c3828084f472_add_indexes_for_messages_and_",
"test_d429de0e5338_adopt_peer_paradigm",

9
tests/alembic/revisions/test_20f89a421aff_rename_metamessage_type_to_label.py
View File

@ -24,9 +24,12 @@ def prepare_metamessage_label(verifier: MigrationVerifier) -> None:
("metamessages", "idx_metamessages_message_lookup"),
)
verifier.assert_indexes_exist(OLD_INDEXES)
verifier.assert_constraint_exists(
"metamessages", "metamessage_type_length", "check"
)
# Check for either naming convention or plain name (backward compatibility)
constraints = verifier.fetch_constraints("metamessages", "check")
assert (
"ck_metamessages_metamessage_type_length" in constraints
or "metamessage_type_length" in constraints
), f"Expected constraint not found. Available: {constraints}"
schema = verifier.schema
connection = verifier.conn

169
tests/alembic/revisions/test_baa22cad81e2_standardize_constraint_names.py Normal file
View File

@ -0,0 +1,169 @@
"""Hooks for revision baa22cad81e2 (standardize_constraint_names)."""
from __future__ import annotations
from sqlalchemy import text
from tests.alembic.registry import register_after_upgrade, register_before_upgrade
from tests.alembic.verifier import MigrationVerifier
@register_before_upgrade("baa22cad81e2")
def prepare_standardize_constraint_names(verifier: MigrationVerifier) -> None:
"""Seed state and assertions before upgrading to baa22cad81e2."""
# Verify old-style unique constraint names exist
verifier.assert_constraint_exists(
"active_queue_sessions", "unique_work_unit_key", "unique"
)
verifier.assert_constraint_exists(
"collections", "unique_observer_observed_collection", "unique"
)
verifier.assert_constraint_exists("peers", "unique_name_workspace_peer", "unique")
verifier.assert_constraint_exists("sessions", "unique_session_name", "unique")
verifier.assert_constraint_exists("messages", "uq_messages_session_seq", "unique")
verifier.assert_constraint_exists("workspaces", "uq_apps_name", "unique")
# Note: FK name checking is flexible because the FK might have been created
# with different names in different migration paths. The important thing is
# that after the migration it has the correct name with CASCADE.
# Verify old-style index names exist
old_indexes = [
("collections", "idx_collections_observed"),
("collections", "idx_collections_observer"),
("documents", "idx_documents_embedding_hnsw"),
("documents", "idx_documents_observed"),
("documents", "idx_documents_observer"),
("documents", "idx_documents_session_name"),
("message_embeddings", "idx_message_embeddings_created_at"),
("message_embeddings", "idx_message_embeddings_embedding_hnsw"),
("message_embeddings", "idx_message_embeddings_message_id"),
("message_embeddings", "idx_message_embeddings_peer_name"),
("message_embeddings", "idx_message_embeddings_session_name"),
("message_embeddings", "idx_message_embeddings_workspace_name"),
("messages", "idx_messages_content_gin"),
("messages", "idx_messages_session_lookup"),
("webhook_endpoints", "idx_webhook_endpoints_workspace_lookup"),
]
verifier.assert_indexes_exist(old_indexes)
# Note: We don't check for redundant indexes in before_upgrade because
# they may or may not exist depending on the migration path. The important
# thing is that after the migration, they should not exist.
@register_after_upgrade("baa22cad81e2")
def verify_standardize_constraint_names(verifier: MigrationVerifier) -> None:
"""Add assertions validating the effects of baa22cad81e2."""
# Verify new-style unique constraint names exist
verifier.assert_constraint_exists(
"active_queue_sessions", "uq_active_queue_sessions_work_unit_key", "unique"
)
verifier.assert_constraint_exists(
"collections", "uq_collections_observer_observed_workspace_name", "unique"
)
verifier.assert_constraint_exists("peers", "uq_peers_name_workspace_name", "unique")
verifier.assert_constraint_exists(
"sessions", "uq_sessions_name_workspace_name", "unique"
)
verifier.assert_constraint_exists(
"messages", "uq_messages_workspace_name_session_name_seq_in_session", "unique"
)
verifier.assert_constraint_exists("workspaces", "uq_workspaces_name", "unique")
# Verify old-style constraint names no longer exist
verifier.assert_constraint_exists(
"active_queue_sessions", "unique_work_unit_key", "unique", exists=False
)
verifier.assert_constraint_exists(
"collections", "unique_observer_observed_collection", "unique", exists=False
)
verifier.assert_constraint_exists(
"peers", "unique_name_workspace_peer", "unique", exists=False
)
verifier.assert_constraint_exists(
"sessions", "unique_session_name", "unique", exists=False
)
verifier.assert_constraint_exists(
"messages", "uq_messages_session_seq", "unique", exists=False
)
verifier.assert_constraint_exists(
"workspaces", "uq_apps_name", "unique", exists=False
)
# Verify new-style foreign key names exist with correct name
verifier.assert_constraint_exists(
"message_embeddings", "fk_message_embeddings_message_id_messages", "foreign_key"
)
# Verify new-style index names exist
new_indexes = [
("collections", "ix_collections_observed"),
("collections", "ix_collections_observer"),
("documents", "ix_documents_embedding_hnsw"),
("documents", "ix_documents_observed"),
("documents", "ix_documents_observer"),
("documents", "ix_documents_session_name"),
("message_embeddings", "ix_message_embeddings_created_at"),
("message_embeddings", "ix_message_embeddings_embedding_hnsw"),
("message_embeddings", "ix_message_embeddings_message_id"),
("message_embeddings", "ix_message_embeddings_peer_name"),
("message_embeddings", "ix_message_embeddings_session_name"),
("message_embeddings", "ix_message_embeddings_workspace_name"),
("messages", "ix_messages_content_gin"),
("messages", "ix_messages_session_lookup"),
("webhook_endpoints", "ix_webhook_endpoints_workspace_name"),
]
verifier.assert_indexes_exist(new_indexes)
# Verify key old-style index names no longer exist
# Note: Only checking a subset because some indexes may not exist in all migration paths
key_old_indexes = [
("collections", "idx_collections_observed"),
("documents", "idx_documents_observed"),
("message_embeddings", "idx_message_embeddings_created_at"),
]
verifier.assert_indexes_not_exist(key_old_indexes)
# Verify redundant indexes were removed
redundant_indexes = [
("messages", "ix_messages_id"),
("messages", "ix_messages_public_id"),
("peers", "ix_peers_name"),
("peers", "idx_peers_workspace_lookup"),
("workspaces", "ix_workspaces_name"),
("active_queue_sessions", "ix_active_queue_sessions_work_unit_key"),
("queue", "ix_queue_workspace_name_processed"),
]
verifier.assert_indexes_not_exist(redundant_indexes)
# Verify new indexes were created
new_queue_indexes = [
("queue", "ix_queue_processed"),
("queue", "ix_queue_work_unit_key_processed_id"),
("sessions", "ix_sessions_workspace_name"),
]
verifier.assert_indexes_exist(new_queue_indexes)
# Verify FK was removed
verifier.assert_constraint_exists(
"messages", "fk_messages_workspace_name_workspaces", "foreign_key", exists=False
)
# Verify FK CASCADE was added to message_embeddings.message_id
# Check using raw SQL since the verifier doesn't have a method for this
result = verifier.conn.execute(
text(
f"""
SELECT confdeltype
FROM pg_constraint
WHERE conrelid = '{verifier.schema}.message_embeddings'::regclass
AND conname = 'fk_message_embeddings_message_id_messages'
"""
)
)
row = result.fetchone()
assert row is not None, "FK fk_message_embeddings_message_id_messages not found"
assert (
row[0] == "c"
), f"FK should have ON DELETE CASCADE (confdeltype='c'), got '{row[0]}'"