From 657fe023b2b939fcacff0edcd87373360bd310da Mon Sep 17 00:00:00 2001
From: Chris Bartholomew <chris.bartholomew@vectorize.io>
Date: Thu, 29 Jan 2026 15:03:20 -0500
Subject: [PATCH] fix: run migrations on tenant schemas at startup and harden
 worker poller (#237)

Tenant schemas were never migrated when new migrations were deployed.
Only the public schema was migrated at startup, and tenant schemas only
got migrations when first provisioned. This meant existing tenants
missed any new columns (e.g. task_payload, worker_id, claimed_at on
async_operations), causing the worker poller to crash silently.

Changes:
- Run migrations on all existing tenant schemas at startup when a
  tenant_extension is configured. Each schema migration is wrapped in
  try/except so one failure doesn't block others.
- Add try/except in WorkerPoller.recover_own_tasks() so a broken
  schema doesn't prevent the polling loop from starting.
- Add try/except in WorkerPoller._claim_batch_for_schema() so a
  broken schema doesn't prevent claiming tasks from other schemas.
---
 .../hindsight_api/engine/memory_engine.py     | 17 +++++++++
 hindsight-api/hindsight_api/worker/poller.py  | 35 ++++++++++++-------
 2 files changed, 40 insertions(+), 12 deletions(-)

diff --git a/hindsight-api/hindsight_api/engine/memory_engine.py b/hindsight-api/hindsight_api/engine/memory_engine.py
index f8e1ebc8..31b9b7cc 100644
--- a/hindsight-api/hindsight_api/engine/memory_engine.py
+++ b/hindsight-api/hindsight_api/engine/memory_engine.py
@@ -889,6 +889,23 @@ class MemoryEngine(MemoryEngineInterface):
             # Use configured database schema for migrations (defaults to "public")
             run_migrations(self.db_url, schema=get_config().database_schema)
 
+            # Migrate all existing tenant schemas (if multi-tenant)
+            if self._tenant_extension is not None:
+                try:
+                    tenants = await self._tenant_extension.list_tenants()
+                    if tenants:
+                        logger.info(f"Running migrations on {len(tenants)} tenant schemas...")
+                        for tenant in tenants:
+                            schema = tenant.schema
+                            if schema and schema != "public":
+                                try:
+                                    run_migrations(self.db_url, schema=schema)
+                                except Exception as e:
+                                    logger.warning(f"Failed to migrate tenant schema {schema}: {e}")
+                        logger.info("Tenant schema migrations completed")
+                except Exception as e:
+                    logger.warning(f"Failed to run tenant schema migrations: {e}")
+
             # Ensure embedding column dimension matches the model's dimension
             # This is done after migrations and after embeddings.initialize()
             ensure_embedding_dimension(self.db_url, self.embeddings.dimension, schema=get_config().database_schema)
diff --git a/hindsight-api/hindsight_api/worker/poller.py b/hindsight-api/hindsight_api/worker/poller.py
index a0e77ec1..3c990c56 100644
--- a/hindsight-api/hindsight_api/worker/poller.py
+++ b/hindsight-api/hindsight_api/worker/poller.py
@@ -132,6 +132,14 @@ class WorkerPoller:
 
     async def _claim_batch_for_schema(self, schema: str | None, limit: int) -> list[ClaimedTask]:
         """Claim tasks from a specific schema."""
+        try:
+            return await self._claim_batch_for_schema_inner(schema, limit)
+        except Exception as e:
+            logger.warning(f"Worker {self._worker_id} failed to claim tasks for schema {schema or 'public'}: {e}")
+            return []
+
+    async def _claim_batch_for_schema_inner(self, schema: str | None, limit: int) -> list[ClaimedTask]:
+        """Inner implementation for claiming tasks from a specific schema."""
         table = fq_table("async_operations", schema)
 
         async with self._pool.acquire() as conn:
@@ -293,20 +301,23 @@ class WorkerPoller:
         total_count = 0
 
         for schema in schemas:
-            table = fq_table("async_operations", schema)
+            try:
+                table = fq_table("async_operations", schema)
 
-            result = await self._pool.execute(
-                f"""
-                UPDATE {table}
-                SET status = 'pending', worker_id = NULL, claimed_at = NULL, updated_at = now()
-                WHERE status = 'processing' AND worker_id = $1
-                """,
-                self._worker_id,
-            )
+                result = await self._pool.execute(
+                    f"""
+                    UPDATE {table}
+                    SET status = 'pending', worker_id = NULL, claimed_at = NULL, updated_at = now()
+                    WHERE status = 'processing' AND worker_id = $1
+                    """,
+                    self._worker_id,
+                )
 
-            # Parse "UPDATE N" to get count
-            count = int(result.split()[-1]) if result else 0
-            total_count += count
+                # Parse "UPDATE N" to get count
+                count = int(result.split()[-1]) if result else 0
+                total_count += count
+            except Exception as e:
+                logger.warning(f"Worker {self._worker_id} failed to recover tasks for schema {schema or 'public'}: {e}")
 
         if total_count > 0:
             logger.info(f"Worker {self._worker_id} recovered {total_count} stale tasks from previous run")