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:
Chris Bartholomew 2026-01-29 15:03:20 -05:00 committed by GitHub
parent 9c95a1ac1d
commit 657fe023b2
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 40 additions and 12 deletions

View file

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

View file

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