feat: use service account for tiled insertion (#1398)

Author: ZohebShaikh

helm/blueapi/config_schema.json

@@ -440,6 +440,28 @@
             "type": "object",
             "$id": "ScratchRepository"
         },
+        "ServiceAccount": {
+            "additionalProperties": false,
+            "properties": {
+                "client_id": {
+                    "default": "",
+                    "description": "Service account client ID",
+                    "title": "Client Id",
+                    "type": "string"
+                },
+                "client_secret": {
+                    "default": "",
+                    "description": "Service account client secret",
+                    "format": "password",
+                    "title": "Client Secret",
+                    "type": "string",
+                    "writeOnly": true
+                }
+            },
+            "title": "ServiceAccount",
+            "type": "object",
+            "$id": "ServiceAccount"
+        },
         "StompConfig": {
             "additionalProperties": false,
             "description": "Config for connecting to stomp broker",
@@ -491,17 +513,21 @@
                     "title": "Url",
                     "type": "string"
                 },
-                "api_key": {
+                "authentication": {
                     "anyOf": [
                         {
                             "type": "string"
                         },
+                        {
+                            "$ref": "ServiceAccount"
+                        },
                         {
                             "type": "null"
                         }
                     ],
                     "default": null,
-                    "title": "Api Key"
+                    "description": "Tiled Authentication can be API_KEY or OIDC Service account",
+                    "title": "Authentication"
                 }
             },
             "title": "TiledConfig",

helm/blueapi/values.schema.json

@@ -872,6 +872,28 @@
             },
             "additionalProperties": false
         },
