Add support for quota overrides

Add a new REST API to set, retrieve, and delete quota overrides. Apply
those overrides on top of the configured quotas. The overrides are
stored in Redis and retrieved on every request currently; we will see
if that's too slow and needs caching.

changelog.d/20250116_152452_rra_DM_48432.md

@@ -0,0 +1,3 @@
+### New features
+
+- Add support for quota overrides. Overrides can be set via a new REST API at `/auth/api/v1/quota-overrides` and take precedence over the configured quotas if present and applicable.

docs/dev/internals.rst

@@ -159,6 +159,9 @@ Python internal API
 .. automodapi:: gafaelfawr.storage.oidc
    :include-all-objects:
 
+.. automodapi:: gafaelfawr.storage.quota
+   :include-all-objects:
+
 .. automodapi:: gafaelfawr.storage.token
    :include-all-objects:
 

src/gafaelfawr/config.py

@@ -53,8 +53,8 @@
 from .constants import MINIMUM_LIFETIME, SCOPE_REGEX, USERNAME_REGEX
 from .exceptions import InvalidTokenError
 from .keypair import RSAKeyPair
+from .models.quota import QuotaConfig
 from .models.token import Token
-from .models.userinfo import Quota
 from .util import group_name_for_github_team
 
 HttpsUrl = Annotated[
@@ -656,69 +656,6 @@ def keypair(self) -> RSAKeyPair:
         return self._keypair
 
 
-class QuotaConfig(BaseModel):
-    """Quota configuration."""
-
-    model_config = ConfigDict(extra="forbid")
-
-    default: Quota = Field(
-        ..., title="Default quota", description="Default quotas for all users"
-    )
-
-    groups: dict[str, Quota] = Field(
-        {},
-        title="Quota grants by group",
-        description="Additional quota grants by group name",
-    )
-
-    bypass: set[str] = Field(
-        set(),
-        title="Groups without quotas",
-        description="Groups whose members bypass all quota restrictions",
-    )
-
-    def calculate_quota(self, groups: set[str]) -> Quota | None:
-        """Calculate user's quota given their group membership.
-
-        Parameters
-        ----------
-        groups
-            Group membership of the user.
-
-        Returns
-        -------
-        Quota or None
-            Quota information for that user or `None` if no quotas apply.
-        """
-        if groups & self.bypass:
-            return None
-
-        # Start with the defaults.
-        api = dict(self.default.api)
-        notebook = None
-        if self.default.notebook:
-            notebook = self.default.notebook.model_copy()
-
-        # Look for group-specific rules.
-        for group in groups & set(self.groups.keys()):
-            extra = self.groups[group]
-            if extra.notebook:
-                if notebook:
-                    notebook.cpu += extra.notebook.cpu
-                    notebook.memory += extra.notebook.memory
-                    notebook.spawn &= extra.notebook.spawn
-                else:
-                    notebook = extra.notebook.model_copy()
-            for service, quota in extra.api.items():
-                if service in api:
-                    api[service] += quota
-                else:
-                    api[service] = quota
-
-        # Return the results.
-        return Quota(api=api, notebook=notebook)
-
-
 class GitHubGroupTeam(BaseModel):
     """Specification for a GitHub team."""
 

src/gafaelfawr/factory.py

@@ -20,7 +20,7 @@
 from redis.backoff import ExponentialBackoff
 from safir.database import create_async_session
 from safir.dependencies.http_client import http_client_dependency
-from safir.redis import EncryptedPydanticRedisStorage
+from safir.redis import EncryptedPydanticRedisStorage, PydanticRedisStorage
 from safir.slack.webhook import SlackWebhookClient
 from sqlalchemy import select
 from sqlalchemy.ext.asyncio import AsyncEngine, async_scoped_session
@@ -39,6 +39,7 @@
 from .exceptions import NotConfiguredError
 from .models.ldap import LDAPUserData
 from .models.oidc import OIDCAuthorization
+from .models.quota import QuotaConfig
 from .models.token import TokenData
 from .models.userinfo import Group
 from .providers.base import Provider
@@ -66,6 +67,7 @@
 )
 from .storage.ldap import LDAPStorage
 from .storage.oidc import OIDCAuthorizationStore
+from .storage.quota import QuotaOverridesStore
 from .storage.token import TokenDatabaseStore, TokenRedisStore
 
 __all__ = ["Factory", "ProcessContext"]
@@ -657,10 +659,18 @@ def create_user_info_service(self) -> UserInfoService:
                 user_cache=self._context.ldap_user_cache,
                 logger=self._logger,
             )
+        quota_overrides_store = QuotaOverridesStore(
+            PydanticRedisStorage(
+                datatype=QuotaConfig, redis=self._context.persistent_redis
+            ),
+            self.create_slack_client(),
+            self._logger,
+        )
         return UserInfoService(
             config=self._context.config,
             ldap=ldap,
             firestore=firestore,
+            quota_overrides_store=quota_overrides_store,
             logger=self._logger,
         )
 

src/gafaelfawr/handlers/api.py

@@ -1,7 +1,7 @@
-"""Route handlers for the ``/auth/api/v1`` API.
+"""Route handlers for the token API.
 
 All the route handlers are intentionally defined in a single file to encourage
