diff --git a/migrations/versions/119a52b73c60_support_external_embeddings.py b/migrations/versions/119a52b73c60_support_external_embeddings.py index e07bc0c0..ab3dc431 100644 --- a/migrations/versions/119a52b73c60_support_external_embeddings.py +++ b/migrations/versions/119a52b73c60_support_external_embeddings.py @@ -9,6 +9,8 @@ This migration: 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 @@ -22,7 +24,7 @@ import sqlalchemy as sa from alembic import op from pgvector.sqlalchemy import Vector -from migrations.utils import column_exists, get_schema, index_exists +from migrations.utils import column_exists, constraint_exists, get_schema, index_exists # revision identifiers, used by Alembic. revision: str = "119a52b73c60" @@ -195,10 +197,92 @@ def upgrade() -> None: 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 + while True: + result = conn.execute( + sa.text( + f""" + DELETE FROM "{schema}".active_queue_sessions + WHERE work_unit_key IN ( + SELECT work_unit_key FROM "{schema}".queue + WHERE workspace_name IS NULL + ) + LIMIT :batch_size + """ + ), + {"batch_size": batch_size}, + ) + if result.rowcount == 0: + break + + # Delete queue items with NULL workspace_name + while True: + result = conn.execute( + sa.text( + f""" + DELETE FROM "{schema}".queue + WHERE workspace_name IS NULL + LIMIT :batch_size + """ + ), + {"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_work_unit_key", inspector): diff --git a/src/deriver/consumer.py b/src/deriver/consumer.py index ae9bb97e..212be5a7 100644 --- a/src/deriver/consumer.py +++ b/src/deriver/consumer.py @@ -34,6 +34,25 @@ async def process_item(queue_item: models.QueueItem) -> None: queue_payload = queue_item.payload workspace_name = queue_item.workspace_name + # Handle reconciler first - it's the only task type that doesn't require workspace_name + if task_type == "reconciler": + with sentry_sdk.start_transaction(name="process_reconciler_task", op="deriver"): + try: + validated = ReconcilerPayload(**queue_payload) + except ValidationError as e: + logger.error( + "Invalid reconciler payload received: %s. Payload: %s", + str(e), + queue_payload, + ) + raise ValueError(f"Invalid payload structure: {str(e)}") from e + await process_reconciler(validated) + return + + # All other task types require a workspace_name + if workspace_name is None: + raise ValueError(f"{task_type} tasks require a workspace_name") + if task_type == "webhook": try: validated = WebhookPayload(**queue_payload) @@ -122,19 +141,6 @@ async def process_item(queue_item: models.QueueItem) -> None: raise ValueError(f"Invalid payload structure: {str(e)}") from e await process_deletion(validated, workspace_name) - elif task_type == "reconciler": - with sentry_sdk.start_transaction(name="process_reconciler_task", op="deriver"): - try: - validated = ReconcilerPayload(**queue_payload) - except ValidationError as e: - logger.error( - "Invalid reconciler payload received: %s. Payload: %s", - str(e), - queue_payload, - ) - raise ValueError(f"Invalid payload structure: {str(e)}") from e - await process_reconciler(validated) - else: raise ValueError(f"Invalid task type: {task_type}") diff --git a/src/deriver/queue_manager.py b/src/deriver/queue_manager.py index 350ca4d8..3593a27c 100644 --- a/src/deriver/queue_manager.py +++ b/src/deriver/queue_manager.py @@ -468,7 +468,10 @@ class QueueManager: if removed and queue_item_count > 0: # Only publish webhook if we actually removed an active session try: - if work_unit.task_type in ["representation", "summary"]: + if ( + work_unit.task_type in ["representation", "summary"] + and work_unit.workspace_name is not None + ): logger.debug( f"Publishing queue.empty event for {work_unit_key} in workspace {work_unit.workspace_name}" ) @@ -730,7 +733,10 @@ class QueueManager: ) await db.commit() - if work_unit.task_type in ["representation", "summary"]: + if ( + work_unit.task_type in ["representation", "summary"] + and work_unit.workspace_name is not None + ): prometheus.DERIVER_QUEUE_ITEMS_PROCESSED.labels( workspace_name=work_unit.workspace_name, task_type=work_unit.task_type, diff --git a/src/models.py b/src/models.py index 88135391..040ee517 100644 --- a/src/models.py +++ b/src/models.py @@ -478,8 +478,8 @@ class QueueItem(Base): created_at: Mapped[datetime.datetime] = mapped_column( DateTime(timezone=True), server_default=func.now(), index=True ) - workspace_name: Mapped[str] = mapped_column( - ForeignKey("workspaces.name"), nullable=False, index=True + workspace_name: Mapped[str | None] = mapped_column( + ForeignKey("workspaces.name"), nullable=True, index=True ) message_id: Mapped[int | None] = mapped_column( BigInteger, ForeignKey("messages.id"), nullable=True diff --git a/src/reconciler/scheduler.py b/src/reconciler/scheduler.py index 2b1eee0e..d4f314a3 100644 --- a/src/reconciler/scheduler.py +++ b/src/reconciler/scheduler.py @@ -15,9 +15,7 @@ from datetime import datetime, timedelta, timezone import sentry_sdk from pydantic import BaseModel from sqlalchemy import exists, select -from sqlalchemy.dialects.postgresql import insert as pg_insert from sqlalchemy.exc import IntegrityError -from sqlalchemy.ext.asyncio import AsyncSession from src import models from src.config import settings @@ -26,9 +24,6 @@ from src.models import QueueItem logger = logging.getLogger(__name__) -# System workspace used for global reconciler tasks -SYSTEM_WORKSPACE_NAME = "__system__" - class ReconcilerTask(BaseModel): """Definition of a reconciler task.""" @@ -216,9 +211,6 @@ class ReconcilerScheduler: True if a task was enqueued, False if skipped """ async with tracked_db("reconciler_enqueue") as db: - # Ensure the system workspace exists (for FK constraint) - await self._ensure_system_workspace(db) - # Check if task is already in progress in_progress_check = select( exists( @@ -258,7 +250,7 @@ class ReconcilerScheduler: }, session_id=None, task_type="reconciler", - workspace_name=SYSTEM_WORKSPACE_NAME, + workspace_name=None, message_id=None, ) db.add(queue_item) @@ -274,10 +266,3 @@ class ReconcilerScheduler: logger.info("Enqueued reconciler task: %s", task.name) return True - - async def _ensure_system_workspace(self, db: AsyncSession) -> None: - """Ensure the system workspace exists for reconciler tasks.""" - # Use upsert to create workspace if it doesn't exist - stmt = pg_insert(models.Workspace).values(name=SYSTEM_WORKSPACE_NAME) - stmt = stmt.on_conflict_do_nothing(index_elements=["name"]) - await db.execute(stmt) diff --git a/src/schemas.py b/src/schemas.py index 857f8c17..c747d3bc 100644 --- a/src/schemas.py +++ b/src/schemas.py @@ -206,13 +206,6 @@ class WorkspaceCreate(WorkspaceBase): model_config = ConfigDict(populate_by_name=True) # pyright: ignore - @field_validator("name") - @classmethod - def validate_not_reserved(cls, v: str) -> str: - if v == "__system__": - raise ValueError("Workspace name is reserved.") - return v - class WorkspaceGet(WorkspaceBase): filters: dict[str, Any] | None = None diff --git a/src/utils/work_unit.py b/src/utils/work_unit.py index 3726cb2f..952ff12c 100644 --- a/src/utils/work_unit.py +++ b/src/utils/work_unit.py @@ -9,7 +9,7 @@ class ParsedWorkUnit(BaseModel): """Parsed work unit components.""" task_type: str - workspace_name: str + workspace_name: str | None session_name: str | None observer: str | None observed: str | None @@ -149,7 +149,7 @@ def parse_work_unit_key(work_unit_key: str) -> ParsedWorkUnit: ) return ParsedWorkUnit( task_type=task_type, - workspace_name="__system__", + workspace_name=None, session_name=None, observer=None, observed=None, diff --git a/tests/alembic/revisions/test_119a52b73c60_support_external_embeddings.py b/tests/alembic/revisions/test_119a52b73c60_support_external_embeddings.py index 8a88a4a3..d9e9bfe5 100644 --- a/tests/alembic/revisions/test_119a52b73c60_support_external_embeddings.py +++ b/tests/alembic/revisions/test_119a52b73c60_support_external_embeddings.py @@ -38,6 +38,9 @@ def prepare_support_external_embeddings( verifier.assert_column_exists("message_embeddings", "last_sync_at", exists=False) verifier.assert_column_exists("message_embeddings", "sync_attempts", exists=False) + # Queue workspace_name should be NOT NULL before migration + verifier.assert_column_exists("queue", "workspace_name", nullable=False) + # Indexes should not exist verifier.assert_indexes_not_exist(INDEXES) @@ -64,5 +67,8 @@ def verify_support_external_embeddings( verifier.assert_column_exists("message_embeddings", "last_sync_at", nullable=True) verifier.assert_column_exists("message_embeddings", "sync_attempts", nullable=False) + # Queue workspace_name should now be nullable for system-level tasks + verifier.assert_column_exists("queue", "workspace_name", nullable=True) + # All indexes should exist verifier.assert_indexes_exist(INDEXES)