+        "ServiceAccount": {
+            "$id": "ServiceAccount",
+            "title": "ServiceAccount",
+            "type": "object",
+            "properties": {
+                "client_id": {
+                    "title": "Client Id",
+                    "description": "Service account client ID",
+                    "default": "",
+                    "type": "string"
+                },
+                "client_secret": {
+                    "title": "Client Secret",
+                    "description": "Service account client secret",
+                    "writeOnly": true,
+                    "default": "",
+                    "type": "string",
+                    "format": "password"
+                }
+            },
+            "additionalProperties": false
+        },
         "StompConfig": {
             "$id": "StompConfig",
             "title": "StompConfig",
@@ -910,12 +932,16 @@
             "title": "TiledConfig",
             "type": "object",
             "properties": {
-                "api_key": {
-                    "title": "Api Key",
+                "authentication": {
+                    "title": "Authentication",
+                    "description": "Tiled Authentication can be API_KEY or OIDC Service account",
                     "anyOf": [
                         {
                             "type": "string"
                         },
+                        {
+                            "$ref": "ServiceAccount"
+                        },
                         {
                             "type": "null"
                         }

pyproject.toml

@@ -13,7 +13,7 @@ classifiers = [
 ]
 description = "Lightweight bluesky-as-a-service wrapper application. Also usable as a library."
 dependencies = [
-    "tiled[client]>=0.2.3",
+    "tiled[client]>=0.2.4",
     "bluesky[plotting]>=1.14.0",                     # plotting includes matplotlib, required for BestEffortCallback in run plans
     "ophyd-async>=0.13.5",
     "aioca",
@@ -68,7 +68,8 @@ dev = [
     "mock",
     "jwcrypto",
     "deepdiff",
-    "tiled[minimal-server]>=0.2.3", # For system-test of dls.py
+    "tiled[minimal-server]>=0.2.4", # For system-test of dls.py
+    "respx"
 ]
 
 [project.scripts]

src/blueapi/config.py

@@ -16,6 +16,7 @@
     BaseModel,
     Field,
     HttpUrl,
+    SecretStr,
     TypeAdapter,
     UrlConstraints,
     ValidationError,
@@ -107,13 +108,26 @@ class StompConfig(BlueapiBaseModel):
     )
 
 
+class ServiceAccount(BlueapiBaseModel):
+    client_id: str = Field(description="Service account client ID", default="")
+    client_secret: SecretStr = Field(
+        description="Service account client secret", default=SecretStr("")
+    )
+    token_url: SkipJsonSchema[str] = Field(
+        description="Field overridden by OIDCConfig.token_endpoint", default=""
+    )
+
+
 class TiledConfig(BlueapiBaseModel):
     enabled: bool = Field(
         description="True if blueapi should forward data to a Tiled instance",
         default=False,
     )
     url: HttpUrl = HttpUrl("http://localhost:8407")
-    api_key: str | None = os.environ.get("TILED_SINGLE_USER_API_KEY", None)
+    authentication: str | ServiceAccount | None = Field(
+        description="Tiled Authentication can be API_KEY or OIDC Service account",
+        default=os.environ.get("TILED_SINGLE_USER_API_KEY", None),
+    )
 
 
 class WorkerEventConfig(BlueapiBaseModel):

src/blueapi/core/context.py

@@ -30,6 +30,7 @@
     DodalSource,
     EnvironmentConfig,
     PlanSource,
+    ServiceAccount,
     TiledConfig,
 )
 from blueapi.core.protocols import DeviceManager
@@ -186,6 +187,13 @@ def _update_scan_num(md: dict[str, Any]) -> int:
                     "Tiled has been configured but `instrument` metadata is not set - "
                     "this field is required to make authorization decisions."
                 )
+            if isinstance(tiled_conf.authentication, ServiceAccount):
+                if configuration.oidc is None:
+                    raise InvalidConfigError(
+                        "Tiled has been configured but oidc configuration is missing "
+                        "this field is required to make authorization decisions."
+                    )
+                tiled_conf.authentication.token_url = configuration.oidc.token_endpoint
             self.tiled_conf = tiled_conf
 
     def find_device(self, addr: str | list[str]) -> Device | None:

src/blueapi/service/authentication.py

@@ -2,6 +2,7 @@
 
 import base64
 import os
+import threading
 import time
 import webbrowser
 from abc import ABC, abstractmethod
@@ -10,12 +11,13 @@
 from pathlib import Path
 from typing import Any, cast
 
+import httpx
 import jwt
 import requests
 from pydantic import TypeAdapter
 from requests.auth import AuthBase
 
-from blueapi.config import OIDCConfig
+from blueapi.config import OIDCConfig, ServiceAccount
 from blueapi.service.model import Cache
 
 DEFAULT_CACHE_DIR = "~/.cache/"
@@ -239,3 +241,28 @@ def __call__(self, request):
         if self.token:
             request.headers["Authorization"] = f"Bearer {self.token}"
         return request
+
+
+class TiledAuth(httpx.Auth):
+    def __init__(self, tiled_auth: ServiceAccount):
+        if tiled_auth.token_url == "":
+            raise RuntimeError("Token URL is not set please check oidc config")
+        self._tiled_auth: ServiceAccount = tiled_auth
+        self._sync_lock = threading.RLock()
+
+    def get_access_token(self):
+        with self._sync_lock:
+            response = httpx.post(
+                self._tiled_auth.token_url,
+                data={
+                    "client_id": self._tiled_auth.client_id,
+                    "client_secret": self._tiled_auth.client_secret.get_secret_value(),
+                    "grant_type": "client_credentials",
+                },
+            )
+            response.raise_for_status()
+            return response.json().get("access_token")
+
+    def sync_auth_flow(self, request):
+        request.headers["Authorization"] = f"Bearer {self.get_access_token()}"
+        yield request

src/blueapi/service/interface.py

@@ -8,10 +8,11 @@
 from tiled.client import from_uri
 
 from blueapi.cli.scratch import get_python_environment
-from blueapi.config import ApplicationConfig, OIDCConfig, StompConfig
+from blueapi.config import ApplicationConfig, OIDCConfig, ServiceAccount, StompConfig
 from blueapi.core.context import BlueskyContext
 from blueapi.core.event import EventStream
 from blueapi.log import set_up_logging
+from blueapi.service.authentication import TiledAuth
 from blueapi.service.model import (
     DeviceModel,
     PlanModel,
@@ -188,11 +189,18 @@ def begin_task(
 
     if tiled_config := active_context.tiled_conf:
         # Tiled queries the root node, so must create an authorized client
-        tiled_client = from_uri(
-            str(tiled_config.url),
-            api_key=tiled_config.api_key,
-            headers=pass_through_headers,
-        )
+        if isinstance(tiled_config.authentication, ServiceAccount):
+            tiled_client = from_uri(
+                str(tiled_config.url),
+                auth=TiledAuth(tiled_auth=tiled_config.authentication),
+            )
+        else:
+            tiled_client = from_uri(
+                str(tiled_config.url),
+                api_key=tiled_config.authentication,
+                headers=pass_through_headers,
+            )
+
         tiled_writer_token = active_context.run_engine.subscribe(
             TiledWriter(tiled_client, batch_size=1)
         )

tests/system_tests/compose.yaml

@@ -39,16 +39,16 @@ services:
       retries: 10
       start_period: 30s
 
-  tiled: 
-    image: ghcr.io/bluesky/tiled:0.2.3
+  tiled:
+    image: ghcr.io/bluesky/tiled:0.2.4
     network_mode: host
     environment:
       - PYTHONPATH=/deploy/
-    volumes: 
+    volumes:
       - ./services/tiled_config:/deploy/config
-    command: ["tiled", "serve", "config", "--host", "0.0.0.0", "--port", "8407"] 
-    depends_on: 
-      keycloak: 
+    command: ["tiled", "serve", "config", "--host", "0.0.0.0", "--port", "8407"]
+    depends_on:
+      keycloak:
         condition: service_healthy
 
   opa:

tests/system_tests/config.yaml

@@ -18,6 +18,9 @@ numtracker:
 tiled:
   enabled: true
   url: http://localhost:8407/api/v1
+  authentication:
+    client_id: "tiled-writer"
+    client_secret: "secret"
 oidc:
   well_known_url: "http://localhost:8081/realms/master/.well-known/openid-configuration"
   client_id: "ixx-cli-blueapi"

tests/system_tests/services/keycloak_config/service-account.json

@@ -0,0 +1,34 @@
+{
+    "protocolMappers": [
+        {
+            "name": "beamline",
+            "protocol": "openid-connect",
+            "protocolMapper": "oidc-hardcoded-claim-mapper",
+            "consentRequired": false,
+            "config": {
+                "introspection.token.claim": "true",
+                "claim.value": "adsim",
+                "userinfo.token.claim": "true",
+                "id.token.claim": "true",
+                "lightweight.claim": "false",
+                "access.token.claim": "true",
+                "claim.name": "beamline",
+                "jsonType.label": "String",
+                "access.tokenResponse.claim": "false"
+            }
+        },
+        {
+            "name": "tiled",
+            "protocol": "openid-connect",
+            "protocolMapper": "oidc-audience-mapper",
+            "consentRequired": false,
+            "config": {
+                "id.token.claim": "false",
+                "lightweight.claim": "false",
+                "access.token.claim": "true",
+                "introspection.token.claim": "true",
+                "included.custom.audience": "tiled-writer"
+            }
+        }
+    ]
+}