digitalkin 0.2.25rc1__py3-none-any.whl → 0.3.0__py3-none-any.whl

digitalkin/modules/job_manager/single_job_manager.py

@@ -1,6 +1,7 @@
 """Background module manager with single instance."""
 
 import asyncio
+import datetime
 import uuid
 from collections.abc import AsyncGenerator, AsyncIterator
 from contextlib import asynccontextmanager
@@ -9,15 +10,18 @@ from typing import Any, Generic
 import grpc
 
 from digitalkin.logger import logger
-from digitalkin.models import ModuleStatus
 from digitalkin.models.module import InputModelT, OutputModelT, SetupModelT
+from digitalkin.models.module.module import ModuleCodeModel
+from digitalkin.models.module.task_monitor import TaskStatus
 from digitalkin.modules._base_module import BaseModule
 from digitalkin.modules.job_manager.base_job_manager import BaseJobManager
-from digitalkin.modules.job_manager.job_manager_models import StreamCodeModel
+from digitalkin.modules.job_manager.surrealdb_repository import SurrealDBConnection
+from digitalkin.modules.job_manager.task_manager import TaskManager
+from digitalkin.modules.job_manager.task_session import TaskSession
 from digitalkin.services.services_models import ServicesMode
 
 
-class SingleJobManager(BaseJobManager, Generic[InputModelT, SetupModelT]):
+class SingleJobManager(BaseJobManager, TaskManager, Generic[InputModelT, OutputModelT, SetupModelT]):
     """Manages a single instance of a module job.
 
     This class ensures that only one instance of a module job is active at a time.
@@ -25,8 +29,10 @@ class SingleJobManager(BaseJobManager, Generic[InputModelT, SetupModelT]):
     to handle their output data.
     """
 
