digitalkin 0.2.26__py3-none-any.whl → 0.3.0.dev1__py3-none-any.whl

digitalkin/modules/job_manager/surrealdb_repository.py

@@ -0,0 +1,228 @@
+"""SurrealDB connection management."""
+
+import datetime
+import os
+from collections.abc import AsyncGenerator
+from typing import Any, Generic, TypeVar
+from uuid import UUID
+
+from surrealdb import AsyncHttpSurrealConnection, AsyncSurreal, AsyncWsSurrealConnection, RecordID
+
+from digitalkin.logger import logger
+
+TSurreal = TypeVar("TSurreal", bound=AsyncHttpSurrealConnection | AsyncWsSurrealConnection)
+
+
+class SurrealDBSetupBadIDError(Exception):
+    """Exception raised when an invalid ID is encountered during the setup process in the SurrealDB repository.
+
+    This error is used to indicate that the provided ID does not meet the
+    expected format or criteria.
+    """
+
+
+class SurrealDBSetupVersionBadIDError(Exception):
+    """Exception raised when an invalid ID is encountered during the setup of a SurrealDB version.
+
+    This error is intended to signal that the provided ID does not meet
+    the expected format or criteria for a valid SurrealDB setup version ID.
+    """
+
+
+class SurrealDBConnection(Generic[TSurreal]):
+    """Base repository for database operations.
+
+    This class provides common database operations that can be used by
+    specific table repositories.
+    """
+
+    db: TSurreal
+    timeout: datetime.timedelta
+
+    @staticmethod
+    def _valid_id(raw_id: str, table_name: str) -> RecordID:
+        """Validate and parse a raw ID string into a RecordID.
+
+        Args:
+            raw_id: The raw ID string to validate
+            table_name: table name to enforce
+
+        Raises:
+            SurrealDBSetupBadIDError: If the raw ID string is not valid
+
+        Returns:
+            RecordID: Parsed RecordID object if valid, None otherwise
+        """
+        try:
+            split_id = raw_id.split(":")
+            if split_id[0] != table_name:
+                msg = f"Invalid table name for ID: {raw_id}"
+                raise SurrealDBSetupBadIDError(msg)
+            return RecordID(split_id[0], split_id[1])
+        except IndexError:
+            raise SurrealDBSetupBadIDError
+
+    def __init__(
+        self,
+        database: str | None = None,
+        timeout: datetime.timedelta = datetime.timedelta(seconds=5),
+    ) -> None:
+        """Initialize the repository.
+
+        Args:
+            database: AsyncSurrealDB connection to a specific database
+            table_name: Name of the table to interact with
+            timeout: Timeout for database operations
+        """
+        self.timeout = timeout
+        base_url = os.getenv("SURREALDB_URL", "ws://localhost").strip()
+        port = (os.getenv("SURREALDB_PORT") or "").strip()
+        self.url = f"{base_url}{f':{port}' if port else ''}/rpc"
+
+        self.username = os.getenv("SURREALDB_USERNAME", "root")
+        self.password = os.getenv("SURREALDB_PASSWORD", "root")
+        self.namespace = os.getenv("SURREALDB_NAMESPACE", "test")
+        self.database = database or os.getenv("SURREALDB_DATABASE", "task_manager")
+
+    async def init_surreal_instance(self) -> None:
+        """Init a SurrealDB connection instance."""
+        logger.debug("Connecting to SurrealDB at %s", self.url)
+        self.db = AsyncSurreal(self.url)  # type: ignore
+        await self.db.signin({"username": self.username, "password": self.password})
+        await self.db.use(self.namespace, self.database)
+        logger.debug("Successfully connected to SurrealDB")
+
+    async def close(self) -> None:
+        """Close the SurrealDB connection if it exists."""
+        logger.debug("Closing SurrealDB connection")
+        await self.db.close()
+
+    async def create(
+        self,
+        table_name: str,
+        data: dict[str, Any],
+    ) -> list[dict[str, Any]] | dict[str, Any]:
+        """Create a new record.
+
+        Args:
+            table_name: Name of the table to insert into
+            data: Data to insert
+
+        Returns:
+            Dict[str, Any]: The created record as returned by the database
+        """
+        logger.debug("Creating record in %s with data: %s", table_name, data)
+        result = await self.db.create(table_name, data)
+        logger.debug("create result: %s", result)
+        return result
+
+    async def merge(
+        self,
+        table_name: str,
+        record_id: str | RecordID,
+        data: dict[str, Any],
+    ) -> list[dict[str, Any]] | dict[str, Any]:
+        """Update an existing record.
+
+        Args:
+            table_name: Name of the table to insert into
+            record_id: record ID to update
+            data: Data to insert
+
+        Returns:
+            Dict[str, Any]: The created record as returned by the database
+        """
+        if isinstance(record_id, str):
+            # validate surrealDB id if raw str
+            record_id = self._valid_id(record_id, table_name)
+        logger.debug("Updating record in %s with data: %s", record_id, data)
+        result = await self.db.merge(record_id, data)
+        logger.debug("update result: %s", result)
+        return result
+
+    async def update(
+        self,
+        table_name: str,
+        record_id: str | RecordID,
+        data: dict[str, Any],
+    ) -> list[dict[str, Any]] | dict[str, Any]:
+        """Update an existing record.
+
+        Args:
+            table_name: Name of the table to insert into
+            record_id: record ID to update
+            data: Data to insert
+
+        Returns:
+            Dict[str, Any]: The created record as returned by the database
+        """
+        if isinstance(record_id, str):
+            # validate surrealDB id if raw str
+            record_id = self._valid_id(record_id, table_name)
+        logger.debug("Updating record in %s with data: %s", record_id, data)
+        result = await self.db.update(record_id, data)
+        logger.debug("update result: %s", result)
+        return result
+
+    async def execute_query(self, query: str, params: dict[str, Any] | None = None) -> list[dict[str, Any]]:
+        """Execute a custom SurrealQL query.
+
+        Args:
+            query: SurrealQL query
+            params: Query parameters
+
+        Returns:
+            List[Dict[str, Any]]: Query results
+        """
+        logger.debug("execute_query: %s with params: %s", query, params)
+        result = await self.db.query(query, params or {})
+        logger.debug("execute_query result: %s", result)
+        return [result] if isinstance(result, dict) else result
+
+    async def select_by_task_id(self, table: str, value: str) -> dict[str, Any]:
+        """Fetch a record from a table by a unique field.
+
+        Args:
+            table: Table name
+            value: Field value to match
+
+        Raises:
+            ValueError: If no records are found
+
+        Returns:
+            Dict with record data if found, else None
+        """
+        query = "SELECT * FROM type::table($table) WHERE task_id = $value;"
+        params = {"table": table, "value": value}
+
+        result = await self.execute_query(query, params)
+        if not result:
+            msg = f"No records found in table '{table}' with task_id '{value}'"
+            logger.error(msg)
+            raise ValueError(msg)
+
+        return result[0]
+
+    async def start_live(
+        self,
+        table_name: str,
+    ) -> tuple[UUID, AsyncGenerator[dict[str, Any], None]]:
+        """Create and subscribe to a live SurrealQL query.
+
+        Args:
+            table_name: Name of the table to insert into
+
+        Returns:
+            List[Dict[str, Any]]: Query results
+        """
+        live_id = await self.db.live(table_name, diff=False)
+        return live_id, await self.db.subscribe_live(live_id)
+
+    async def stop_live(self, live_id: UUID) -> None:
+        """Kill a live SurrealQL query.
+
+        Args:
+            live_id: record ID to watch for
+        """
+        logger.debug("KILL Subscribe live for: %s", live_id)
+        await self.db.kill(live_id)

