fleet-memory/hindsight-api/hindsight_api/engine/task_backend.py

"""
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")
        payload_json = json.dumps(task_dict)

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