feat: Add OpaClient to wrap auth checks

helm/blueapi/config_schema.json

@@ -330,6 +330,22 @@
             "type": "object",
             "$id": "OIDCConfig"
         },
+        "OpaConfig": {
+            "additionalProperties": false,
+            "properties": {
+                "root": {
+                    "default": "http://localhost:8181/",
+                    "format": "uri",
+                    "maxLength": 2083,
+                    "minLength": 1,
+                    "title": "Root",
+                    "type": "string"
+                }
+            },
+            "title": "OpaConfig",
+            "type": "object",
+            "$id": "OpaConfig"
+        },
         "PlanSource": {
             "additionalProperties": false,
             "properties": {
@@ -612,6 +628,17 @@
                 }
             ],
             "default": null
+        },
+        "opa": {
+            "anyOf": [
+                {
+                    "$ref": "OpaConfig"
+                },
+                {
+                    "type": "null"
+                }
+            ],
+            "default": null
         }
     },
     "title": "ApplicationConfig",

helm/blueapi/values.schema.json

@@ -751,6 +751,22 @@
             },
             "additionalProperties": false
         },
+        "OpaConfig": {
+            "$id": "OpaConfig",
+            "title": "OpaConfig",
+            "type": "object",
+            "properties": {
+                "root": {
+                    "title": "Root",
+                    "default": "http://localhost:8181/",
+                    "type": "string",
+                    "format": "uri",
+                    "maxLength": 2083,
+                    "minLength": 1
+                }
+            },
+            "additionalProperties": false
+        },
         "PlanSource": {
             "$id": "PlanSource",
             "title": "PlanSource",
@@ -1011,6 +1027,16 @@
                         }
                     ]
                 },
