Merge pull request #1219 from lsst-sqre:tickets/DM-48432

DM-48432: Add support for quota overrides

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)