-the implementation to be very short.  All the business logic should be defined
+the implementation to be very short. All the business logic should be defined
 in manager objects and the output formatting should be handled by response
 models.
 """
@@ -30,6 +30,7 @@
 from ..models.auth import APIConfig, APILoginResponse, Scope
 from ..models.enums import TokenType
 from ..models.history import TokenChangeHistoryCursor, TokenChangeHistoryEntry
+from ..models.quota import QuotaConfig
 from ..models.token import (
     AdminTokenRequest,
     NewToken,
@@ -298,6 +299,74 @@ async def get_login(
     )
 
 
+@router.get(
+    "/auth/api/v1/quota-overrides",
+    description="Return the current quota overrides if any",
+    response_model_exclude_none=True,
+    responses={
+        404: {"description": "No quota overrides set", "model": ErrorModel}
+    },
+    summary="Get quota overrides",
+    tags=["admin"],
+)
+async def get_quota_overrides(
+    *,
+    auth_data: Annotated[TokenData, Depends(authenticate_admin_read)],
+    context: Annotated[RequestContext, Depends(context_dependency)],
+) -> QuotaConfig:
+    user_info_service = context.factory.create_user_info_service()
+    overrides = await user_info_service.get_quota_overrides()
+    if overrides:
+        return overrides
+    else:
+        raise HTTPException(
+            status_code=status.HTTP_404_NOT_FOUND,
+            detail=[{"type": "not_found", "msg": "No quota overrides set"}],
+        )
+
+
+@router.delete(
+    "/auth/api/v1/quota-overrides",
+    description="Remove any existing quota overrides",
+    responses={
+        404: {"description": "No quota overrides set", "model": ErrorModel}
+    },
+    status_code=204,
+    summary="Remove quota overrides",
+    tags=["admin"],
+)
+async def delete_quota_overrides(
+    *,
+    auth_data: Annotated[TokenData, Depends(authenticate_admin_write)],
+    context: Annotated[RequestContext, Depends(context_dependency)],
+) -> None:
+    user_info_service = context.factory.create_user_info_service()
+    success = await user_info_service.delete_quota_overrides()
+    if not success:
+        raise HTTPException(
+            status_code=status.HTTP_404_NOT_FOUND,
+            detail=[{"type": "not_found", "msg": "No quota overrides set"}],
+        )
+
+
+@router.put(
+    "/auth/api/v1/quota-overrides",
+    description="Set the quota overrides",
+    response_model_exclude_none=True,
+    summary="Set quota overrides",
+    tags=["admin"],
+)
+async def put_quota_overrides(
+    overrides: QuotaConfig,
+    *,
+    auth_data: Annotated[TokenData, Depends(authenticate_admin_write)],
+    context: Annotated[RequestContext, Depends(context_dependency)],
+) -> QuotaConfig:
+    user_info_service = context.factory.create_user_info_service()
+    await user_info_service.set_quota_overrides(overrides)
+    return overrides
+
+
 @router.get(
     "/auth/api/v1/token-info",
     description="Return metadata about the authentication token",

src/gafaelfawr/models/quota.py

@@ -0,0 +1,124 @@
+"""Models for user quotas."""
+
+from __future__ import annotations
+
+from pydantic import BaseModel, ConfigDict, Field
+
+__all__ = [
+    "NotebookQuota",
+    "Quota",
+    "QuotaConfig",
+]
+
+
+class NotebookQuota(BaseModel):
+    """Notebook Aspect quota information for a user."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    cpu: float = Field(..., title="CPU equivalents", examples=[4.0])
+
+    memory: float = Field(
+        ..., title="Maximum memory use (GiB)", examples=[16.0]
+    )
+
+    spawn: bool = Field(
+        True,
+        title="Spawning allowed",
+        description="Whether the user is allowed to spawn a notebook",
+    )
+
+
+class Quota(BaseModel):
+    """Quota information for a user."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    api: dict[str, int] = Field(
+        {},
+        title="API quotas",
+        description=(
+            "Mapping of service names to allowed requests per 15 minutes."
+        ),
+        examples=[
+            {
+                "datalinker": 500,
+                "hips": 2000,
+                "tap": 500,
+                "vo-cutouts": 100,
+            }
+        ],
+    )
+
+    notebook: NotebookQuota | None = Field(
+        None, title="Notebook Aspect quotas"
+    )
+
+
+class QuotaConfig(BaseModel):
+    """Quota configuration."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    default: Quota = Field(
+        ..., title="Default quota", description="Default quotas for all users"
+    )
+
+    groups: dict[str, Quota] = Field(
+        {},
+        title="Quota grants by group",
+        description="Additional quota grants by group name",
+    )
+
+    bypass: set[str] = Field(
+        set(),
+        title="Groups without quotas",
+        description="Groups whose members bypass all quota restrictions",
+    )
+
+    def calculate_quota(self, groups: set[str]) -> Quota | None:
+        """Calculate user's quota given their group membership.
+
+        Parameters
+        ----------
+        groups
+            Group membership of the user.
+
+        Returns
+        -------
+        Quota or None
+            Quota information for that user or `None` if no quotas apply. If
+            the user bypasses quotas, a `~gafaelfawr.models.quota.Quota` model
+            with quotas set to `None` or an empty dictionary is returned rather
+            than `None`.
+        """
+        if groups & self.bypass:
+            return Quota()
+
+        # Start with the defaults.
+        api = dict(self.default.api)
+        notebook = None
+        if self.default.notebook:
+            notebook = self.default.notebook.model_copy()
+
+        # Look for group-specific rules.
+        for group in groups & set(self.groups.keys()):
+            extra = self.groups[group]
+            if extra.notebook:
+                if notebook:
+                    notebook.cpu += extra.notebook.cpu
+                    notebook.memory += extra.notebook.memory
+                    notebook.spawn &= extra.notebook.spawn
+                else:
+                    notebook = extra.notebook.model_copy()
+            for service, quota in extra.api.items():
+                if service in api:
+                    api[service] += quota
+                else:
+                    api[service] = quota
+
+        # Return the results.
+        if not notebook and not api:
+            return None
+        else:
+            return Quota(api=api, notebook=notebook)