DM-38425: Split out Gafaelfawr storage layer and use models

src/mobu/constants.py

@@ -1,7 +1,34 @@
 """Global constants for mobu."""
 
+from __future__ import annotations
+
+from datetime import timedelta
+
+__all__ = [
+    "NOTEBOOK_REPO_URL",
+    "NOTEBOOK_REPO_BRANCH",
+    "TOKEN_LIFETIME",
+    "USERNAME_REGEX",
+]
+
 NOTEBOOK_REPO_URL = "https://github.com/lsst-sqre/notebook-demo.git"
 """Default notebook repository for NotebookRunner."""
 
 NOTEBOOK_REPO_BRANCH = "prod"
 """Default repository branch for NotebookRunner."""
+
+TOKEN_LIFETIME = timedelta(days=365)
+"""Token lifetime for mobu's service tokens.
+
+mobu currently has no mechanism for refreshing tokens while running, so this
+should be long enough that mobu will be restarted before the tokens expire.
+An expiration exists primarily to ensure that the tokens don't accumulate
+forever.
+"""
+
+# This must be kept in sync with Gafaelfawr until we can import the models
+# from Gafaelfawr directly.
+USERNAME_REGEX = (
+    "^[a-z0-9](?:[a-z0-9]|-[a-z0-9])*[a-z](?:[a-z0-9]|-[a-z0-9])*$"
+)
+"""Regex matching all valid usernames."""

src/mobu/exceptions.py

@@ -8,6 +8,7 @@
 
 from fastapi import status
 from httpx_ws import HTTPXWSException, WebSocketDisconnect
+from pydantic import ValidationError
 from safir.datetime import format_datetime_for_logging
 from safir.fastapi import ClientRequestError
 from safir.models import ErrorLocation
@@ -29,6 +30,7 @@
     "CachemachineError",
     "CodeExecutionError",
     "FlockNotFoundError",
+    "GafaelfawrParseError",
     "GafaelfawrWebError",
     "JupyterTimeoutError",
     "JupyterWebError",
@@ -62,6 +64,60 @@ def _remove_ansi_escapes(string: str) -> str:
     return _ANSI_REGEX.sub("", string)
 
 
+class GafaelfawrParseError(SlackException):
+    """Unable to parse the reply from Gafaelfawr.
+
+    Parameters
+    ----------
+    message
+        Summary error message.
+    error
+        Detailed error message, possibly multi-line.
+    user
+        Username of the user involved.
+    """
+
+    @classmethod
+    def from_exception(
+        cls, exc: ValidationError, user: Optional[str] = None
+    ) -> Self:
+        """Create an exception from a Pydantic parse failure.
+
+        Parameters
+        ----------
+        exc
+            Pydantic exception.
+        user
+            Username of the user involved.
+
+        Returns
+        -------
+        GafaelfawrParseError
+            Constructed exception.
+        """
+        error = f"{type(exc).__name__}: {str(exc)}"
+        return cls("Unable to parse reply from Gafalefawr", error, user)
+
+    def __init__(
+        self, message: str, error: str, user: Optional[str] = None
+    ) -> None:
+        super().__init__(message, user)
+        self.error = error
+
+    def to_slack(self) -> SlackMessage:
+        """Convert to a Slack message for Slack alerting.
+
+        Returns
+        -------
+        SlackMessage
+            Slack message suitable for posting as an alert.
+        """
+        message = super().to_slack()
+        block = SlackCodeBlock(heading="Error", code=self.error)
+        message.blocks.append(block)
+        return message
+
+
 class GafaelfawrWebError(SlackWebException):
     """An API call to Gafaelfawr failed."""
 

src/mobu/factory.py

@@ -13,6 +13,7 @@
 from .models.solitary import SolitaryConfig
 from .services.manager import FlockManager
 from .services.solitary import Solitary
+from .storage.gafaelfawr import GafaelfawrStorage
 
 __all__ = ["Factory", "ProcessContext"]
 
