Refacors state handling to use that-depends context over env vars

Author: chrisk314

plugboard/cli/process/__init__.py

@@ -1,7 +1,6 @@
 """Plugboard Process CLI."""
 
 import asyncio
-import os
 from pathlib import Path
 
 import msgspec
@@ -71,8 +70,7 @@ def run(
     config_spec = _read_yaml(config)
 
     if job_id:
-        # Override job ID in env and config file if set
-        os.environ["PLUGBOARD_JOB_ID"] = job_id
+        # Override job ID in config file if set
         config_spec.plugboard.process.args.state.args.job_id = job_id
 
     with Progress(

plugboard/process/process_builder.py

@@ -1,6 +1,5 @@
 """Provides `ProcessBuilder` to build `Process` objects."""
 
-import os
 from pydoc import locate
 import typing as _t
 
@@ -9,7 +8,7 @@
 from plugboard.connector.connector_builder import ConnectorBuilder
 from plugboard.events.event_connector_builder import EventConnectorBuilder
 from plugboard.process.process import Process
-from plugboard.schemas import ProcessSpec, StateBackendSpec
+from plugboard.schemas import ProcessSpec
 from plugboard.state import StateBackend
 from plugboard.utils import DI
 
@@ -48,28 +47,8 @@ def _build_statebackend(cls, spec: ProcessSpec) -> StateBackend:
         statebackend_class: _t.Optional[_t.Any] = locate(state_spec.type)
         if not statebackend_class or not issubclass(statebackend_class, StateBackend):
             raise ValueError(f"StateBackend class {spec.args.state.type} not found.")
-        cls._handle_job_id(state_spec)
         return statebackend_class(**dict(spec.args.state.args))
 
-    @staticmethod
-    def _handle_job_id(state_spec: StateBackendSpec) -> None:
-        """Handle job ID for the state backend.
-
-        If a job ID is provided in the state spec, it will be set as an environment variable.
-        If the job ID is already set in the environment, it will be checked against the one in the
-        state spec. If they do not match, a RuntimeError will be raised.
-        """
-        if state_spec.args.job_id is None:
-            return
-        if (
-            env_job_id := os.environ.get("PLUGBOARD_JOB_ID")
-        ) is not None and env_job_id != state_spec.args.job_id:
-            raise RuntimeError(
-                f"Job ID {state_spec.args.job_id} does not match environment variable "
-                f"PLUGBOARD_JOB_ID={env_job_id}"
-            )
-        os.environ["PLUGBOARD_JOB_ID"] = state_spec.args.job_id
-
     @classmethod
     def _build_components(cls, spec: ProcessSpec) -> list[Component]:
         for c in spec.args.components:

plugboard/state/state_backend.py

@@ -3,11 +3,14 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
+from contextlib import AsyncExitStack
 from datetime import datetime, timezone
-import os
 from types import TracebackType
 import typing as _t
 
+from that_depends import Provide, container_context, inject
+
+from plugboard.exceptions import NotFoundError
 from plugboard.utils import DI, ExportMixin
 
 
@@ -33,13 +36,18 @@ def __init__(
         self._local_state = {"job_id": job_id, "metadata": metadata, **kwargs}
         self._logger = DI.logger.sync_resolve().bind(cls=self.__class__.__name__, job_id=job_id)
         self._logger.info("StateBackend created")
+        self._ctx = AsyncExitStack()
 
     async def init(self) -> None:
         """Initialises the `StateBackend`."""
+        job_id = self._local_state.pop("job_id", None)
+        container_cm = container_context(global_context={"job_id": job_id})
+        await self._ctx.enter_async_context(container_cm)
         await self._initialise_data(**self._local_state)
 
     async def destroy(self) -> None:
         """Destroys the `StateBackend`."""
+        await self._ctx.aclose()
         pass
 
     async def __aenter__(self) -> StateBackend:
@@ -56,34 +64,23 @@ async def __aexit__(
         """Exits the context manager."""
         await self.destroy()
 
+    @inject
     async def _initialise_data(
-        self, job_id: _t.Optional[str] = None, metadata: _t.Optional[dict] = None, **kwargs: _t.Any
+        self, job_id: str = Provide[DI.job_id], metadata: _t.Optional[dict] = None, **kwargs: _t.Any
     ) -> None:
         """Initialises the state data."""
-        if (_job_id := self._resolve_job_id(job_id)) is not None:
-            job_data = await self._get_job(_job_id)
-        else:
+        try:
+            # TODO : Requires state for if this is a new job to conditionally raise exception?
+            job_data = await self._get_job(job_id)
+        except NotFoundError:
             job_data = {
-                "job_id": DI.job_id.sync_resolve(),
+                "job_id": job_id,
                 "created_at": datetime.now(timezone.utc).isoformat(),
                 "metadata": metadata or dict(),
             }
             await self._upsert_job(job_data)
         self._local_state.update(job_data)
 
-    @staticmethod
-    def _resolve_job_id(job_id: _t.Optional[str] = None) -> _t.Optional[str]:
-        """Resolves the job id from the environment or argument if present."""
-        env_job_id = os.environ.get("PLUGBOARD_JOB_ID")
-        if job_id is None:
-            return env_job_id
-        if env_job_id is not None and job_id != env_job_id:
-            raise RuntimeError(
-                f"Job ID {job_id} does not match environment variable PLUGBOARD_JOB_ID={env_job_id}"
-            )
-        os.environ["PLUGBOARD_JOB_ID"] = job_id
-        return job_id
-
     @abstractmethod
     async def _get(self, key: str | tuple[str, ...], value: _t.Optional[_t.Any] = None) -> _t.Any:
         """Returns a value from the state."""

plugboard/utils/di.py

@@ -6,8 +6,8 @@
 
 import aio_pika
 import structlog
-from that_depends import BaseContainer
-from that_depends.providers import Resource, Singleton
+from that_depends import BaseContainer, fetch_context_item
+from that_depends.providers import ContextResource, Resource, Singleton
 from yarl import URL
 
 from plugboard._zmq.zmq_proxy import ZMQProxy
@@ -59,17 +59,29 @@ async def _rabbitmq_conn(
         await conn.close()  # pragma: no cover
 
 
-def _job_id() -> str:
+def _job_id() -> _t.Iterator[str]:
     """Returns a job ID which uniquely identifies the current plugboard run.
 
-    If a job ID is set in the environment variable `PLUGBOARD_JOB_ID`, it will be used.
-    Otherwise, a new unique job ID will be generated and set in the environment.
+    If a job ID is available in the context (from the cli, the state spec, or an argument to the
+    StateBackend), it will take precedence. If the job ID is set in the env var `PLUGBOARD_JOB_ID`,
+    it will be checked against the one in the context, if present. If they do not match, a
+    RuntimeError will be raised. If the job ID is not set in the context or the env var, a new
+    unique job ID will be generated.
     """
-    # TODO : Should the env var be unset on DI teardown? Consider notebook execution.
-    #      : Where multiple process runs may be executed from the same os process.
-    if (job_id := os.environ.get("PLUGBOARD_JOB_ID")) is None:
-        os.environ["PLUGBOARD_JOB_ID"] = job_id = EntityIdGen.job_id()
-    return job_id
+    arg_job_id = fetch_context_item("job_id")
+    env_job_id = os.environ.get("PLUGBOARD_JOB_ID")
+    if arg_job_id is not None:
+        if env_job_id is not None and arg_job_id != env_job_id:
+            raise RuntimeError(
+                f"Job ID {arg_job_id} does not match environment variable "
+                f"PLUGBOARD_JOB_ID={env_job_id}"
+            )
+        job_id = arg_job_id
+    elif env_job_id is not None:
+        job_id = env_job_id
+    else:
+        job_id = EntityIdGen.job_id()
+    yield job_id
 
 
 class DI(BaseContainer):
@@ -84,4 +96,4 @@ class DI(BaseContainer):
     rabbitmq_conn: Resource[aio_pika.abc.AbstractRobustConnection] = Resource(
         _rabbitmq_conn, logger, url=settings.rabbitmq.url
     )
-    job_id: Singleton[str] = Singleton(_job_id)
+    job_id: ContextResource[str] = ContextResource(_job_id)