360 lines
12 KiB
Python
360 lines
12 KiB
Python
"""make embeddings nullable, add soft delete, add vector sync state
|
|
|
|
This migration:
|
|
1. Makes embedding columns nullable in both message_embeddings and documents tables
|
|
since embeddings are now stored in external vector stores instead of PostgreSQL.
|
|
2. Adds deleted_at column to documents table for soft delete support, enabling
|
|
hybrid sync/soft delete pattern for vector store consistency.
|
|
3. Adds sync_state, last_sync_at, and sync_attempts columns to documents and
|
|
message_embeddings tables for tracking vector store synchronization status.
|
|
4. Adds partial unique index on queue table for reconciler task deduplication,
|
|
ensuring only one pending reconciler task exists per work_unit_key.
|
|
5. Makes workspace_name nullable on queue table for system-level tasks (e.g., reconciler)
|
|
that don't belong to any specific workspace.
|
|
|
|
Revision ID: 119a52b73c60
|
|
Revises: 7c0d9a4e3b1f
|
|
Create Date: 2025-11-24 12:00:00.000000
|
|
|
|
"""
|
|
|
|
from collections.abc import Sequence
|
|
|
|
import sqlalchemy as sa
|
|
from alembic import op
|
|
from pgvector.sqlalchemy import Vector
|
|
|
|
from migrations.utils import column_exists, constraint_exists, get_schema, index_exists
|
|
|
|
# revision identifiers, used by Alembic.
|
|
revision: str = "119a52b73c60"
|
|
down_revision: str | None = "7c0d9a4e3b1f"
|
|
branch_labels: str | Sequence[str] | None = None
|
|
depends_on: str | Sequence[str] | None = None
|
|
|
|
schema = get_schema()
|
|
|
|
|
|
def upgrade() -> None:
|
|
"""Make embeddings nullable, add deleted_at, add sync state."""
|
|
inspector = sa.inspect(op.get_bind())
|
|
|
|
op.alter_column(
|
|
"message_embeddings",
|
|
"embedding",
|
|
existing_type=Vector(1536),
|
|
nullable=True,
|
|
schema=schema,
|
|
)
|
|
|
|
op.alter_column(
|
|
"documents",
|
|
"embedding",
|
|
existing_type=Vector(1536),
|
|
nullable=True,
|
|
schema=schema,
|
|
)
|
|
|
|
# Add deleted_at column to documents for soft delete support
|
|
if not column_exists("documents", "deleted_at", inspector):
|
|
op.add_column(
|
|
"documents",
|
|
sa.Column(
|
|
"deleted_at",
|
|
sa.DateTime(timezone=True),
|
|
nullable=True,
|
|
),
|
|
schema=schema,
|
|
)
|
|
# Create partial index for efficient cleanup queries (only index non-null values)
|
|
op.create_index(
|
|
"ix_documents_deleted_at",
|
|
"documents",
|
|
["deleted_at"],
|
|
schema=schema,
|
|
postgresql_where=sa.text("deleted_at IS NOT NULL"),
|
|
)
|
|
|
|
# Add sync state columns to documents table
|
|
if not column_exists("documents", "sync_state", inspector):
|
|
op.add_column(
|
|
"documents",
|
|
sa.Column(
|
|
"sync_state",
|
|
sa.TEXT(),
|
|
nullable=False,
|
|
server_default=sa.text(
|
|
"'pending'"
|
|
), # Existing records need reconciliation
|
|
),
|
|
schema=schema,
|
|
)
|
|
op.create_index(
|
|
"ix_documents_sync_state",
|
|
"documents",
|
|
["sync_state"],
|
|
schema=schema,
|
|
)
|
|
|
|
if not column_exists("documents", "last_sync_at", inspector):
|
|
op.add_column(
|
|
"documents",
|
|
sa.Column(
|
|
"last_sync_at",
|
|
sa.DateTime(timezone=True),
|
|
nullable=True,
|
|
),
|
|
schema=schema,
|
|
)
|
|
|
|
if not column_exists("documents", "sync_attempts", inspector):
|
|
op.add_column(
|
|
"documents",
|
|
sa.Column(
|
|
"sync_attempts",
|
|
sa.Integer(),
|
|
nullable=False,
|
|
server_default=sa.text("0"),
|
|
),
|
|
schema=schema,
|
|
)
|
|
|
|
# Add composite index for efficient reconciliation queries after both columns exist
|
|
# Reconciliation orders by: WHERE sync_state='pending' ORDER BY last_sync_at
|
|
if not index_exists("documents", "ix_documents_sync_state_last_sync_at", inspector):
|
|
op.create_index(
|
|
"ix_documents_sync_state_last_sync_at",
|
|
"documents",
|
|
["sync_state", "last_sync_at"],
|
|
schema=schema,
|
|
)
|
|
|
|
# Add sync state columns to message_embeddings table
|
|
if not column_exists("message_embeddings", "sync_state", inspector):
|
|
op.add_column(
|
|
"message_embeddings",
|
|
sa.Column(
|
|
"sync_state",
|
|
sa.TEXT(),
|
|
nullable=False,
|
|
server_default=sa.text(
|
|
"'pending'"
|
|
), # Existing records need reconciliation
|
|
),
|
|
schema=schema,
|
|
)
|
|
op.create_index(
|
|
"ix_message_embeddings_sync_state",
|
|
"message_embeddings",
|
|
["sync_state"],
|
|
schema=schema,
|
|
)
|
|
|
|
if not column_exists("message_embeddings", "last_sync_at", inspector):
|
|
op.add_column(
|
|
"message_embeddings",
|
|
sa.Column(
|
|
"last_sync_at",
|
|
sa.DateTime(timezone=True),
|
|
nullable=True,
|
|
),
|
|
schema=schema,
|
|
)
|
|
|
|
if not column_exists("message_embeddings", "sync_attempts", inspector):
|
|
op.add_column(
|
|
"message_embeddings",
|
|
sa.Column(
|
|
"sync_attempts",
|
|
sa.Integer(),
|
|
nullable=False,
|
|
server_default=sa.text("0"),
|
|
),
|
|
schema=schema,
|
|
)
|
|
|
|
# Add composite index for efficient reconciliation queries after both columns exist
|
|
# Reconciliation orders by: WHERE sync_state='pending' ORDER BY last_sync_at
|
|
if not index_exists(
|
|
"message_embeddings", "ix_message_embeddings_sync_state_last_sync_at", inspector
|
|
):
|
|
op.create_index(
|
|
"ix_message_embeddings_sync_state_last_sync_at",
|
|
"message_embeddings",
|
|
["sync_state", "last_sync_at"],
|
|
schema=schema,
|
|
)
|
|
|
|
# Add partial unique index on queue table for reconciler task deduplication
|
|
# This ensures only one pending reconciler task exists per work_unit_key
|
|
if not index_exists(
|
|
"queue", "uq_queue_reconciler_pending_work_unit_key", inspector
|
|
):
|
|
op.create_index(
|
|
"uq_queue_reconciler_pending_work_unit_key",
|
|
"queue",
|
|
["work_unit_key"],
|
|
unique=True,
|
|
schema=schema,
|
|
postgresql_where=sa.text("task_type = 'reconciler' AND processed = false"),
|
|
)
|
|
|
|
# Make workspace_name nullable on queue table for system-level tasks
|
|
# This requires dropping and recreating the FK constraint
|
|
if constraint_exists("queue", "fk_queue_workspace_name", "foreignkey", inspector):
|
|
op.drop_constraint(
|
|
"fk_queue_workspace_name", "queue", type_="foreignkey", schema=schema
|
|
)
|
|
op.alter_column(
|
|
"queue",
|
|
"workspace_name",
|
|
existing_type=sa.TEXT(),
|
|
nullable=True,
|
|
schema=schema,
|
|
)
|
|
op.create_foreign_key(
|
|
"fk_queue_workspace_name",
|
|
"queue",
|
|
"workspaces",
|
|
["workspace_name"],
|
|
["name"],
|
|
source_schema=schema,
|
|
referent_schema=schema,
|
|
)
|
|
|
|
|
|
def downgrade() -> None:
|
|
"""Remove deleted_at columns and revert embedding columns."""
|
|
inspector = sa.inspect(op.get_bind())
|
|
conn = op.get_bind()
|
|
|
|
# Delete system-level queue items (with NULL workspace_name) before reverting to NOT NULL
|
|
# First delete any active_queue_sessions referencing these queue items
|
|
batch_size = 5000
|
|
|
|
# Delete active_queue_sessions for queue items with NULL workspace_name
|
|
# Use CTE-based delete since PostgreSQL doesn't support DELETE ... LIMIT
|
|
while True:
|
|
result = conn.execute(
|
|
sa.text(
|
|
f"""
|
|
WITH del AS (
|
|
SELECT work_unit_key FROM "{schema}".queue
|
|
WHERE workspace_name IS NULL
|
|
LIMIT :batch_size
|
|
)
|
|
DELETE FROM "{schema}".active_queue_sessions
|
|
WHERE work_unit_key IN (SELECT work_unit_key FROM del)
|
|
"""
|
|
),
|
|
{"batch_size": batch_size},
|
|
)
|
|
if result.rowcount == 0:
|
|
break
|
|
|
|
# Delete queue items with NULL workspace_name
|
|
# Use CTE with ctid for safe batched delete in PostgreSQL
|
|
while True:
|
|
result = conn.execute(
|
|
sa.text(
|
|
f"""
|
|
WITH del AS (
|
|
SELECT ctid FROM "{schema}".queue
|
|
WHERE workspace_name IS NULL
|
|
LIMIT :batch_size
|
|
)
|
|
DELETE FROM "{schema}".queue
|
|
WHERE ctid IN (SELECT ctid FROM del)
|
|
"""
|
|
),
|
|
{"batch_size": batch_size},
|
|
)
|
|
if result.rowcount == 0:
|
|
break
|
|
|
|
# Revert workspace_name to NOT NULL on queue table
|
|
op.drop_constraint(
|
|
"fk_queue_workspace_name", "queue", type_="foreignkey", schema=schema
|
|
)
|
|
op.alter_column(
|
|
"queue",
|
|
"workspace_name",
|
|
existing_type=sa.TEXT(),
|
|
nullable=False,
|
|
schema=schema,
|
|
)
|
|
op.create_foreign_key(
|
|
"fk_queue_workspace_name",
|
|
"queue",
|
|
"workspaces",
|
|
["workspace_name"],
|
|
["name"],
|
|
source_schema=schema,
|
|
referent_schema=schema,
|
|
)
|
|
|
|
# Drop reconciler queue index if it exists
|
|
if index_exists("queue", "uq_queue_reconciler_pending_work_unit_key", inspector):
|
|
op.drop_index(
|
|
"uq_queue_reconciler_pending_work_unit_key",
|
|
table_name="queue",
|
|
schema=schema,
|
|
)
|
|
|
|
# Drop message_embeddings indexes if they exist
|
|
if index_exists(
|
|
"message_embeddings", "ix_message_embeddings_sync_state_last_sync_at", inspector
|
|
):
|
|
op.drop_index(
|
|
"ix_message_embeddings_sync_state_last_sync_at",
|
|
table_name="message_embeddings",
|
|
schema=schema,
|
|
)
|
|
if index_exists(
|
|
"message_embeddings", "ix_message_embeddings_sync_state", inspector
|
|
):
|
|
op.drop_index(
|
|
"ix_message_embeddings_sync_state",
|
|
table_name="message_embeddings",
|
|
schema=schema,
|
|
)
|
|
|
|
if column_exists("message_embeddings", "sync_state", inspector):
|
|
op.drop_column("message_embeddings", "sync_state", schema=schema)
|
|
|
|
if column_exists("message_embeddings", "sync_attempts", inspector):
|
|
op.drop_column("message_embeddings", "sync_attempts", schema=schema)
|
|
|
|
if column_exists("message_embeddings", "last_sync_at", inspector):
|
|
op.drop_column("message_embeddings", "last_sync_at", schema=schema)
|
|
|
|
# Drop documents indexes if they exist
|
|
if index_exists("documents", "ix_documents_sync_state_last_sync_at", inspector):
|
|
op.drop_index(
|
|
"ix_documents_sync_state_last_sync_at",
|
|
table_name="documents",
|
|
schema=schema,
|
|
)
|
|
if index_exists("documents", "ix_documents_sync_state", inspector):
|
|
op.drop_index("ix_documents_sync_state", table_name="documents", schema=schema)
|
|
|
|
if column_exists("documents", "sync_state", inspector):
|
|
op.drop_column("documents", "sync_state", schema=schema)
|
|
|
|
# Remove sync state columns from documents
|
|
if column_exists("documents", "sync_attempts", inspector):
|
|
op.drop_column("documents", "sync_attempts", schema=schema)
|
|
|
|
if column_exists("documents", "last_sync_at", inspector):
|
|
op.drop_column("documents", "last_sync_at", schema=schema)
|
|
|
|
# Remove deleted_at column and index from documents
|
|
if index_exists("documents", "ix_documents_deleted_at", inspector):
|
|
op.drop_index("ix_documents_deleted_at", table_name="documents", schema=schema)
|
|
if column_exists("documents", "deleted_at", inspector):
|
|
op.drop_column("documents", "deleted_at", schema=schema)
|
|
|
|
# NOTE: This downgrade does NOT restore the NOT NULL constraint on embedding columns
|
|
# in message_embeddings and documents tables. This is intentional to avoid migration
|
|
# failures if NULL embedding values exist (which is expected when using external vector stores).
|