-    modules: dict[str, BaseModule]
-    queue: dict[str, asyncio.Queue]
+    async def start(self) -> None:
+        """Start manager."""
+        self.channel = SurrealDBConnection("task_manager", datetime.timedelta(seconds=5))
+        await self.channel.init_surreal_instance()
 
     def __init__(
         self,
@@ -40,12 +46,9 @@ class SingleJobManager(BaseJobManager, Generic[InputModelT, SetupModelT]):
             services_mode: The mode of operation for the services (e.g., ASYNC or SYNC).
         """
         super().__init__(module_class, services_mode)
-
         self._lock = asyncio.Lock()
-        self.modules: dict[str, BaseModule] = {}
-        self.queues: dict[str, asyncio.Queue] = {}
 
-    async def generate_config_setup_module_response(self, job_id: str) -> SetupModelT:
+    async def generate_config_setup_module_response(self, job_id: str) -> SetupModelT | ModuleCodeModel:
         """Generate a stream consumer for a module's output data.
 
         This method creates an asynchronous generator that streams output data
@@ -56,16 +59,19 @@ class SingleJobManager(BaseJobManager, Generic[InputModelT, SetupModelT]):
             job_id: The unique identifier of the job.
 
         Returns:
-            SetupModelT: the SetupModelT object fully processed.
+            SetupModelT | ModuleCodeModel: the SetupModelT object fully processed.
         """
-        module = self.modules.get(job_id, None)
-        logger.debug("Module %s found: %s", job_id, module)
+        if (session := self.tasks_sessions.get(job_id, None)) is None:
+            return ModuleCodeModel(
+                code=str(grpc.StatusCode.NOT_FOUND),
+                message=f"Module {job_id} not found",
+            )
 
+        logger.debug("Module %s found: %s", job_id, session.module)
         try:
-            return await self.queues[job_id].get()
+            return await session.queue.get()
         finally:
-            logger.info(f"{job_id=}: {self.queues[job_id].empty()}")
-            del self.queues[job_id]
+            logger.info(f"{job_id=}: {session.queue.empty()}")
 
     async def create_config_setup_instance_job(
         self,
@@ -95,8 +101,7 @@ class SingleJobManager(BaseJobManager, Generic[InputModelT, SetupModelT]):
         job_id = str(uuid.uuid4())
         # TODO: Ensure the job_id is unique.
         module = self.module_class(job_id, mission_id=mission_id, setup_id=setup_id, setup_version_id=setup_version_id)
-        self.modules[job_id] = module
-        self.queues[job_id] = asyncio.Queue()
+        self.tasks_sessions[job_id] = TaskSession(job_id, self.channel, module)
 
         try:
             await module.start_config_setup(
@@ -106,13 +111,13 @@ class SingleJobManager(BaseJobManager, Generic[InputModelT, SetupModelT]):
             logger.debug("Module %s (%s) started successfully", job_id, module.name)
         except Exception:
             # Remove the module from the manager in case of an error.
-            del self.modules[job_id]
+            del self.tasks_sessions[job_id]
             logger.exception("Failed to start module %s: %s", job_id)
             raise
         else:
             return job_id
 
-    async def add_to_queue(self, job_id: str, output_data: OutputModelT) -> None:  # type: ignore
+    async def add_to_queue(self, job_id: str, output_data: OutputModelT | ModuleCodeModel) -> None:
         """Add output data to the queue for a specific job.
 
         This method is used as a callback to handle output data generated by a module job.
@@ -121,7 +126,7 @@ class SingleJobManager(BaseJobManager, Generic[InputModelT, SetupModelT]):
             job_id: The unique identifier of the job.
             output_data: The output data produced by the job.
         """
-        await self.queues[job_id].put(output_data.model_dump())
+        await self.tasks_sessions[job_id].queue.put(output_data.model_dump())
 
     @asynccontextmanager  # type: ignore
     async def generate_stream_consumer(self, job_id: str) -> AsyncIterator[AsyncGenerator[dict[str, Any], None]]:  # type: ignore
@@ -137,39 +142,48 @@ class SingleJobManager(BaseJobManager, Generic[InputModelT, SetupModelT]):
         Yields:
             AsyncGenerator: A stream of output data or error messages.
         """
-        module = self.modules.get(job_id, None)
-
-        logger.debug("Module %s found: %s", job_id, module)
+        if (session := self.tasks_sessions.get(job_id, None)) is None:
 
-        async def _stream() -> AsyncGenerator[dict[str, Any], Any]:
-            """Stream output data from the module.
+            async def _error_gen() -> AsyncGenerator[dict[str, Any], None]:  # noqa: RUF029
+                """Generate an error message for a non-existent module.
 
-            Yields:
-                dict: Output data generated by the module.
-            """
-            if module is None:
+                Yields:
+                    AsyncGenerator: A generator yielding an error message.
+                """
                 yield {
                     "error": {
                         "error_message": f"Module {job_id} not found",
                         "code": grpc.StatusCode.NOT_FOUND,
                     }
                 }
-                return
 
-            try:
-                while module.status == ModuleStatus.RUNNING or (
-                    not self.queues[job_id].empty()
-                    and module.status
-                    in {
-                        ModuleStatus.STOPPED,
-                        ModuleStatus.STOPPING,
-                    }
-                ):
-                    logger.info(f"{job_id=}: {module.status=}")
-                    yield await self.queues[job_id].get()
+            yield _error_gen()
+            return
+
+        logger.debug("Session: %s with Module %s", job_id, session.module)
+
+        async def _stream() -> AsyncGenerator[dict[str, Any], Any]:
+            """Stream output data from the module.
 
-            finally:
-                del self.queues[job_id]
+            Yields:
+                dict: Output data generated by the module.
+            """
+            while True:
+                # if queue is empty but producer not finished yet, block on get()
+                msg = await session.queue.get()
+                try:
+                    yield msg
+                finally:
+                    session.queue.task_done()
+
+                # If the producer marked finished and no more items, break soon:
+                if (
+                    session.is_cancelled.is_set()
+                    or (session.status is TaskStatus.COMPLETED and session.queue.empty())
+                    or session.status is TaskStatus.FAILED
+                ):
+                    # and session.queue.empty():
+                    break
 
         yield _stream()
 
@@ -200,32 +214,21 @@ class SingleJobManager(BaseJobManager, Generic[InputModelT, SetupModelT]):
             Exception: If the module fails to start.
         """
         job_id = str(uuid.uuid4())
-        # TODO: Ensure the job_id is unique.
         module = self.module_class(
             job_id,
             mission_id=mission_id,
             setup_id=setup_id,
             setup_version_id=setup_version_id,
         )
-        self.modules[job_id] = module
-        self.queues[job_id] = asyncio.Queue()
         callback = await self.job_specific_callback(self.add_to_queue, job_id)
 
-        try:
-            await module.start(
-                input_data,
-                setup_data,
-                callback,
-                done_callback=lambda _: asyncio.create_task(callback(StreamCodeModel(code="__END_OF_STREAM__"))),
-            )
-            logger.debug("Module %s (%s) started successfully", job_id, module.name)
-        except Exception:
-            # Remove the module from the manager in case of an error.
-            del self.modules[job_id]
-            logger.exception("Failed to start module %s: %s", job_id)
-            raise
-        else:
-            return job_id
+        await self.create_task(
+            job_id,
+            module,
+            module.start(input_data, setup_data, callback, done_callback=None),
+        )
+        logger.info("Managed task started: '%s'", job_id, extra={"task_id": job_id})
+        return job_id
 
     async def stop_module(self, job_id: str) -> bool:
         """Stop a running module job.
@@ -239,34 +242,37 @@ class SingleJobManager(BaseJobManager, Generic[InputModelT, SetupModelT]):
         Raises:
             Exception: If an error occurs while stopping the module.
         """
+        logger.critical(f"STOP {job_id=} | {self.tasks_sessions.keys()}")
+
         async with self._lock:
-            module = self.modules.get(job_id)
-            if not module:
-                logger.warning(f"Module {job_id} not found")
+            session = self.tasks_sessions.get(job_id)
+
+            if not session:
+                logger.warning(f"session with id: {job_id} not found")
                 return False
             try:
-                await module.stop()
-                # should maybe be added in finally
-                del self.queues[job_id]
-                del self.modules[job_id]
-                logger.debug(f"Module {job_id} ({module.name}) stopped successfully")
+                await session.module.stop()
+
+                if job_id in self.tasks:
+                    await self.cancel_task(job_id)
+                logger.debug(f"session {job_id} ({session.module.name}) stopped successfully")
             except Exception as e:
                 logger.error(f"Error while stopping module {job_id}: {e}")
                 raise
             else:
                 return True
 
-    async def get_module_status(self, job_id: str) -> ModuleStatus | None:
+    async def get_module_status(self, job_id: str) -> TaskStatus:
         """Retrieve the status of a module job.
 
         Args:
             job_id: The unique identifier of the job.
 
         Returns:
-            ModuleStatus | None: The status of the module, or None if it does not exist.
+            ModuleStatus: The status of the module.
         """
-        module = self.modules.get(job_id)
-        return module.status if module else None
+        session = self.tasks_sessions.get(job_id, None)
+        return session.status if session is not None else TaskStatus.FAILED
 
     async def stop_all_modules(self) -> None:
         """Stop all currently running module jobs.
@@ -274,7 +280,7 @@ class SingleJobManager(BaseJobManager, Generic[InputModelT, SetupModelT]):
         This method ensures that all active jobs are gracefully terminated.
         """
         async with self._lock:
-            stop_tasks = [self.stop_module(job_id) for job_id in list(self.modules.keys())]
+            stop_tasks = [self.stop_module(job_id) for job_id in list(self.tasks_sessions.keys())]
             if stop_tasks:
                 await asyncio.gather(*stop_tasks, return_exceptions=True)
 
@@ -286,9 +292,9 @@ class SingleJobManager(BaseJobManager, Generic[InputModelT, SetupModelT]):
         """
         return {
             job_id: {
-                "name": module.name,
-                "status": module.status,
-                "class": module.__class__.__name__,
+                "name": session.module.name,
+                "status": session.module.status,
+                "class": session.module.__class__.__name__,
             }
-            for job_id, module in self.modules.items()
+            for job_id, session in self.tasks_sessions.items()
         }

digitalkin/modules/job_manager/surrealdb_repository.py

@@ -0,0 +1,225 @@
+"""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
+        self.url = f"{os.getenv('SURREALDB_URL', 'ws://localhost')}:{os.getenv('SURREALDB_PORT', '8000')}/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)