align DB state to sqlalchemy model definitions (#245)
* feat: align DB schema with sqlalchemy model definitions * test: add test for new migration file * fix: add pk for message embeddings table * fix: add naming constraint * fix: align more indexes + unique constraints --> rm redundancy * fix: add missing indexes * Fix Server Default Migrations and add in appropriate Server Defaults (#252) * fix: Improve alembic migration reliability * fix: Ensure alembic migrations use a session pooler (#247) * fix: Ensure alembic migrations use a session pooler * fix: bump batch size; fix document delete in 08894082221a (#249) * chore: add alembic logging * fix (alembic): remove order by in batches * fix: fkey -> fk * chore: code rabbit --------- Co-authored-by: Rajat Ahuja <rahuja445@gmail.com> * fix: Add server defaults to appropriate columns * chore: (tests) Add alembic migration test --------- Co-authored-by: Rajat Ahuja <rahuja445@gmail.com> * fix: align models with alembic check --------- Co-authored-by: Vineeth Voruganti <13438633+VVoruganti@users.noreply.github.com>
This commit is contained in:
Rajat Ahuja
GitHub
parent
1df47e61c8
commit
097f3b31a0
|
|
@ -1,4 +1,4 @@
|
|||
import logging
|
||||
import logging # noqa: I001
|
||||
import sys
|
||||
from logging.config import fileConfig
|
||||
from pathlib import Path
|
||||
|
|
@ -12,6 +12,10 @@ from src.config import settings
|
|||
# Import your models
|
||||
from src.db import Base
|
||||
|
||||
# Import all models so they register with Base.metadata
|
||||
import src.models # noqa: F401
|
||||
|
||||
|
||||
# Set up logging more verbosely
|
||||
logging.basicConfig()
|
||||
logging.getLogger("sqlalchemy.engine").setLevel(logging.INFO)
|
||||
|
|
@ -166,6 +170,13 @@ def run_migrations_online() -> None:
|
|||
connection=connection,
|
||||
target_metadata=target_metadata,
|
||||
version_table_schema=target_metadata.schema,
|
||||
include_schemas=True,
|
||||
include_object=lambda obj, name, type_, reflected, compare_to: (
|
||||
# Only include objects from our target schema
|
||||
getattr(obj, "schema", None) == target_metadata.schema
|
||||
if hasattr(obj, "schema")
|
||||
else True
|
||||
),
|
||||
)
|
||||
|
||||
with context.begin_transaction():
|
||||
|
|
|
|||
|
|
@ -5,18 +5,20 @@ Revises: ${down_revision | comma,n}
|
|||
Create Date: ${create_date}
|
||||
|
||||
"""
|
||||
from typing import Sequence, Union
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import op
|
||||
import sqlalchemy as sa
|
||||
${imports if imports else ""}
|
||||
from migrations.utils import get_schema
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = ${repr(up_revision)}
|
||||
down_revision: Union[str, None] = ${repr(down_revision)}
|
||||
branch_labels: Union[str, Sequence[str], None] = ${repr(branch_labels)}
|
||||
depends_on: Union[str, Sequence[str], None] = ${repr(depends_on)}
|
||||
down_revision: str | None = ${repr(down_revision)}
|
||||
branch_labels: str | Sequence[str] | None = ${repr(branch_labels)}
|
||||
depends_on: str | Sequence[str] | None = ${repr(depends_on)}
|
||||
|
||||
schema = get_schema()
|
||||
|
||||
def upgrade() -> None:
|
||||
${upgrades if upgrades else "pass"}
|
||||
|
|
|
|||
|
|
@ -71,3 +71,89 @@ def constraint_exists(
|
|||
else:
|
||||
raise ValueError(f"Invalid constraint type: {type}")
|
||||
return any(constraint["name"] == constraint_name for constraint in constraints)
|
||||
|
||||
|
||||
def make_column_non_nullable_safe(table_name: str, column_name: str) -> None:
|
||||
"""
|
||||
Make a column non-nullable using a non-blocking approach to minimize lock duration.
|
||||
|
||||
WARNING: Only use this if you can guarantee that:
|
||||
1. No NULL values currently exist in the column
|
||||
2. The application code is already writing non-NULL values to this column or
|
||||
3. The column has never accepted NULLs in practice
|
||||
|
||||
This uses a 4-step process to avoid long exclusive locks:
|
||||
1. Add CHECK constraint with NOT VALID (instant, no scan)
|
||||
2. Validate the constraint (scans but allows concurrent read/writes to the table)
|
||||
3. Set column NOT NULL (fast since we've validated the constraint)
|
||||
4. Drop the redundant CHECK constraint
|
||||
|
||||
Args:
|
||||
table_name: The name of the table
|
||||
column_name: The name of the column to make non-nullable
|
||||
"""
|
||||
schema = get_schema()
|
||||
conn = op.get_bind()
|
||||
constraint_name = f"{table_name}_{column_name}_not_null"
|
||||
|
||||
# Step 1: Check if the column is already non-nullable
|
||||
inspector = sa.inspect(op.get_bind())
|
||||
columns = inspector.get_columns(table_name, schema=schema)
|
||||
column_info = next((col for col in columns if col["name"] == column_name), None)
|
||||
if column_info is None:
|
||||
raise ValueError(f"Column {table_name}.{column_name} does not exist")
|
||||
if not column_info["nullable"]:
|
||||
print(f"Column {table_name}.{column_name} is already non-nullable, skipping...")
|
||||
return
|
||||
|
||||
# Step 2: Add CHECK constraint without validation (instant)
|
||||
# Note: op.create_check_constraint() doesn't support NOT VALID, so use raw SQL
|
||||
|
||||
# Get the identifier preparer for safe quoting
|
||||
dialect = conn.dialect
|
||||
preparer = dialect.identifier_preparer
|
||||
|
||||
quoted_schema = preparer.quote(schema)
|
||||
quoted_table = preparer.quote(table_name)
|
||||
quoted_constraint = preparer.quote(constraint_name)
|
||||
quoted_column = preparer.quote(column_name)
|
||||
|
||||
# Step 2: Add CHECK constraint without validation (instant)
|
||||
# Note: op.create_check_constraint() doesn't support NOT VALID, so use raw SQL
|
||||
if not constraint_exists(table_name, constraint_name, "check"):
|
||||
conn.execute(
|
||||
sa.text(
|
||||
f"""
|
||||
ALTER TABLE {quoted_schema}.{quoted_table}
|
||||
ADD CONSTRAINT {quoted_constraint}
|
||||
CHECK ({quoted_column} IS NOT NULL)
|
||||
NOT VALID
|
||||
"""
|
||||
)
|
||||
)
|
||||
|
||||
# Step 3: Validate constraint (scans but allows concurrent operations)
|
||||
conn.execute(
|
||||
sa.text(
|
||||
f"""
|
||||
ALTER TABLE {quoted_schema}.{quoted_table}
|
||||
VALIDATE CONSTRAINT {quoted_constraint}
|
||||
"""
|
||||
)
|
||||
)
|
||||
|
||||
# Step 4: Set NOT NULL (fast with validated constraint)
|
||||
op.alter_column(
|
||||
table_name,
|
||||
column_name,
|
||||
nullable=False,
|
||||
schema=schema,
|
||||
)
|
||||
|
||||
# Step 5: Drop the redundant CHECK constraint
|
||||
op.drop_constraint(
|
||||
constraint_name,
|
||||
table_name,
|
||||
type_="check",
|
||||
schema=schema,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -0,0 +1,341 @@
|
|||
"""align_schema_with_declarative_models
|
||||
|
||||
Revision ID: 066e87ca5b07
|
||||
Revises: bb6fb3a7a643
|
||||
Create Date: 2025-10-27 12:36:51.614959
|
||||
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
|
||||
from migrations.utils import (
|
||||
column_exists,
|
||||
constraint_exists,
|
||||
fk_exists,
|
||||
get_schema,
|
||||
index_exists,
|
||||
make_column_non_nullable_safe,
|
||||
)
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = "066e87ca5b07"
|
||||
down_revision: str | None = "bb6fb3a7a643"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
schema = get_schema()
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
"""
|
||||
The application code has been previously updated to ensure none of the following columns have NULL values but the actual DB schema is out of sync with our SQLAlchemy model definitions.
|
||||
This migration fixes this by making the columns non-nullable using a non-blocking approach to minimize lock duration.
|
||||
"""
|
||||
conn = op.get_bind()
|
||||
|
||||
# Make peers.workspace_name non-nullable
|
||||
if column_exists("peers", "workspace_name"):
|
||||
make_column_non_nullable_safe("peers", "workspace_name")
|
||||
|
||||
# Make sessions.workspace_name non-nullable
|
||||
if column_exists("sessions", "workspace_name"):
|
||||
make_column_non_nullable_safe("sessions", "workspace_name")
|
||||
|
||||
# Make active_queue_sessions.work_unit_key non-nullable
|
||||
if column_exists("active_queue_sessions", "work_unit_key"):
|
||||
make_column_non_nullable_safe("active_queue_sessions", "work_unit_key")
|
||||
|
||||
# Make documents.embedding non-nullable
|
||||
if column_exists("documents", "embedding"):
|
||||
make_column_non_nullable_safe("documents", "embedding")
|
||||
|
||||
# Add primary key constraint to message_embeddings.id
|
||||
if column_exists("message_embeddings", "id") and not constraint_exists(
|
||||
"message_embeddings", "pk_message_embeddings", "primary"
|
||||
):
|
||||
conn.execute(
|
||||
sa.text(
|
||||
f"""
|
||||
ALTER TABLE {schema}.message_embeddings
|
||||
ADD CONSTRAINT pk_message_embeddings
|
||||
PRIMARY KEY (id)
|
||||
"""
|
||||
)
|
||||
)
|
||||
|
||||
# Rename indexes on peers table
|
||||
inspector = sa.inspect(conn)
|
||||
index_renames = [
|
||||
("peers", "ix_users_created_at", "ix_peers_created_at"),
|
||||
("peers", "ix_users_name", "ix_peers_name"),
|
||||
("workspaces", "ix_apps_created_at", "ix_workspaces_created_at"),
|
||||
("workspaces", "ix_apps_name", "ix_workspaces_name"),
|
||||
]
|
||||
for table_name, old_name, new_name in index_renames:
|
||||
if index_exists(table_name, old_name, inspector):
|
||||
conn.execute(
|
||||
sa.text(f"ALTER INDEX {schema}.{old_name} RENAME TO {new_name}")
|
||||
)
|
||||
|
||||
# Drop redundant indexes
|
||||
if index_exists("workspaces", "ix_apps_public_id", inspector):
|
||||
op.drop_index("ix_apps_public_id", table_name="workspaces", schema=schema)
|
||||
if index_exists("peers", "ix_users_public_id", inspector):
|
||||
op.drop_index("ix_users_public_id", table_name="peers", schema=schema)
|
||||
if index_exists("sessions", "ix_sessions_public_id", inspector):
|
||||
op.drop_index("ix_sessions_public_id", table_name="sessions", schema=schema)
|
||||
if index_exists("documents", "ix_documents_public_id", inspector):
|
||||
op.drop_index("ix_documents_public_id", table_name="documents", schema=schema)
|
||||
if index_exists("collections", "ix_collections_public_id", inspector):
|
||||
op.drop_index(
|
||||
"ix_collections_public_id", table_name="collections", schema=schema
|
||||
)
|
||||
|
||||
# Drop redundant unique constraints
|
||||
if constraint_exists("workspaces", "uq_apps_public_id", "unique", inspector):
|
||||
op.drop_constraint(
|
||||
"uq_apps_public_id", "workspaces", type_="unique", schema=schema
|
||||
)
|
||||
|
||||
if constraint_exists("peers", "uq_users_public_id", "unique", inspector):
|
||||
op.drop_constraint("uq_users_public_id", "peers", type_="unique", schema=schema)
|
||||
|
||||
if constraint_exists("sessions", "uq_sessions_public_id", "unique", inspector):
|
||||
op.drop_constraint(
|
||||
"uq_sessions_public_id", "sessions", type_="unique", schema=schema
|
||||
)
|
||||
|
||||
if constraint_exists(
|
||||
"collections", "uq_collections_public_id", "unique", inspector
|
||||
):
|
||||
op.drop_constraint(
|
||||
"uq_collections_public_id", "collections", type_="unique", schema=schema
|
||||
)
|
||||
|
||||
if constraint_exists("documents", "uq_documents_public_id", "unique", inspector):
|
||||
op.drop_constraint(
|
||||
"uq_documents_public_id", "documents", type_="unique", schema=schema
|
||||
)
|
||||
|
||||
# Drop unnecessary index on active queue
|
||||
if index_exists(
|
||||
"active_queue_sessions",
|
||||
f"ix_{schema}_active_queue_sessions_work_unit_key",
|
||||
inspector,
|
||||
):
|
||||
op.drop_index(
|
||||
f"ix_{schema}_active_queue_sessions_work_unit_key",
|
||||
table_name="active_queue_sessions",
|
||||
schema=schema,
|
||||
)
|
||||
|
||||
# Add FK constraint on queue.session_id to sessions.id
|
||||
if not fk_exists("queue", "fk_queue_session_id"):
|
||||
# Add constraint without validation (fast, doesn't scan)
|
||||
conn.execute(
|
||||
sa.text(
|
||||
f"""
|
||||
ALTER TABLE {schema}.queue
|
||||
ADD CONSTRAINT fk_queue_session_id
|
||||
FOREIGN KEY (session_id)
|
||||
REFERENCES {schema}.sessions(id)
|
||||
NOT VALID
|
||||
"""
|
||||
)
|
||||
)
|
||||
# Validate constraint (scans but allows concurrent reads)
|
||||
conn.execute(
|
||||
sa.text(
|
||||
f"ALTER TABLE {schema}.queue VALIDATE CONSTRAINT fk_queue_session_id"
|
||||
)
|
||||
)
|
||||
|
||||
# Create missing indexes
|
||||
if not index_exists("peers", "ix_peers_workspace_name", inspector):
|
||||
op.create_index(
|
||||
"ix_peers_workspace_name", "peers", ["workspace_name"], schema=schema
|
||||
)
|
||||
|
||||
if not index_exists("collections", "ix_collections_workspace_name", inspector):
|
||||
op.create_index(
|
||||
"ix_collections_workspace_name",
|
||||
"collections",
|
||||
["workspace_name"],
|
||||
schema=schema,
|
||||
)
|
||||
|
||||
if not index_exists("documents", "ix_documents_workspace_name", inspector):
|
||||
op.create_index(
|
||||
"ix_documents_workspace_name",
|
||||
"documents",
|
||||
["workspace_name"],
|
||||
schema=schema,
|
||||
)
|
||||
|
||||
if not fk_exists("session_peers", "fk_session_peers_workspace_name", inspector):
|
||||
# Add constraint without validation (fast, doesn't scan)
|
||||
conn.execute(
|
||||
sa.text(
|
||||
f"""
|
||||
ALTER TABLE {schema}.session_peers
|
||||
ADD CONSTRAINT fk_session_peers_workspace_name
|
||||
FOREIGN KEY (workspace_name)
|
||||
REFERENCES {schema}.workspaces(name)
|
||||
NOT VALID
|
||||
"""
|
||||
)
|
||||
)
|
||||
# Validate constraint (scans but allows concurrent reads)
|
||||
conn.execute(
|
||||
sa.text(
|
||||
f"ALTER TABLE {schema}.session_peers VALIDATE CONSTRAINT fk_session_peers_workspace_name"
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
conn = op.get_bind()
|
||||
inspector = sa.inspect(conn)
|
||||
|
||||
if fk_exists("session_peers", "fk_session_peers_workspace_name", inspector):
|
||||
op.drop_constraint(
|
||||
"fk_session_peers_workspace_name",
|
||||
table_name="session_peers",
|
||||
type_="foreignkey",
|
||||
schema=schema,
|
||||
)
|
||||
|
||||
if index_exists("documents", "ix_documents_workspace_name", inspector):
|
||||
op.drop_index(
|
||||
"ix_documents_workspace_name", table_name="documents", schema=schema
|
||||
)
|
||||
|
||||
if index_exists("peers", "ix_peers_workspace_name", inspector):
|
||||
op.drop_index("ix_peers_workspace_name", table_name="peers", schema=schema)
|
||||
if index_exists("collections", "ix_collections_workspace_name", inspector):
|
||||
op.drop_index(
|
||||
"ix_collections_workspace_name", table_name="collections", schema=schema
|
||||
)
|
||||
|
||||
# First, drop the FK constraint (we'll recreate it later if needed)
|
||||
if fk_exists("queue", "fk_queue_session_id"):
|
||||
op.drop_constraint(
|
||||
"fk_queue_session_id",
|
||||
"queue",
|
||||
type_="foreignkey",
|
||||
schema=schema,
|
||||
)
|
||||
|
||||
if not index_exists(
|
||||
"active_queue_sessions",
|
||||
f"ix_{schema}_active_queue_sessions_work_unit_key",
|
||||
inspector,
|
||||
):
|
||||
op.create_index(
|
||||
f"ix_{schema}_active_queue_sessions_work_unit_key",
|
||||
table_name="active_queue_sessions",
|
||||
columns=["work_unit_key"],
|
||||
schema=schema,
|
||||
)
|
||||
|
||||
# Recreate the redundant unique constraints
|
||||
if not constraint_exists("sessions", "uq_sessions_public_id", "unique", inspector):
|
||||
op.create_unique_constraint(
|
||||
"uq_sessions_public_id", "sessions", ["id"], schema=schema
|
||||
)
|
||||
if not constraint_exists("peers", "uq_users_public_id", "unique", inspector):
|
||||
op.create_unique_constraint(
|
||||
"uq_users_public_id", "peers", ["id"], schema=schema
|
||||
)
|
||||
if not constraint_exists("workspaces", "uq_apps_public_id", "unique", inspector):
|
||||
op.create_unique_constraint(
|
||||
"uq_apps_public_id", "workspaces", ["id"], schema=schema
|
||||
)
|
||||
|
||||
if not constraint_exists(
|
||||
"collections", "uq_collections_public_id", "unique", inspector
|
||||
):
|
||||
op.create_unique_constraint(
|
||||
"uq_collections_public_id", "collections", ["id"], schema=schema
|
||||
)
|
||||
if not constraint_exists(
|
||||
"documents", "uq_documents_public_id", "unique", inspector
|
||||
):
|
||||
op.create_unique_constraint(
|
||||
"uq_documents_public_id", "documents", ["id"], schema=schema
|
||||
)
|
||||
|
||||
# Recreate the redundant indexes
|
||||
if not index_exists("sessions", "ix_sessions_public_id", inspector):
|
||||
op.create_index("ix_sessions_public_id", "sessions", ["id"], schema=schema)
|
||||
if not index_exists("peers", "ix_users_public_id", inspector):
|
||||
op.create_index("ix_users_public_id", "peers", ["id"], schema=schema)
|
||||
if not index_exists("workspaces", "ix_apps_public_id", inspector):
|
||||
op.create_index("ix_apps_public_id", "workspaces", ["id"], schema=schema)
|
||||
if not index_exists("documents", "ix_documents_public_id", inspector):
|
||||
op.create_index("ix_documents_public_id", "documents", ["id"], schema=schema)
|
||||
if not index_exists("collections", "ix_collections_public_id", inspector):
|
||||
op.create_index(
|
||||
"ix_collections_public_id",
|
||||
"collections",
|
||||
["id"],
|
||||
schema=schema,
|
||||
)
|
||||
|
||||
# Rename indexes on peers table back to original names
|
||||
index_renames = [
|
||||
("peers", "ix_peers_created_at", "ix_users_created_at"),
|
||||
("peers", "ix_peers_name", "ix_users_name"),
|
||||
("workspaces", "ix_workspaces_created_at", "ix_apps_created_at"),
|
||||
("workspaces", "ix_workspaces_name", "ix_apps_name"),
|
||||
]
|
||||
for table_name, new_name, old_name in index_renames:
|
||||
if index_exists(table_name, new_name, inspector):
|
||||
conn.execute(
|
||||
sa.text(f"ALTER INDEX {schema}.{new_name} RENAME TO {old_name}")
|
||||
)
|
||||
|
||||
# Drop primary key constraint from message_embeddings.id
|
||||
if constraint_exists("message_embeddings", "pk_message_embeddings", "primary"):
|
||||
op.drop_constraint(
|
||||
"pk_message_embeddings", "message_embeddings", "primary", schema=schema
|
||||
)
|
||||
|
||||
# Make documents.embedding nullable
|
||||
if column_exists("documents", "embedding"):
|
||||
op.alter_column(
|
||||
"documents",
|
||||
"embedding",
|
||||
nullable=True,
|
||||
schema=schema,
|
||||
)
|
||||
|
||||
# Make active_queue_sessions.work_unit_key nullable
|
||||
if column_exists("active_queue_sessions", "work_unit_key"):
|
||||
op.alter_column(
|
||||
"active_queue_sessions",
|
||||
"work_unit_key",
|
||||
nullable=True,
|
||||
schema=schema,
|
||||
)
|
||||
|
||||
# Make sessions.workspace_name nullable
|
||||
if column_exists("sessions", "workspace_name"):
|
||||
op.alter_column(
|
||||
"sessions",
|
||||
"workspace_name",
|
||||
nullable=True,
|
||||
schema=schema,
|
||||
)
|
||||
|
||||
# Make peers.workspace_name nullable
|
||||
if column_exists("peers", "workspace_name"):
|
||||
op.alter_column(
|
||||
"peers",
|
||||
"workspace_name",
|
||||
nullable=True,
|
||||
schema=schema,
|
||||
)
|
||||
|
|
@ -53,7 +53,7 @@ def upgrade() -> None:
|
|||
connection.execute(
|
||||
text(
|
||||
f"""
|
||||
INSERT INTO {schema}.sessions (id, name, workspace_name, is_active) VALUES (:session_id, '__global_observations__', :workspace_name, true) ON CONFLICT DO NOTHING
|
||||
INSERT INTO {schema}.sessions (id, name, workspace_name, is_active, metadata, internal_metadata, configuration, created_at) VALUES (:session_id, '__global_observations__', :workspace_name, true, '{{}}', '{{}}', '{{}}', NOW()) ON CONFLICT DO NOTHING
|
||||
"""
|
||||
),
|
||||
{"session_id": session_id, "workspace_name": workspace_name},
|
||||
|
|
@ -446,7 +446,10 @@ def upgrade() -> None:
|
|||
schema=schema,
|
||||
)
|
||||
|
||||
# Step 17: Drop the name column from collections
|
||||
# Step 17: Drop the name_length check constraint before dropping the name column from collections
|
||||
if constraint_exists("collections", "name_length", "check", inspector):
|
||||
op.drop_constraint("name_length", "collections", schema=schema)
|
||||
|
||||
if column_exists("collections", "name", inspector):
|
||||
op.drop_column("collections", "name", schema=schema)
|
||||
|
||||
|
|
|
|||
|
|
@ -11,6 +11,8 @@ from collections.abc import Sequence
|
|||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
|
||||
from migrations.utils import constraint_exists
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = "20f89a421aff"
|
||||
down_revision: str | None = "556a16564f50"
|
||||
|
|
@ -61,7 +63,8 @@ def upgrade() -> None:
|
|||
)
|
||||
|
||||
# Rename check constraint
|
||||
op.execute("ALTER TABLE metamessages DROP CONSTRAINT metamessage_type_length;")
|
||||
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")
|
||||
# ### end Alembic commands ###
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,387 @@
|
|||
"""add server defaults to timestamp boolean and jsonb columns
|
||||
|
||||
Revision ID: e9b705f9adf9
|
||||
Revises: 066e87ca5b07
|
||||
Create Date: 2025-10-29 12:08:36.803611
|
||||
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
|
||||
from migrations.utils import get_schema
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = "e9b705f9adf9"
|
||||
down_revision: str | None = "066e87ca5b07"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
schema = get_schema()
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
# Add server defaults for timestamp columns
|
||||
op.alter_column(
|
||||
"workspaces",
|
||||
"created_at",
|
||||
server_default=sa.func.now(),
|
||||
schema=schema,
|
||||
)
|
||||
op.alter_column(
|
||||
"peers",
|
||||
"created_at",
|
||||
server_default=sa.func.now(),
|
||||
schema=schema,
|
||||
)
|
||||
op.alter_column(
|
||||
"sessions",
|
||||
"created_at",
|
||||
server_default=sa.func.now(),
|
||||
schema=schema,
|
||||
)
|
||||
op.alter_column(
|
||||
"messages",
|
||||
"created_at",
|
||||
server_default=sa.func.now(),
|
||||
schema=schema,
|
||||
)
|
||||
op.alter_column(
|
||||
"message_embeddings",
|
||||
"created_at",
|
||||
server_default=sa.func.now(),
|
||||
schema=schema,
|
||||
)
|
||||
op.alter_column(
|
||||
"collections",
|
||||
"created_at",
|
||||
server_default=sa.func.now(),
|
||||
schema=schema,
|
||||
)
|
||||
op.alter_column(
|
||||
"documents",
|
||||
"created_at",
|
||||
server_default=sa.func.now(),
|
||||
schema=schema,
|
||||
)
|
||||
op.alter_column(
|
||||
"queue",
|
||||
"created_at",
|
||||
server_default=sa.func.now(),
|
||||
schema=schema,
|
||||
)
|
||||
op.alter_column(
|
||||
"webhook_endpoints",
|
||||
"created_at",
|
||||
server_default=sa.func.now(),
|
||||
schema=schema,
|
||||
)
|
||||
op.alter_column(
|
||||
"session_peers",
|
||||
"joined_at",
|
||||
server_default=sa.func.now(),
|
||||
schema=schema,
|
||||
)
|
||||
op.alter_column(
|
||||
"active_queue_sessions",
|
||||
"last_updated",
|
||||
server_default=sa.func.now(),
|
||||
schema=schema,
|
||||
)
|
||||
|
||||
# Add server defaults for JSONB columns
|
||||
op.alter_column(
|
||||
"workspaces",
|
||||
"metadata",
|
||||
server_default=sa.text("'{}'::jsonb"),
|
||||
schema=schema,
|
||||
)
|
||||
op.alter_column(
|
||||
"workspaces",
|
||||
"internal_metadata",
|
||||
server_default=sa.text("'{}'::jsonb"),
|
||||
schema=schema,
|
||||
)
|
||||
op.alter_column(
|
||||
"workspaces",
|
||||
"configuration",
|
||||
server_default=sa.text("'{}'::jsonb"),
|
||||
schema=schema,
|
||||
)
|
||||
op.alter_column(
|
||||
"peers",
|
||||
"metadata",
|
||||
server_default=sa.text("'{}'::jsonb"),
|
||||
schema=schema,
|
||||
)
|
||||
op.alter_column(
|
||||
"peers",
|
||||
"internal_metadata",
|
||||
server_default=sa.text("'{}'::jsonb"),
|
||||
schema=schema,
|
||||
)
|
||||
op.alter_column(
|
||||
"peers",
|
||||
"configuration",
|
||||
server_default=sa.text("'{}'::jsonb"),
|
||||
schema=schema,
|
||||
)
|
||||
op.alter_column(
|
||||
"sessions",
|
||||
"metadata",
|
||||
server_default=sa.text("'{}'::jsonb"),
|
||||
schema=schema,
|
||||
)
|
||||
op.alter_column(
|
||||
"sessions",
|
||||
"internal_metadata",
|
||||
server_default=sa.text("'{}'::jsonb"),
|
||||
schema=schema,
|
||||
)
|
||||
op.alter_column(
|
||||
"sessions",
|
||||
"configuration",
|
||||
server_default=sa.text("'{}'::jsonb"),
|
||||
schema=schema,
|
||||
)
|
||||
op.alter_column(
|
||||
"messages",
|
||||
"metadata",
|
||||
server_default=sa.text("'{}'::jsonb"),
|
||||
schema=schema,
|
||||
)
|
||||
op.alter_column(
|
||||
"messages",
|
||||
"internal_metadata",
|
||||
server_default=sa.text("'{}'::jsonb"),
|
||||
schema=schema,
|
||||
)
|
||||
op.alter_column(
|
||||
"collections",
|
||||
"metadata",
|
||||
server_default=sa.text("'{}'::jsonb"),
|
||||
schema=schema,
|
||||
)
|
||||
op.alter_column(
|
||||
"collections",
|
||||
"internal_metadata",
|
||||
server_default=sa.text("'{}'::jsonb"),
|
||||
schema=schema,
|
||||
)
|
||||
op.alter_column(
|
||||
"documents",
|
||||
"internal_metadata",
|
||||
server_default=sa.text("'{}'::jsonb"),
|
||||
schema=schema,
|
||||
)
|
||||
op.alter_column(
|
||||
"session_peers",
|
||||
"configuration",
|
||||
server_default=sa.text("'{}'::jsonb"),
|
||||
schema=schema,
|
||||
)
|
||||
op.alter_column(
|
||||
"session_peers",
|
||||
"internal_metadata",
|
||||
server_default=sa.text("'{}'::jsonb"),
|
||||
schema=schema,
|
||||
)
|
||||
|
||||
# Add server defaults for boolean columns
|
||||
op.alter_column(
|
||||
"sessions",
|
||||
"is_active",
|
||||
server_default=sa.text("true"),
|
||||
schema=schema,
|
||||
)
|
||||
op.alter_column(
|
||||
"queue",
|
||||
"processed",
|
||||
server_default=sa.text("false"),
|
||||
schema=schema,
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
# Remove server defaults for timestamp columns
|
||||
op.alter_column(
|
||||
"workspaces",
|
||||
"created_at",
|
||||
server_default=None,
|
||||
schema=schema,
|
||||
)
|
||||
op.alter_column(
|
||||
"peers",
|
||||
"created_at",
|
||||
server_default=None,
|
||||
schema=schema,
|
||||
)
|
||||
op.alter_column(
|
||||
"sessions",
|
||||
"created_at",
|
||||
server_default=None,
|
||||
schema=schema,
|
||||
)
|
||||
op.alter_column(
|
||||
"messages",
|
||||
"created_at",
|
||||
server_default=None,
|
||||
schema=schema,
|
||||
)
|
||||
op.alter_column(
|
||||
"message_embeddings",
|
||||
"created_at",
|
||||
server_default=None,
|
||||
schema=schema,
|
||||
)
|
||||
op.alter_column(
|
||||
"collections",
|
||||
"created_at",
|
||||
server_default=None,
|
||||
schema=schema,
|
||||
)
|
||||
op.alter_column(
|
||||
"documents",
|
||||
"created_at",
|
||||
server_default=None,
|
||||
schema=schema,
|
||||
)
|
||||
op.alter_column(
|
||||
"queue",
|
||||
"created_at",
|
||||
server_default=None,
|
||||
schema=schema,
|
||||
)
|
||||
op.alter_column(
|
||||
"webhook_endpoints",
|
||||
"created_at",
|
||||
server_default=None,
|
||||
schema=schema,
|
||||
)
|
||||
op.alter_column(
|
||||
"session_peers",
|
||||
"joined_at",
|
||||
server_default=None,
|
||||
schema=schema,
|
||||
)
|
||||
op.alter_column(
|
||||
"active_queue_sessions",
|
||||
"last_updated",
|
||||
server_default=None,
|
||||
schema=schema,
|
||||
)
|
||||
|
||||
# Remove server defaults for JSONB columns
|
||||
op.alter_column(
|
||||
"workspaces",
|
||||
"metadata",
|
||||
server_default=None,
|
||||
schema=schema,
|
||||
)
|
||||
op.alter_column(
|
||||
"workspaces",
|
||||
"internal_metadata",
|
||||
server_default=None,
|
||||
schema=schema,
|
||||
)
|
||||
op.alter_column(
|
||||
"workspaces",
|
||||
"configuration",
|
||||
server_default=None,
|
||||
schema=schema,
|
||||
)
|
||||
op.alter_column(
|
||||
"peers",
|
||||
"metadata",
|
||||
server_default=None,
|
||||
schema=schema,
|
||||
)
|
||||
op.alter_column(
|
||||
"peers",
|
||||
"internal_metadata",
|
||||
server_default=None,
|
||||
schema=schema,
|
||||
)
|
||||
op.alter_column(
|
||||
"peers",
|
||||
"configuration",
|
||||
server_default=None,
|
||||
schema=schema,
|
||||
)
|
||||
op.alter_column(
|
||||
"sessions",
|
||||
"metadata",
|
||||
server_default=None,
|
||||
schema=schema,
|
||||
)
|
||||
op.alter_column(
|
||||
"sessions",
|
||||
"internal_metadata",
|
||||
server_default=None,
|
||||
schema=schema,
|
||||
)
|
||||
op.alter_column(
|
||||
"sessions",
|
||||
"configuration",
|
||||
server_default=None,
|
||||
schema=schema,
|
||||
)
|
||||
op.alter_column(
|
||||
"messages",
|
||||
"metadata",
|
||||
server_default=None,
|
||||
schema=schema,
|
||||
)
|
||||
op.alter_column(
|
||||
"messages",
|
||||
"internal_metadata",
|
||||
server_default=None,
|
||||
schema=schema,
|
||||
)
|
||||
op.alter_column(
|
||||
"collections",
|
||||
"metadata",
|
||||
server_default=None,
|
||||
schema=schema,
|
||||
)
|
||||
op.alter_column(
|
||||
"collections",
|
||||
"internal_metadata",
|
||||
server_default=None,
|
||||
schema=schema,
|
||||
)
|
||||
op.alter_column(
|
||||
"documents",
|
||||
"internal_metadata",
|
||||
server_default=None,
|
||||
schema=schema,
|
||||
)
|
||||
op.alter_column(
|
||||
"session_peers",
|
||||
"configuration",
|
||||
server_default=None,
|
||||
schema=schema,
|
||||
)
|
||||
op.alter_column(
|
||||
"session_peers",
|
||||
"internal_metadata",
|
||||
server_default=None,
|
||||
schema=schema,
|
||||
)
|
||||
|
||||
# Remove server defaults for boolean columns
|
||||
op.alter_column(
|
||||
"sessions",
|
||||
"is_active",
|
||||
server_default=None,
|
||||
schema=schema,
|
||||
)
|
||||
op.alter_column(
|
||||
"queue",
|
||||
"processed",
|
||||
server_default=None,
|
||||
schema=schema,
|
||||
)
|
||||
|
|
@ -46,6 +46,8 @@ SessionLocal = async_sessionmaker(
|
|||
)
|
||||
|
||||
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.schema = table_schema
|
||||
Base = declarative_base(metadata=meta)
|
||||
|
|
|
|||
170
src/models.py
170
src/models.py
|
|
@ -51,13 +51,25 @@ session_peers_table = Table(
|
|||
nullable=False,
|
||||
),
|
||||
Column("peer_name", TEXT, primary_key=True, nullable=False),
|
||||
Column("configuration", JSONB, default=dict),
|
||||
Column("internal_metadata", JSONB, default=dict),
|
||||
Column(
|
||||
"configuration",
|
||||
JSONB,
|
||||
default=dict,
|
||||
nullable=False,
|
||||
server_default=text("'{}'::jsonb"),
|
||||
),
|
||||
Column(
|
||||
"internal_metadata",
|
||||
JSONB,
|
||||
default=dict,
|
||||
nullable=False,
|
||||
server_default=text("'{}'::jsonb"),
|
||||
),
|
||||
Column(
|
||||
"joined_at",
|
||||
DateTime(timezone=True),
|
||||
nullable=False,
|
||||
default=func.now(),
|
||||
server_default=func.now(),
|
||||
),
|
||||
Column(
|
||||
"left_at",
|
||||
|
|
@ -81,22 +93,28 @@ session_peers_table = Table(
|
|||
class Workspace(Base):
|
||||
__tablename__: str = "workspaces"
|
||||
id: Mapped[str] = mapped_column(TEXT, default=generate_nanoid, primary_key=True)
|
||||
name: Mapped[str] = mapped_column(TEXT, index=True, unique=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), index=True, default=func.now()
|
||||
DateTime(timezone=True), server_default=func.now()
|
||||
)
|
||||
h_metadata: Mapped[dict[str, Any]] = mapped_column(
|
||||
"metadata", JSONB, default=dict, server_default=text("'{}'::jsonb")
|
||||
)
|
||||
h_metadata: Mapped[dict[str, Any]] = mapped_column("metadata", JSONB, default=dict)
|
||||
internal_metadata: Mapped[dict[str, Any]] = mapped_column(
|
||||
"internal_metadata", JSONB, default=dict
|
||||
)
|
||||
configuration: Mapped[dict[str, Any]] = mapped_column(JSONB, default=dict)
|
||||
configuration: Mapped[dict[str, Any]] = mapped_column(
|
||||
JSONB, default=dict, server_default=text("'{}'::jsonb")
|
||||
)
|
||||
|
||||
__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"),
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -104,18 +122,22 @@ class Workspace(Base):
|
|||
class Peer(Base):
|
||||
__tablename__: str = "peers"
|
||||
id: Mapped[str] = mapped_column(TEXT, default=generate_nanoid, primary_key=True)
|
||||
name: Mapped[str] = mapped_column(TEXT, index=True)
|
||||
h_metadata: Mapped[dict[str, Any]] = mapped_column("metadata", JSONB, default=dict)
|
||||
name: Mapped[str] = mapped_column(TEXT, nullable=False)
|
||||
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")
|
||||
)
|
||||
created_at: Mapped[datetime.datetime] = mapped_column(
|
||||
DateTime(timezone=True), index=True, default=func.now()
|
||||
DateTime(timezone=True), server_default=func.now()
|
||||
)
|
||||
workspace_name: Mapped[str] = mapped_column(
|
||||
ForeignKey("workspaces.name"), index=True, nullable=False
|
||||
ForeignKey("workspaces.name"), nullable=False
|
||||
)
|
||||
configuration: Mapped[dict[str, Any]] = mapped_column(
|
||||
JSONB, default=dict, server_default=text("'{}'::jsonb")
|
||||
)
|
||||
configuration: Mapped[dict[str, Any]] = mapped_column(JSONB, default=dict)
|
||||
|
||||
workspace = relationship("Workspace", back_populates="peers")
|
||||
sessions = relationship(
|
||||
|
|
@ -128,6 +150,9 @@ class Peer(Base):
|
|||
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:
|
||||
|
|
@ -138,20 +163,24 @@ class Peer(Base):
|
|||
class Session(Base):
|
||||
__tablename__: str = "sessions"
|
||||
id: Mapped[str] = mapped_column(TEXT, primary_key=True, default=generate_nanoid)
|
||||
name: Mapped[str] = mapped_column(TEXT, index=True)
|
||||
is_active: Mapped[bool] = mapped_column(default=True)
|
||||
h_metadata: Mapped[dict[str, Any]] = mapped_column("metadata", JSONB, default=dict)
|
||||
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(
|
||||
"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")
|
||||
)
|
||||
created_at: Mapped[datetime.datetime] = mapped_column(
|
||||
DateTime(timezone=True), index=True, default=func.now()
|
||||
DateTime(timezone=True), server_default=func.now()
|
||||
)
|
||||
messages = relationship("Message", back_populates="session")
|
||||
workspace_name: Mapped[str] = mapped_column(
|
||||
ForeignKey("workspaces.name"), index=True, nullable=False
|
||||
ForeignKey("workspaces.name"), nullable=False
|
||||
)
|
||||
configuration: Mapped[dict[str, Any]] = mapped_column(
|
||||
JSONB, default=dict, server_default=text("'{}'::jsonb")
|
||||
)
|
||||
configuration: Mapped[dict[str, Any]] = mapped_column(JSONB, default=dict)
|
||||
|
||||
peers = relationship(
|
||||
"Peer", secondary=session_peers_table, back_populates="sessions"
|
||||
|
|
@ -162,6 +191,7 @@ class Session(Base):
|
|||
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:
|
||||
|
|
@ -175,26 +205,30 @@ class Message(Base):
|
|||
BigInteger, Identity(), primary_key=True, autoincrement=True
|
||||
)
|
||||
public_id: Mapped[str] = mapped_column(
|
||||
TEXT, index=True, unique=True, default=generate_nanoid
|
||||
TEXT,
|
||||
unique=True,
|
||||
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(index=True, nullable=False)
|
||||
session_name: Mapped[str] = mapped_column(TEXT, nullable=False)
|
||||
content: Mapped[str] = mapped_column(TEXT)
|
||||
h_metadata: Mapped[dict[str, Any]] = mapped_column("metadata", JSONB, default=dict)
|
||||
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")
|
||||
)
|
||||
token_count: Mapped[int] = mapped_column(Integer, default=0, nullable=False)
|
||||
seq_in_session: Mapped[int] = mapped_column(BigInteger, nullable=False, index=True)
|
||||
seq_in_session: Mapped[int] = mapped_column(BigInteger, nullable=False)
|
||||
|
||||
created_at: Mapped[datetime.datetime] = mapped_column(
|
||||
DateTime(timezone=True), index=True, default=func.now()
|
||||
DateTime(timezone=True), server_default=func.now()
|
||||
)
|
||||
session = relationship("Session", back_populates="messages")
|
||||
peer_name: Mapped[str] = mapped_column(index=True)
|
||||
peer_name: Mapped[str] = mapped_column(TEXT)
|
||||
workspace_name: Mapped[str] = mapped_column(
|
||||
ForeignKey("workspaces.name"), index=True
|
||||
ForeignKey("workspaces.name"),
|
||||
)
|
||||
|
||||
__table_args__ = (
|
||||
|
|
@ -229,6 +263,11 @@ class Message(Base):
|
|||
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
|
||||
|
|
@ -246,15 +285,15 @@ 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"), index=True
|
||||
ForeignKey("messages.public_id"),
|
||||
)
|
||||
workspace_name: Mapped[str] = mapped_column(
|
||||
ForeignKey("workspaces.name"), index=True
|
||||
ForeignKey("workspaces.name"),
|
||||
)
|
||||
session_name: Mapped[str] = mapped_column(TEXT, index=True, nullable=False)
|
||||
peer_name: Mapped[str | None] = mapped_column(TEXT, index=True)
|
||||
session_name: Mapped[str] = mapped_column(TEXT, nullable=False)
|
||||
peer_name: Mapped[str] = mapped_column(TEXT)
|
||||
created_at: Mapped[datetime.datetime] = mapped_column(
|
||||
DateTime(timezone=True), index=True, default=func.now()
|
||||
DateTime(timezone=True), server_default=func.now()
|
||||
)
|
||||
|
||||
# Relationship to Message
|
||||
|
|
@ -278,6 +317,11 @@ class MessageEmbedding(Base):
|
|||
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"),
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -286,20 +330,22 @@ 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)
|
||||
observer: Mapped[str] = mapped_column(TEXT)
|
||||
observed: Mapped[str] = mapped_column(TEXT)
|
||||
created_at: Mapped[datetime.datetime] = mapped_column(
|
||||
DateTime(timezone=True), index=True, default=func.now()
|
||||
DateTime(timezone=True), server_default=func.now()
|
||||
)
|
||||
h_metadata: Mapped[dict[str, Any]] = mapped_column(
|
||||
"metadata", JSONB, default=dict, server_default=text("'{}'::jsonb")
|
||||
)
|
||||
h_metadata: Mapped[dict[str, Any]] = mapped_column("metadata", JSONB, default=dict)
|
||||
internal_metadata: Mapped[dict[str, Any]] = mapped_column(
|
||||
"internal_metadata", JSONB, default=dict
|
||||
"internal_metadata", JSONB, default=dict, server_default=text("'{}'::jsonb")
|
||||
)
|
||||
documents = relationship(
|
||||
"Document", back_populates="collection", cascade="all, delete, delete-orphan"
|
||||
)
|
||||
workspace_name: Mapped[str] = mapped_column(
|
||||
ForeignKey("workspaces.name"), index=True
|
||||
ForeignKey("workspaces.name"),
|
||||
)
|
||||
|
||||
__table_args__ = (
|
||||
|
|
@ -321,6 +367,10 @@ 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"),
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -329,20 +379,18 @@ class Document(Base):
|
|||
__tablename__: str = "documents"
|
||||
id: Mapped[str] = mapped_column(TEXT, default=generate_nanoid, primary_key=True)
|
||||
internal_metadata: Mapped[dict[str, Any]] = mapped_column(
|
||||
"internal_metadata", JSONB, default=dict
|
||||
"internal_metadata", JSONB, default=dict, server_default=text("'{}'::jsonb")
|
||||
)
|
||||
content: Mapped[str] = mapped_column(TEXT)
|
||||
embedding: MappedColumn[Any] = mapped_column(Vector(1536))
|
||||
created_at: Mapped[datetime.datetime] = mapped_column(
|
||||
DateTime(timezone=True), index=True, default=func.now()
|
||||
DateTime(timezone=True), server_default=func.now()
|
||||
)
|
||||
|
||||
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"), index=True
|
||||
)
|
||||
session_name: Mapped[str] = mapped_column(TEXT, 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)
|
||||
collection = relationship("Collection", back_populates="documents")
|
||||
|
||||
__table_args__ = (
|
||||
|
|
@ -383,6 +431,11 @@ 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"),
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -395,17 +448,22 @@ class QueueItem(Base):
|
|||
id: Mapped[int] = mapped_column(
|
||||
BigInteger, Identity(), primary_key=True, autoincrement=True
|
||||
)
|
||||
session_id: Mapped[str] = mapped_column(
|
||||
ForeignKey("sessions.id"), index=True, nullable=True
|
||||
)
|
||||
session_id: Mapped[str] = mapped_column(ForeignKey("sessions.id"), nullable=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)
|
||||
processed: Mapped[bool] = mapped_column(
|
||||
Boolean, default=False, server_default=text("false")
|
||||
)
|
||||
error: Mapped[str | None] = mapped_column(TEXT, nullable=True)
|
||||
created_at: Mapped[datetime.datetime] = mapped_column(
|
||||
DateTime(timezone=True), index=True, default=func.now()
|
||||
DateTime(timezone=True), server_default=func.now()
|
||||
)
|
||||
|
||||
__table_args__ = (
|
||||
Index("ix_queue_created_at", "created_at"),
|
||||
Index("ix_queue_session_id", "session_id"),
|
||||
)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
|
|
@ -418,10 +476,10 @@ class ActiveQueueSession(Base):
|
|||
|
||||
id: Mapped[str] = mapped_column(TEXT, default=generate_nanoid, primary_key=True)
|
||||
|
||||
work_unit_key: Mapped[str] = mapped_column(TEXT, unique=True, index=True)
|
||||
work_unit_key: Mapped[str] = mapped_column(TEXT, unique=True)
|
||||
|
||||
last_updated: Mapped[datetime.datetime] = mapped_column(
|
||||
DateTime(timezone=True), default=func.now(), onupdate=func.now()
|
||||
DateTime(timezone=True), server_default=func.now(), onupdate=func.now()
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -430,11 +488,11 @@ 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"), index=True, nullable=False
|
||||
ForeignKey("workspaces.name"), nullable=False
|
||||
)
|
||||
url: Mapped[str] = mapped_column(TEXT, nullable=False)
|
||||
created_at: Mapped[datetime.datetime] = mapped_column(
|
||||
DateTime(timezone=True), default=func.now()
|
||||
DateTime(timezone=True), server_default=func.now()
|
||||
)
|
||||
|
||||
workspace = relationship("Workspace", back_populates="webhook_endpoints")
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@
|
|||
|
||||
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_20f89a421aff_rename_metamessage_type_to_label,
|
||||
test_66e63cf2cf77_add_indexes_to_documents_table,
|
||||
|
|
@ -19,6 +20,7 @@ from . import (
|
|||
|
||||
__all__ = [
|
||||
"test_05486ce795d5_make_session_name_required_on_messages",
|
||||
"test_066e87ca5b07_align_schema_with_declarative_models",
|
||||
"test_08894082221a_replace_collection_name_with_observer_",
|
||||
"test_20f89a421aff_rename_metamessage_type_to_label",
|
||||
"test_556a16564f50_add_user_id_and_app_id_to_tables",
|
||||
|
|
|
|||
|
|
@ -0,0 +1,42 @@
|
|||
"""Hooks for revision 066e87ca5b07 (align_schema_with_declarative_models)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from tests.alembic.registry import register_after_upgrade, register_before_upgrade
|
||||
from tests.alembic.verifier import MigrationVerifier
|
||||
|
||||
|
||||
@register_before_upgrade("066e87ca5b07")
|
||||
def prepare_align_schema_with_declarative_models(verifier: MigrationVerifier) -> None:
|
||||
"""Seed state and assertions before upgrading to 066e87ca5b07."""
|
||||
# Assert columns exist but are nullable before migration
|
||||
verifier.assert_column_exists(
|
||||
"active_queue_sessions", "work_unit_key", exists=True, nullable=True
|
||||
)
|
||||
verifier.assert_column_exists("documents", "embedding", exists=True, nullable=True)
|
||||
|
||||
# Assert FK constraint does not exist yet
|
||||
verifier.assert_constraint_exists(
|
||||
"queue", "fk_queue_session_id", "foreign_key", exists=False
|
||||
)
|
||||
|
||||
|
||||
@register_after_upgrade("066e87ca5b07")
|
||||
def verify_align_schema_with_declarative_models(verifier: MigrationVerifier) -> None:
|
||||
"""Add assertions validating the effects of 066e87ca5b07."""
|
||||
# Assert columns are now non-nullable
|
||||
verifier.assert_column_exists(
|
||||
"peers", "workspace_name", exists=True, nullable=False
|
||||
)
|
||||
verifier.assert_column_exists(
|
||||
"sessions", "workspace_name", exists=True, nullable=False
|
||||
)
|
||||
verifier.assert_column_exists(
|
||||
"active_queue_sessions", "work_unit_key", exists=True, nullable=False
|
||||
)
|
||||
verifier.assert_column_exists("documents", "embedding", exists=True, nullable=False)
|
||||
|
||||
# Assert FK constraint now exists
|
||||
verifier.assert_constraint_exists(
|
||||
"queue", "fk_queue_session_id", "foreign_key", exists=True
|
||||
)
|
||||
|
|
@ -0,0 +1,284 @@
|
|||
"""Hooks for revision e9b705f9adf9 (add server defaults to timestamp, boolean, and jsonb columns)."""
|
||||
|
||||
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
|
||||
|
||||
# Test data IDs
|
||||
WORKSPACE_ID = generate_nanoid()
|
||||
PEER_ID = generate_nanoid()
|
||||
SESSION_ID = generate_nanoid()
|
||||
MESSAGE_ID = generate_nanoid()
|
||||
COLLECTION_ID = generate_nanoid()
|
||||
DOCUMENT_ID = generate_nanoid()
|
||||
|
||||
|
||||
@register_before_upgrade("e9b705f9adf9")
|
||||
def prepare_add_server_defaults(verifier: MigrationVerifier) -> None:
|
||||
"""Seed state before upgrading to e9b705f9adf9.
|
||||
|
||||
This migration adds server defaults to timestamp, JSONB, and boolean columns.
|
||||
We verify that columns exist but don't have server defaults before the migration.
|
||||
"""
|
||||
conn = verifier.conn
|
||||
schema = verifier.schema
|
||||
inspector = verifier.get_inspector()
|
||||
|
||||
# Sample timestamp columns to check - they should exist but without server defaults
|
||||
for table, column in [
|
||||
("workspaces", "created_at"),
|
||||
("peers", "created_at"),
|
||||
("sessions", "created_at"),
|
||||
("messages", "created_at"),
|
||||
("collections", "created_at"),
|
||||
("documents", "created_at"),
|
||||
("queue", "created_at"),
|
||||
]:
|
||||
columns = inspector.get_columns(table, schema=schema)
|
||||
col_info = next((c for c in columns if c["name"] == column), None)
|
||||
assert (
|
||||
col_info is not None
|
||||
), f"Column {table}.{column} should exist before migration"
|
||||
|
||||
# Create test data to ensure existing rows work after migration
|
||||
conn.execute(
|
||||
text(
|
||||
f'INSERT INTO "{schema}"."workspaces" '
|
||||
+ '("id", "name", "created_at", "metadata", "internal_metadata", "configuration") '
|
||||
+ "VALUES (:id, :name, NOW(), :metadata, :internal_metadata, :configuration)"
|
||||
),
|
||||
{
|
||||
"id": WORKSPACE_ID,
|
||||
"name": "test-workspace",
|
||||
"metadata": "{}",
|
||||
"internal_metadata": "{}",
|
||||
"configuration": "{}",
|
||||
},
|
||||
)
|
||||
|
||||
conn.execute(
|
||||
text(
|
||||
f'INSERT INTO "{schema}"."peers" '
|
||||
+ '("id", "name", "workspace_name", "created_at", "metadata", "internal_metadata", "configuration") '
|
||||
+ "VALUES (:id, :name, :workspace_name, NOW(), :metadata, :internal_metadata, :configuration)"
|
||||
),
|
||||
{
|
||||
"id": PEER_ID,
|
||||
"name": "test-peer",
|
||||
"workspace_name": "test-workspace",
|
||||
"metadata": "{}",
|
||||
"internal_metadata": "{}",
|
||||
"configuration": "{}",
|
||||
},
|
||||
)
|
||||
|
||||
conn.execute(
|
||||
text(
|
||||
f'INSERT INTO "{schema}"."sessions" '
|
||||
+ '("id", "name", "workspace_name", "created_at", "is_active", "metadata", "internal_metadata", "configuration") '
|
||||
+ "VALUES (:id, :name, :workspace_name, NOW(), true, :metadata, :internal_metadata, :configuration)"
|
||||
),
|
||||
{
|
||||
"id": SESSION_ID,
|
||||
"name": "test-session",
|
||||
"workspace_name": "test-workspace",
|
||||
"metadata": "{}",
|
||||
"internal_metadata": "{}",
|
||||
"configuration": "{}",
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@register_after_upgrade("e9b705f9adf9")
|
||||
def verify_add_server_defaults(verifier: MigrationVerifier) -> None:
|
||||
"""Validate server defaults were added correctly to all columns."""
|
||||
conn = verifier.conn
|
||||
schema = verifier.schema
|
||||
inspector = verifier.get_inspector()
|
||||
|
||||
# Verify timestamp columns have server defaults (now() function)
|
||||
timestamp_columns = [
|
||||
("workspaces", "created_at"),
|
||||
("peers", "created_at"),
|
||||
("sessions", "created_at"),
|
||||
("messages", "created_at"),
|
||||
("message_embeddings", "created_at"),
|
||||
("collections", "created_at"),
|
||||
("documents", "created_at"),
|
||||
("queue", "created_at"),
|
||||
("webhook_endpoints", "created_at"),
|
||||
("session_peers", "joined_at"),
|
||||
("active_queue_sessions", "last_updated"),
|
||||
]
|
||||
|
||||
for table, column in timestamp_columns:
|
||||
columns = inspector.get_columns(table, schema=schema)
|
||||
col_info = next((c for c in columns if c["name"] == column), None)
|
||||
assert (
|
||||
col_info is not None
|
||||
), f"Column {table}.{column} not found after migration"
|
||||
|
||||
# Check that a server default exists
|
||||
default = col_info.get("default")
|
||||
assert default is not None, (
|
||||
f"Column {table}.{column} should have a server default after migration, "
|
||||
f"but default is None"
|
||||
)
|
||||
|
||||
# Verify JSONB columns have server defaults (empty object '{}')
|
||||
jsonb_columns = [
|
||||
("workspaces", "metadata"),
|
||||
("workspaces", "internal_metadata"),
|
||||
("workspaces", "configuration"),
|
||||
("peers", "metadata"),
|
||||
("peers", "internal_metadata"),
|
||||
("peers", "configuration"),
|
||||
("sessions", "metadata"),
|
||||
("sessions", "internal_metadata"),
|
||||
("sessions", "configuration"),
|
||||
("messages", "metadata"),
|
||||
("messages", "internal_metadata"),
|
||||
("collections", "metadata"),
|
||||
("collections", "internal_metadata"),
|
||||
("documents", "internal_metadata"),
|
||||
("session_peers", "configuration"),
|
||||
("session_peers", "internal_metadata"),
|
||||
]
|
||||
|
||||
for table, column in jsonb_columns:
|
||||
columns = inspector.get_columns(table, schema=schema)
|
||||
col_info = next((c for c in columns if c["name"] == column), None)
|
||||
assert (
|
||||
col_info is not None
|
||||
), f"Column {table}.{column} not found after migration"
|
||||
|
||||
# Check that a server default exists
|
||||
default = col_info.get("default")
|
||||
assert default is not None, (
|
||||
f"Column {table}.{column} should have a server default after migration, "
|
||||
f"but default is None"
|
||||
)
|
||||
|
||||
# Verify boolean columns have server defaults
|
||||
boolean_columns = [
|
||||
("sessions", "is_active", "true"),
|
||||
("queue", "processed", "false"),
|
||||
]
|
||||
|
||||
for table, column, expected_default in boolean_columns:
|
||||
columns = inspector.get_columns(table, schema=schema)
|
||||
col_info = next((c for c in columns if c["name"] == column), None)
|
||||
assert (
|
||||
col_info is not None
|
||||
), f"Column {table}.{column} not found after migration"
|
||||
|
||||
# Check that a server default exists
|
||||
default = col_info.get("default")
|
||||
assert default is not None, (
|
||||
f"Column {table}.{column} should have a server default after migration, "
|
||||
f"but default is None"
|
||||
)
|
||||
|
||||
assert (
|
||||
default == expected_default
|
||||
), f"Column {table}.{column} should have a server default of {expected_default} after migration, but default is {default}"
|
||||
|
||||
# Test that defaults actually work by inserting rows without explicit values
|
||||
test_workspace_id = generate_nanoid()
|
||||
conn.execute(
|
||||
text(
|
||||
f'INSERT INTO "{schema}"."workspaces" ("id", "name") '
|
||||
+ "VALUES (:id, :name)"
|
||||
),
|
||||
{"id": test_workspace_id, "name": "test-defaults-workspace"},
|
||||
)
|
||||
|
||||
# Verify the inserted workspace has default values
|
||||
workspace = conn.execute(
|
||||
text(
|
||||
'SELECT "created_at", "metadata", "internal_metadata", "configuration" '
|
||||
+ f'FROM "{schema}"."workspaces" WHERE "id" = :id'
|
||||
),
|
||||
{"id": test_workspace_id},
|
||||
).one()
|
||||
|
||||
assert workspace.created_at is not None, "created_at should be auto-populated"
|
||||
assert workspace.metadata == {}, "metadata should default to empty object"
|
||||
assert (
|
||||
workspace.internal_metadata == {}
|
||||
), "internal_metadata should default to empty object"
|
||||
assert workspace.configuration == {}, "configuration should default to empty object"
|
||||
|
||||
# Test peer defaults
|
||||
test_peer_id = generate_nanoid()
|
||||
conn.execute(
|
||||
text(
|
||||
f'INSERT INTO "{schema}"."peers" ("id", "name", "workspace_name") '
|
||||
+ "VALUES (:id, :name, :workspace_name)"
|
||||
),
|
||||
{
|
||||
"id": test_peer_id,
|
||||
"name": "test-defaults-peer",
|
||||
"workspace_name": "test-defaults-workspace",
|
||||
},
|
||||
)
|
||||
|
||||
peer = conn.execute(
|
||||
text(
|
||||
'SELECT "created_at", "metadata", "internal_metadata", "configuration" '
|
||||
+ f'FROM "{schema}"."peers" WHERE "id" = :id'
|
||||
),
|
||||
{"id": test_peer_id},
|
||||
).one()
|
||||
|
||||
assert peer.created_at is not None, "peer created_at should be auto-populated"
|
||||
assert peer.metadata == {}, "peer metadata should default to empty object"
|
||||
assert (
|
||||
peer.internal_metadata == {}
|
||||
), "peer internal_metadata should default to empty object"
|
||||
assert peer.configuration == {}, "peer configuration should default to empty object"
|
||||
|
||||
# Test session defaults (including boolean is_active)
|
||||
test_session_id = generate_nanoid()
|
||||
conn.execute(
|
||||
text(
|
||||
f'INSERT INTO "{schema}"."sessions" ("id", "name", "workspace_name") '
|
||||
+ "VALUES (:id, :name, :workspace_name)"
|
||||
),
|
||||
{
|
||||
"id": test_session_id,
|
||||
"name": "test-defaults-session",
|
||||
"workspace_name": "test-defaults-workspace",
|
||||
},
|
||||
)
|
||||
|
||||
session = conn.execute(
|
||||
text(
|
||||
'SELECT "created_at", "is_active", "metadata", "internal_metadata", "configuration" '
|
||||
+ f'FROM "{schema}"."sessions" WHERE "id" = :id'
|
||||
),
|
||||
{"id": test_session_id},
|
||||
).one()
|
||||
|
||||
assert session.created_at is not None, "session created_at should be auto-populated"
|
||||
assert session.is_active is True, "session is_active should default to true"
|
||||
assert session.metadata == {}, "session metadata should default to empty object"
|
||||
assert (
|
||||
session.internal_metadata == {}
|
||||
), "session internal_metadata should default to empty object"
|
||||
assert (
|
||||
session.configuration == {}
|
||||
), "session configuration should default to empty object"
|
||||
|
||||
# Verify pre-existing data still exists
|
||||
existing_workspace = conn.execute(
|
||||
text(f'SELECT "id" FROM "{schema}"."workspaces" WHERE "id" = :id'),
|
||||
{"id": WORKSPACE_ID},
|
||||
).one_or_none()
|
||||
assert (
|
||||
existing_workspace is not None
|
||||
), "Pre-existing workspace should still exist after migration"
|
||||
Loading…
Reference in New Issue