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")