diff --git a/helm/blueapi/config_schema.json b/helm/blueapi/config_schema.json index b5d0c9bf3..1ad4e82bf 100644 --- a/helm/blueapi/config_schema.json +++ b/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", diff --git a/helm/blueapi/values.schema.json b/helm/blueapi/values.schema.json index 74deedadb..6de532cc8 100644 --- a/helm/blueapi/values.schema.json +++ b/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": [ { diff --git a/pyproject.toml b/pyproject.toml index 659779994..9f4231c76 100644 --- a/pyproject.toml +++ b/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" diff --git a/src/blueapi/config.py b/src/blueapi/config.py index 83d6d7021..4c2431cf1 100644 --- a/src/blueapi/config.py +++ b/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 diff --git a/src/blueapi/service/authorization.py b/src/blueapi/service/authorization.py new file mode 100644 index 000000000..b27c304c9 --- /dev/null +++ b/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 diff --git a/src/blueapi/service/main.py b/src/blueapi/service/main.py index 546155f5c..462c9318f 100644 --- a/src/blueapi/service/main.py +++ b/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 diff --git a/tests/unit_tests/service/test_authorization.py b/tests/unit_tests/service/test_authorization.py new file mode 100644 index 000000000..aa4417d50 --- /dev/null +++ b/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) diff --git a/tests/unit_tests/test_config.py b/tests/unit_tests/test_config.py index 30fe551c4..7ce98f6fe 100644 --- a/tests/unit_tests/test_config.py +++ b/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, diff --git a/uv.lock b/uv.lock index b2d10b5bd..eab94a867 100644 --- a/uv.lock +++ b/uv.lock @@ -420,6 +420,7 @@ name = "blueapi" source = { editable = "." } dependencies = [ { name = "aioca" }, + { name = "aiohttp" }, { name = "bluesky", extra = ["plotting"] }, { name = "bluesky-stomp" }, { name = "click" }, @@ -481,6 +482,7 @@ dev = [ [package.metadata] requires-dist = [ { name = "aioca" }, + { name = "aiohttp", specifier = ">=3.13.5" }, { name = "bluesky", extras = ["plotting"], specifier = ">=1.14.0" }, { name = "bluesky-stomp", specifier = ">=0.1.6" }, { name = "click", specifier = ">=8.2.0" },