342 lines
12 KiB
Python
342 lines
12 KiB
Python
"""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,
|
|
)
|