@@ -38,7 +39,9 @@ class ProcessContext:
 
     def __init__(self, http_client: AsyncClient) -> None:
         self.http_client = http_client
-        self.manager = FlockManager(http_client, structlog.get_logger("mobu"))
+        logger = structlog.get_logger("mobu")
+        gafaelfawr = GafaelfawrStorage(http_client, logger)
+        self.manager = FlockManager(gafaelfawr, http_client, logger)
 
     async def aclose(self) -> None:
         """Clean up a process context.
@@ -93,5 +96,10 @@ def create_solitary(self, solitary_config: SolitaryConfig) -> Solitary:
             Newly-created solitary manager.
         """
         return Solitary(
-            solitary_config, self._context.http_client, self._logger
+            solitary_config=solitary_config,
+            gafaelfawr_storage=GafaelfawrStorage(
+                self._context.http_client, self._logger
+            ),
+            http_client=self._context.http_client,
+            logger=self._logger,
         )

src/mobu/models/user.py

@@ -2,16 +2,15 @@
 
 from __future__ import annotations
 
-import time
-from typing import Any, Optional, Self
+from typing import Optional
 
-from httpx import AsyncClient, HTTPError
 from pydantic import BaseModel, Field
 
-from ..config import config
-from ..exceptions import GafaelfawrWebError
-
-__all__ = ["AuthenticatedUser", "User", "UserSpec"]
+__all__ = [
+    "AuthenticatedUser",
+    "User",
+    "UserSpec",
+]
 
 
 class User(BaseModel):
@@ -91,46 +90,3 @@ class AuthenticatedUser(User):
         title="Authentication token for user",
         example="gt-1PhgAeB-9Fsa-N1NhuTu_w.oRvMvAQp1bWfx8KCJKNohg",
     )
-
-    @classmethod
-    async def create(
-        cls, user: User, scopes: list[str], http_client: AsyncClient
-    ) -> Self:
-        if not config.environment_url:
-            raise RuntimeError("environment_url not set")
-        token_url = (
-            str(config.environment_url).rstrip("/") + "/auth/api/v1/tokens"
-        )
-        data: dict[str, Any] = {
-            "username": user.username,
-            "name": "Mobu Test User",
-            "token_type": "user",
-            "token_name": f"mobu {str(float(time.time()))}",
-            "scopes": scopes,
-            "expires": int(time.time() + 60 * 60 * 24 * 365),
-        }
-        if user.uidnumber is not None:
-            data["uid"] = user.uidnumber
-            if user.gidnumber is not None:
-                data["gid"] = user.gidnumber
-            else:
-                data["gid"] = user.uidnumber
-        elif user.gidnumber is not None:
-            data["gid"] = user.gidnumber
-        try:
-            r = await http_client.post(
-                token_url,
-                headers={"Authorization": f"Bearer {config.gafaelfawr_token}"},
-                json=data,
-            )
-            r.raise_for_status()
-        except HTTPError as e:
-            raise GafaelfawrWebError.from_exception(e, user.username)
-        body = r.json()
-        return cls(
-            username=user.username,
-            uidnumber=data["uid"] if "uid" in data else None,
-            gidnumber=data["gid"] if "gid" in data else None,
-            token=body["token"],
-            scopes=scopes,
-        )

src/mobu/services/flock.py

@@ -15,6 +15,7 @@
 from ..exceptions import MonkeyNotFoundError
 from ..models.flock import FlockConfig, FlockData, FlockSummary
 from ..models.user import AuthenticatedUser, User, UserSpec
+from ..storage.gafaelfawr import GafaelfawrStorage
 from .monkey import Monkey
 
 __all__ = ["Flock"]
@@ -29,6 +30,8 @@ class Flock:
         Configuration for this flock of monkeys.
     scheduler
         Job scheduler used to manage the tasks for the monkeys.