+                "opa": {
+                    "anyOf": [
+                        {
+                            "$ref": "OpaConfig"
+                        },
+                        {
+                            "type": "null"
+                        }
+                    ]
+                },
                 "scratch": {
                     "anyOf": [
                         {

pyproject.toml

@@ -37,6 +37,7 @@ dependencies = [
     "tomlkit",
     "graypy>=2.1.0",
     "httpx>=0.28.1",
+    "aiohttp>=3.13.5",
 ]
 dynamic = ["version"]
 license.file = "LICENSE"

src/blueapi/config.py

@@ -296,6 +296,10 @@ class Tag(StrEnum):
     META = "Meta"
 
 
+class OpaConfig(BlueapiBaseModel):
+    root: HttpUrl = HttpUrl("http://localhost:8181")
+
+
 class ApplicationConfig(BlueapiBaseModel):
     """
     Config for the worker application as a whole. Root of
@@ -335,6 +339,7 @@ class ApplicationConfig(BlueapiBaseModel):
     oidc: OIDCConfig | None = None
     auth_token_path: Path | None = None
     numtracker: NumtrackerConfig | None = None
+    opa: OpaConfig | None = None
 
     def __eq__(self, other: object) -> bool:
         if isinstance(other, ApplicationConfig):
@@ -343,6 +348,7 @@ def __eq__(self, other: object) -> bool:
                 & (self.env == other.env)
                 & (self.logging == other.logging)
                 & (self.api == other.api)
+                & (self.opa == other.opa)
             )
         return False
 

src/blueapi/service/authorization.py

@@ -0,0 +1,63 @@
+import logging
+from collections.abc import Mapping
+from contextlib import AbstractAsyncContextManager, aclosing, nullcontext
+from typing import Any, Self
+
+import aiohttp
+from aiohttp import ClientSession
+
+from blueapi.config import OpaConfig
+
+LOGGER = logging.getLogger(__name__)
+
+
+class OpaClient:
+    client: aiohttp.ClientSession
+
+    def __init__(self, instrument: str, config: OpaConfig):
+        LOGGER.info("Creating OpaClient for %s with config %s", instrument, config)
+        self._instrument = instrument
+        self._conf = config
+        self._session = ClientSession(base_url=config.root.encoded_string())
+
+    async def aclose(self):
+        LOGGER.info("Closing OPA session")
+        await self._session.close()
+
+    async def _call_opa(self, endpoint, data: Mapping[str, Any]) -> bool:
+        try:
+            resp = await self._session.post(
+                endpoint,
+                json={
+                    "input": {
+                        "beamline": self._instrument,
+                        "audience": "account",
+                        **data,
+                    }
+                },
+            )
+            return (await resp.json())["result"]
+        except Exception:
+            LOGGER.exception("Failed to run check", exc_info=True)
+            raise
+
+    @classmethod
+    def for_config(
+        cls, instrument: str | None, config: OpaConfig | None
+    ) -> AbstractAsyncContextManager[Self | None]:
+        if config:
+            if not instrument:
+                raise ValueError("Instrument name is required for OPA client")
+            return aclosing(cls(instrument, config))
+        LOGGER.info("No OPA config provided - not creating OpaClient")
+        return nullcontext()
+
+
+
+class OpaUserClient:
+    client: OpaClient
+    token: str
+
+    def __init__(self, client: OpaClient, token: str):
+        self.client = client
+        self.token = token

src/blueapi/service/main.py

@@ -40,6 +40,7 @@
 from blueapi.worker import TrackableTask, WorkerState
 from blueapi.worker.event import TaskStatusEnum
 
+from .authorization import OpaClient
 from .model import (
     DeviceModel,
     DeviceResponse,
@@ -93,8 +94,11 @@ def teardown_runner():
 def lifespan(config: ApplicationConfig):
     @asynccontextmanager
     async def inner(app: FastAPI):
+        meta = config.env.metadata
         setup_runner(config)
-        yield
+        async with OpaClient.for_config(meta and meta.instrument, config.opa) as opa:
+            app.state.authz = opa
+            yield
         teardown_runner()
 
     return inner

tests/unit_tests/service/test_authorization.py

@@ -0,0 +1,48 @@
+from unittest.mock import MagicMock, patch
+
+import pytest
+from pydantic import HttpUrl
+
+from blueapi.config import OpaConfig
+from blueapi.service.authorization import (
+    OpaClient,
+)
+
+# Reusable client patch decorator
+patch_client_session = patch(
+    "blueapi.service.authorization.ClientSession",
+    name="mock_client_session",
+    spec=True,
+)
+
+
+@pytest.fixture(scope="module")
+def opa_config() -> OpaConfig:
+    return OpaConfig(
+        root=HttpUrl("http://auth.example.com"),
+    )
+
+
+@patch_client_session
+async def test_session_closed(session: MagicMock, opa_config: OpaConfig):
+    async with OpaClient.for_config("p45", opa_config):
+        pass
+    session().close.assert_called_once()
+
+
+@patch_client_session
+async def test_opa_client_for_config(session: MagicMock, opa_config: OpaConfig):
+    async with OpaClient.for_config("p45", opa_config) as opa:
+        assert opa is not None
+        session.assert_called_once_with(base_url="http://auth.example.com/")
+
+
+@pytest.mark.parametrize("instrument", [None, "p99"])
+async def test_opa_client_without_config(instrument: str | None):
+    async with OpaClient.for_config(instrument, None) as opa:
+        assert opa is None
+
+
+async def test_opa_fails_without_instrument(opa_config: OpaConfig):
+    with pytest.raises(ValueError, match="Instrument name is required"):
+        OpaClient.for_config(None, opa_config)

tests/unit_tests/test_config.py

@@ -337,6 +337,9 @@ def test_config_yaml_parsed(temp_yaml_config_file):
                     }
                 ],
             },
+            "opa": {
+                "root": "http://opa.example.com/",
+            },
         },
         {
             "stomp": {
@@ -392,6 +395,7 @@ def test_config_yaml_parsed(temp_yaml_config_file):
                     }
                 ],
             },
+            "opa": None,
         },
     ],
     indirect=True,