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.
This commit is contained in:
parent
9c95a1ac1d
commit
657fe023b2
2 changed files with 40 additions and 12 deletions
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
Loading…
Reference in a new issue