+    gafaelfawr_storage
+        Gafaelfawr storage client.
     http_client
         Shared HTTP client.
     logger
@@ -40,12 +43,14 @@ def __init__(
         *,
         flock_config: FlockConfig,
         scheduler: Scheduler,
+        gafaelfawr_storage: GafaelfawrStorage,
         http_client: AsyncClient,
         logger: BoundLogger,
     ) -> None:
         self.name = flock_config.name
         self._config = flock_config
         self._scheduler = scheduler
+        self._gafaelfawr = gafaelfawr_storage
         self._http_client = http_client
         self._logger = logger.bind(flock=self.name)
         self._monkeys: dict[str, Monkey] = {}
@@ -146,7 +151,7 @@ async def _create_users(self) -> list[AuthenticatedUser]:
             users = self._users_from_spec(self._config.user_spec, count)
         scopes = self._config.scopes
         return [
-            await AuthenticatedUser.create(u, scopes, self._http_client)
+            await self._gafaelfawr.create_service_token(u, scopes)
             for u in users
         ]
 

src/mobu/services/manager.py

@@ -12,6 +12,7 @@
 from ..config import config
 from ..exceptions import FlockNotFoundError
 from ..models.flock import FlockConfig, FlockSummary
+from ..storage.gafaelfawr import GafaelfawrStorage
 from .flock import Flock
 
 __all__ = ["FlockManager"]
@@ -26,13 +27,21 @@ class FlockManager:
 
     Parameters
     ----------
+    gafaelfawr_storage
+        Gafaelfawr storage client.
     http_client
         Shared HTTP client.
     logger
         Global logger to use for process-wide (not monkey) logging.
     """
 
-    def __init__(self, http_client: AsyncClient, logger: BoundLogger) -> None:
+    def __init__(
+        self,
+        gafaelfawr_storage: GafaelfawrStorage,
+        http_client: AsyncClient,
+        logger: BoundLogger,
+    ) -> None:
+        self._gafaelfawr = gafaelfawr_storage
         self._http_client = http_client
         self._logger = logger
         self._flocks: dict[str, Flock] = {}
@@ -74,6 +83,7 @@ async def start_flock(self, flock_config: FlockConfig) -> Flock:
         flock = Flock(
             flock_config=flock_config,
             scheduler=self._scheduler,
+            gafaelfawr_storage=self._gafaelfawr,
             http_client=self._http_client,
             logger=self._logger,
         )

src/mobu/services/solitary.py

@@ -8,7 +8,7 @@
 from structlog.stdlib import BoundLogger
 
 from ..models.solitary import SolitaryConfig, SolitaryResult
-from ..models.user import AuthenticatedUser
+from ..storage.gafaelfawr import GafaelfawrStorage
 from .monkey import Monkey
 
 __all__ = ["Solitary"]
@@ -21,6 +21,8 @@ class Solitary:
     ----------
     solitary_config
         Configuration for the monkey.
+    gafaelfawr_storage
+        Gafaelfawr storage client.
     http_client
         Shared HTTP client.
     logger
@@ -29,11 +31,14 @@ class Solitary:
 
     def __init__(
         self,
+        *,
         solitary_config: SolitaryConfig,
+        gafaelfawr_storage: GafaelfawrStorage,
         http_client: AsyncClient,
         logger: BoundLogger,
     ) -> None:
         self._config = solitary_config
+        self._gafaelfawr = gafaelfawr_storage
         self._http_client = http_client
         self._logger = logger
 
@@ -45,8 +50,8 @@ async def run(self) -> SolitaryResult:
         SolitaryResult
             Result of monkey run.
         """
-        user = await AuthenticatedUser.create(
-            self._config.user, self._config.scopes, self._http_client
+        user = await self._gafaelfawr.create_service_token(
+            self._config.user, self._config.scopes
         )
         monkey = Monkey(
             name=f"solitary-{user.username}",