266 lines
8.3 KiB
Python
266 lines
8.3 KiB
Python
"""
|
|
Task backend for distributed task processing.
|
|
|
|
This provides an abstraction for task storage and execution:
|
|
- BrokerTaskBackend: Uses PostgreSQL as broker (production)
|
|
- SyncTaskBackend: Executes tasks immediately (testing/embedded)
|
|
"""
|
|
|
|
import json
|
|
import logging
|
|
from abc import ABC, abstractmethod
|
|
from collections.abc import Awaitable, Callable
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
if TYPE_CHECKING:
|
|
import asyncpg
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def fq_table(table: str, schema: str | None = None) -> str:
|
|
"""Get fully-qualified table name with optional schema prefix."""
|
|
if schema:
|
|
return f'"{schema}".{table}'
|
|
return table
|
|
|
|
|
|
class TaskBackend(ABC):
|
|
"""
|
|
Abstract base class for task execution backends.
|
|
|
|
Implementations must:
|
|
1. Store/publish task events (as serializable dicts)
|
|
2. Execute tasks through a provided executor callback (optional)
|
|
|
|
The backend treats tasks as pure dictionaries that can be serialized
|
|
and stored in the database. The executor (typically MemoryEngine.execute_task)
|
|
receives the dict and routes it to the appropriate handler.
|
|
"""
|
|
|
|
def __init__(self):
|
|
"""Initialize the task backend."""
|
|
self._executor: Callable[[dict[str, Any]], Awaitable[None]] | None = None
|
|
self._initialized = False
|
|
|
|
def set_executor(self, executor: Callable[[dict[str, Any]], Awaitable[None]]):
|
|
"""
|
|
Set the executor callback for processing tasks.
|
|
|
|
Args:
|
|
executor: Async function that takes a task dict and executes it
|
|
"""
|
|
self._executor = executor
|
|
|
|
@abstractmethod
|
|
async def initialize(self):
|
|
"""
|
|
Initialize the backend (e.g., connect to database).
|
|
"""
|
|
pass
|
|
|
|
@abstractmethod
|
|
async def submit_task(self, task_dict: dict[str, Any]):
|
|
"""
|
|
Submit a task for execution.
|
|
|
|
Args:
|
|
task_dict: Task as a dictionary (must be serializable)
|
|
"""
|
|
pass
|
|
|
|
@abstractmethod
|
|
async def shutdown(self):
|
|
"""
|
|
Shutdown the backend gracefully.
|
|
"""
|
|
pass
|
|
|
|
async def _execute_task(self, task_dict: dict[str, Any]):
|
|
"""
|
|
Execute a task through the registered executor.
|
|
|
|
Args:
|
|
task_dict: Task dictionary to execute
|
|
"""
|
|
if self._executor is None:
|
|
task_type = task_dict.get("type", "unknown")
|
|
logger.warning(f"No executor registered, skipping task {task_type}")
|
|
return
|
|
|
|
try:
|
|
await self._executor(task_dict)
|
|
except Exception as e:
|
|
task_type = task_dict.get("type", "unknown")
|
|
logger.error(f"Error executing task {task_type}: {e}")
|
|
import traceback
|
|
|
|
traceback.print_exc()
|
|
|
|
|
|
class SyncTaskBackend(TaskBackend):
|
|
"""
|
|
Synchronous task backend that executes tasks immediately.
|
|
|
|
This is useful for tests and embedded/CLI usage where we don't want
|
|
background workers. Tasks are executed inline rather than being queued.
|
|
"""
|
|
|
|
async def initialize(self):
|
|
"""No-op for sync backend."""
|
|
self._initialized = True
|
|
logger.debug("SyncTaskBackend initialized")
|
|
|
|
async def submit_task(self, task_dict: dict[str, Any]):
|
|
"""
|
|
Execute the task immediately (synchronously).
|
|
|
|
Args:
|
|
task_dict: Task dictionary to execute
|
|
"""
|
|
if not self._initialized:
|
|
await self.initialize()
|
|
|
|
await self._execute_task(task_dict)
|
|
|
|
async def shutdown(self):
|
|
"""No-op for sync backend."""
|
|
self._initialized = False
|
|
logger.debug("SyncTaskBackend shutdown")
|
|
|
|
|
|
class BrokerTaskBackend(TaskBackend):
|
|
"""
|
|
Task backend using PostgreSQL as broker.
|
|
|
|
submit_task() stores task_payload in async_operations table.
|
|
Actual polling and execution is handled separately by WorkerPoller.
|
|
|
|
This backend is used by the API to store tasks. Workers poll
|
|
the database separately to claim and execute tasks.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
pool_getter: Callable[[], "asyncpg.Pool"],
|
|
schema: str | None = None,
|
|
schema_getter: Callable[[], str | None] | None = None,
|
|
):
|
|
"""
|
|
Initialize the broker task backend.
|
|
|
|
Args:
|
|
pool_getter: Callable that returns the asyncpg connection pool
|
|
schema: Database schema for multi-tenant support (optional, static)
|
|
schema_getter: Callable that returns current schema dynamically (optional).
|
|
If set, takes precedence over static schema for submit_task.
|
|
"""
|
|
super().__init__()
|
|
self._pool_getter = pool_getter
|
|
self._schema = schema
|
|
self._schema_getter = schema_getter
|
|
|
|
async def initialize(self):
|
|
"""Initialize the backend."""
|
|
self._initialized = True
|
|
logger.info("BrokerTaskBackend initialized")
|
|
|
|
async def submit_task(self, task_dict: dict[str, Any]):
|
|
"""
|
|
Store task payload in async_operations table.
|
|
|
|
The task_dict should contain an 'operation_id' if updating an existing
|
|
operation record, otherwise a new operation will be created.
|
|
|
|
Args:
|
|
task_dict: Task dictionary to store (must be JSON serializable)
|
|
"""
|
|
if not self._initialized:
|
|
await self.initialize()
|
|
|
|
pool = self._pool_getter()
|
|
operation_id = task_dict.get("operation_id")
|
|
task_type = task_dict.get("type", "unknown")
|
|
bank_id = task_dict.get("bank_id")
|
|
|
|
# Custom encoder to handle datetime objects
|
|
from datetime import datetime
|
|
|
|
def datetime_encoder(obj):
|
|
if isinstance(obj, datetime):
|
|
return obj.isoformat()
|
|
raise TypeError(f"Object of type {type(obj).__name__} is not JSON serializable")
|
|
|
|
payload_json = json.dumps(task_dict, default=datetime_encoder)
|
|
|
|
schema = self._schema_getter() if self._schema_getter else self._schema
|
|
table = fq_table("async_operations", schema)
|
|
|
|
if operation_id:
|
|
# Update existing operation with task payload
|
|
await pool.execute(
|
|
f"""
|
|
UPDATE {table}
|
|
SET task_payload = $1::jsonb, updated_at = now()
|
|
WHERE operation_id = $2
|
|
""",
|
|
payload_json,
|
|
operation_id,
|
|
)
|
|
logger.debug(f"Updated task payload for operation {operation_id}")
|
|
else:
|
|
# Insert new operation (for tasks without pre-created records)
|
|
# e.g., access_count_update tasks
|
|
import uuid
|
|
|
|
new_id = uuid.uuid4()
|
|
await pool.execute(
|
|
f"""
|
|
INSERT INTO {table} (operation_id, bank_id, operation_type, status, task_payload)
|
|
VALUES ($1, $2, $3, 'pending', $4::jsonb)
|
|
""",
|
|
new_id,
|
|
bank_id,
|
|
task_type,
|
|
payload_json,
|
|
)
|
|
logger.debug(f"Created new operation {new_id} for task type {task_type}")
|
|
|
|
async def shutdown(self):
|
|
"""Shutdown the backend."""
|
|
self._initialized = False
|
|
logger.info("BrokerTaskBackend shutdown")
|
|
|
|
async def wait_for_pending_tasks(self, timeout: float = 120.0):
|
|
"""
|
|
Wait for pending tasks to be processed.
|
|
|
|
In the broker model, this polls the database to check if tasks
|
|
for this process have been completed. This is useful in tests
|
|
when worker_enabled=True (API processes its own tasks).
|
|
|
|
Args:
|
|
timeout: Maximum time to wait in seconds
|
|
"""
|
|
import asyncio
|
|
|
|
pool = self._pool_getter()
|
|
schema = self._schema_getter() if self._schema_getter else self._schema
|
|
table = fq_table("async_operations", schema)
|
|
|
|
start_time = asyncio.get_event_loop().time()
|
|
while asyncio.get_event_loop().time() - start_time < timeout:
|
|
# Check if there are any pending tasks with payloads
|
|
count = await pool.fetchval(
|
|
f"""
|
|
SELECT COUNT(*) FROM {table}
|
|
WHERE status = 'pending' AND task_payload IS NOT NULL
|
|
"""
|
|
)
|
|
|
|
if count == 0:
|
|
return
|
|
|
|
await asyncio.sleep(0.5)
|
|
|
|
logger.warning(f"Timeout waiting for pending tasks after {timeout}s")
|