digitalkin/modules/job_manager/task_manager.py

@@ -0,0 +1,389 @@
+"""Task manager with comprehensive lifecycle management."""
+
+import asyncio
+import contextlib
+import datetime
+from collections.abc import Coroutine
+from typing import Any
+
+from digitalkin.logger import logger
+from digitalkin.models.module.task_monitor import SignalMessage, SignalType, TaskStatus
+from digitalkin.modules._base_module import BaseModule
+from digitalkin.modules.job_manager.task_session import SurrealDBConnection, TaskSession
+
+
+class TaskManager:
+    """Task manager with comprehensive lifecycle management."""
+
+    tasks: dict[str, asyncio.Task]
+    tasks_sessions: dict[str, TaskSession]
+    channel: SurrealDBConnection
+    default_timeout: float
+    max_concurrent_tasks: int
+    _shutdown_event: asyncio.Event
+
+    def __init__(self, default_timeout: float = 10.0, max_concurrent_tasks: int = 100) -> None:
+        """."""
+        self.tasks = {}
+        self.tasks_sessions = {}
+        self.default_timeout = default_timeout
+        self.max_concurrent_tasks = max_concurrent_tasks
+        self._shutdown_event = asyncio.Event()
+
+        logger.info(
+            "TaskManager initialized with max_concurrent_tasks: %d, default_timeout: %.1f",
+            max_concurrent_tasks,
+            default_timeout,
+            extra={"max_concurrent_tasks": max_concurrent_tasks, "default_timeout": default_timeout},
+        )
+
+    @property
+    def task_count(self) -> int:
+        """."""
+        return len(self.tasks_sessions)
+
+    @property
+    def running_tasks(self) -> set[str]:
+        """."""
+        return {task_id for task_id, task in self.tasks.items() if not task.done()}
+
+    async def _cleanup_task(self, task_id: str) -> None:
+        """Clean up task resources."""
+        logger.debug("Cleaning up resources for task: '%s'", task_id, extra={"task_id": task_id})
+        if task_id in self.tasks_sessions:
+            await self.tasks_sessions[task_id].db.close()
+        # Remove from collections
+
+    async def _task_wrapper(  # noqa: C901, PLR0915
+        self,
+        task_id: str,
+        coro: Coroutine[Any, Any, None],
+        session: TaskSession,
+    ) -> asyncio.Task[None]:
+        """Task wrapper that runs main, heartbeat, and listener concurrently.
+
+        The first to finish determines the outcome. Returns a Task that the
+        caller can await externally.
+
+        Returns:
+            asyncio.Task[None]: The supervisor task managing the lifecycle.
+        """
+
+        async def signal_wrapper() -> None:
+            try:
+                await self.channel.create(
+                    "tasks",
+                    SignalMessage(
+                        task_id=task_id,
+                        status=session.status,
+                        action=SignalType.START,
+                    ).model_dump(),
+                )
+                await session.listen_signals()
+            except asyncio.CancelledError:
+                logger.debug("Signal listener cancelled", extra={"task_id": task_id})
+            finally:
+                await self.channel.create(
+                    "tasks",
+                    SignalMessage(
+                        task_id=task_id,
+                        status=session.status,
+                        action=SignalType.STOP,
+                    ).model_dump(),
+                )
+                logger.info("Signal listener ended", extra={"task_id": task_id})
+
+        async def heartbeat_wrapper() -> None:
+            try:
+                await session.generate_heartbeats()
+            except asyncio.CancelledError:
+                logger.debug("Signal listener cancelled", extra={"task_id": task_id})
+            finally:
+                logger.info("Heartbeat task ended", extra={"task_id": task_id})
+
+        async def supervisor() -> None:
+            session.started_at = datetime.datetime.now(datetime.timezone.utc)
+            session.status = TaskStatus.RUNNING
+
+            main_task = asyncio.create_task(coro, name=f"{task_id}_main")
+            hb_task = asyncio.create_task(heartbeat_wrapper(), name=f"{task_id}_heartbeat")
+            sig_task = asyncio.create_task(signal_wrapper(), name=f"{task_id}_listener")
+
+            try:
+                done, pending = await asyncio.wait(
+                    [main_task, sig_task, hb_task],
+                    return_when=asyncio.FIRST_COMPLETED,
+                )
+
+                # One task completed -> cancel the others
+                for t in pending:
+                    t.cancel()
+
+                # Propagate exception/result from the finished task
+                completed = next(iter(done))
+                await completed
+
+                if completed is main_task:
+                    session.status = TaskStatus.COMPLETED
+                elif completed is sig_task or (completed is hb_task and sig_task.done()):
+                    logger.debug(f"Task cancelled due to signal {sig_task=}")
+                    session.status = TaskStatus.CANCELLED
+                elif completed is hb_task:
+                    session.status = TaskStatus.FAILED
+                    msg = f"Heartbeat stopped for {task_id}"
+                    raise RuntimeError(msg)  # noqa: TRY301
+
+            except asyncio.CancelledError:
+                session.status = TaskStatus.CANCELLED
+                raise
+            except Exception:
+                session.status = TaskStatus.FAILED
+                raise
+            finally:
+                session.completed_at = datetime.datetime.now(datetime.timezone.utc)
+                # Ensure all tasks are cleaned up
+                for t in [main_task, hb_task, sig_task]:
+                    if not t.done():
+                        t.cancel()
+                await asyncio.gather(main_task, hb_task, sig_task, return_exceptions=True)
+
+        # Return the supervisor task to be awaited outside
+        return asyncio.create_task(supervisor(), name=f"{task_id}_supervisor")
+
+    async def create_task(
+        self,
+        task_id: str,
+        module: BaseModule,
+        coro: Coroutine[Any, Any, None],
+        heartbeat_interval: datetime.timedelta = datetime.timedelta(seconds=2),
+        connection_timeout: datetime.timedelta = datetime.timedelta(seconds=5),
+    ) -> None:
+        """Create and start a new managed task.
+
+        Raises:
+            ValueError: task_id duplicated
+            RuntimeError: task overload
+        """
+        if task_id in self.tasks:
+            # close Coroutine during runtime
+            coro.close()
+            logger.warning("Task creation failed - task already exists: '%s'", task_id, extra={"task_id": task_id})
+            msg = f"Task {task_id} already exists"
+            raise ValueError(msg)
+
+        if len(self.tasks) >= self.max_concurrent_tasks:
+            coro.close()
+            logger.error(
+                "Task creation failed - max concurrent tasks reached: %d",
+                self.max_concurrent_tasks,
+                extra={
+                    "task_id": task_id,
+                    "current_count": len(self.tasks),
+                    "max_concurrent": self.max_concurrent_tasks,
+                },
+            )
+            msg = f"Maximum concurrent tasks ({self.max_concurrent_tasks}) reached"
+            raise RuntimeError(msg)
+
+        logger.info(
+            "Creating new task: '%s'",
+            task_id,
+            extra={
+                "task_id": task_id,
+                "heartbeat_interval": heartbeat_interval,
+                "connection_timeout": connection_timeout,
+            },
+        )
+
+        try:
+            # Initialize components
+            channel: SurrealDBConnection = SurrealDBConnection("task_manager", connection_timeout)
+            await channel.init_surreal_instance()
+            session = TaskSession(task_id, channel, module, heartbeat_interval)
+
+            self.tasks_sessions[task_id] = session
+
+            # Create wrapper task
+            self.tasks[task_id] = asyncio.create_task(self._task_wrapper(task_id, coro, session), name=task_id)
+
+            logger.info(
+                "Task created successfully: '%s'",
+                task_id,
+                extra={
+                    "task_id": task_id,
+                    "total_tasks": len(self.tasks),
+                },
+            )
+
+        except Exception as e:
+            logger.error(
+                "Failed to create task: '%s'", task_id, extra={"task_id": task_id, "error": str(e)}, exc_info=True
+            )
+            # Cleanup on failure
+            await self._cleanup_task(task_id)
+            raise
+
+    async def send_signal(self, task_id: str, signal_type: str, payload: dict) -> bool:
+        """Send signal to a specific task.
+
+        Returns:
+            bool: True if the task sent successfully the given signal, False otherwise.
+        """
+        if task_id not in self.tasks_sessions:
+            logger.warning(
+                "Cannot send signal - task not found: '%s'",
+                task_id,
+                extra={"task_id": task_id, "signal_type": signal_type},
+            )
+            return False
+
+        logger.info(
+            "Sending signal '%s' to task: '%s'",
+            signal_type,
+            task_id,
+            extra={"task_id": task_id, "signal_type": signal_type, "payload": payload},
+        )
+
+        await self.channel.update("tasks", signal_type, payload)
+        return True
+
+    async def cancel_task(self, task_id: str, timeout: float | None = None) -> bool:
+        """Cancel a task with graceful shutdown and fallback.
+
+        Returns:
+            bool: True if the task was cancelled successfully, False otherwise.
+        """
+        if task_id not in self.tasks:
+            logger.warning("Cannot cancel - task not found: '%s'", task_id, extra={"task_id": task_id})
+            return True
+
+        timeout = timeout or self.default_timeout
+        task = self.tasks[task_id]
+
+        logger.info(
+            "Initiating task cancellation: '%s', timeout: %.1fs",
+            task_id,
+            timeout,
+            extra={"task_id": task_id, "timeout": timeout},
+        )
+
+        try:
+            # Phase 1: Cooperative cancellation
+            # await self.send_signal(task_id, "cancel")  # noqa: ERA001
+
+            # Wait for graceful shutdown
+            await asyncio.wait_for(task, timeout=timeout)
+
+            logger.info("Task cancelled gracefully: '%s'", task_id, extra={"task_id": task_id})
+
+        except asyncio.TimeoutError:
+            logger.warning(
+                "Graceful cancellation timed out for task: '%s', forcing cancellation",
+                task_id,
+                extra={"task_id": task_id, "timeout": timeout},
+            )
+
+            # Phase 2: Force cancellation
+            task.cancel()
+            with contextlib.suppress(asyncio.CancelledError):
+                await task
+
+            logger.warning("Task force-cancelled: '%s'", task_id, extra={"task_id": task_id})
+            return True
+
+        except Exception as e:
+            logger.error(
+                "Error during task cancellation: '%s'",
+                task_id,
+                extra={"task_id": task_id, "error": str(e)},
+                exc_info=True,
+            )
+            return False
+        return True
+
+    async def clean_session(self, task_id: str) -> bool:
+        """Clean up task session without cancelling the task.
+
+        Returns:
+            bool: True if the task was cleaned successfully, False otherwise.
+        """
+        if task_id not in self.tasks_sessions:
+            logger.warning("Cannot clean session - task not found: '%s'", task_id, extra={"task_id": task_id})
+            return False
+
+        await self.tasks_sessions[task_id].module.stop()
+        await self.cancel_task(task_id)
+
+        logger.info("Cleaning up session for task: '%s'", task_id, extra={"task_id": task_id})
+        self.tasks_sessions.pop(task_id, None)
+        return True
+
+    async def pause_task(self, task_id: str) -> bool:
+        """Pause a running task.
+
+        Returns:
+            bool: True if the task was paused successfully, False otherwise.
+        """
+        return await self.send_signal(task_id, "pause", {})
+
+    async def resume_task(self, task_id: str) -> bool:
+        """Resume a paused task.
+
+        Returns:
+            bool: True if the task was paused successfully, False otherwise.
+        """
+        return await self.send_signal(task_id, "resume", {})
+
+    async def get_task_status(self, task_id: str) -> bool:
+        """Request status from a task.
+
+        Returns:
+            bool: True if the task was paused successfully, False otherwise.
+        """
+        return await self.send_signal(task_id, "status", {})
+
+    async def cancel_all_tasks(self, timeout: float | None = None) -> dict[str, bool]:
+        """Cancel all running tasks.
+
+        Returns:
+            dict[str: bool]: True if the tasks were paused successfully, False otherwise.
+        """
+        timeout = timeout or self.default_timeout
+        task_ids = list(self.running_tasks)
+
+        logger.info(
+            "Cancelling all tasks: %d tasks", len(task_ids), extra={"task_count": len(task_ids), "timeout": timeout}
+        )
+
+        results = {}
+        for task_id in task_ids:
+            results[task_id] = await self.cancel_task(task_id, timeout)
+
+        return results
+
+    async def shutdown(self, timeout: float = 30.0) -> None:
+        """Graceful shutdown of all tasks."""
+        logger.info(
+            "TaskManager shutdown initiated, timeout: %.1fs",
+            timeout,
+            extra={"timeout": timeout, "active_tasks": len(self.running_tasks)},
+        )
+
+        self._shutdown_event.set()
+        results = await self.cancel_all_tasks(timeout)
+
+        failed_tasks = [task_id for task_id, success in results.items() if not success]
+        if failed_tasks:
+            logger.error(
+                "Failed to cancel %d tasks during shutdown: %s",
+                len(failed_tasks),
+                failed_tasks,
+                extra={"failed_tasks": failed_tasks, "failed_count": len(failed_tasks)},
+            )
+
+        logger.info(
+            "TaskManager shutdown completed, cancelled: %d, failed: %d",
+            len(results) - len(failed_tasks),
+            len(failed_tasks),
+            extra={"cancelled_count": len(results) - len(failed_tasks), "failed_count": len(failed_tasks)},
+        )