From 4076b067ba34db9c9dd704d1b76484fb7c38eb2b Mon Sep 17 00:00:00 2001 From: Russ Allbery Date: Thu, 16 Jan 2025 17:13:29 -0800 Subject: [PATCH] 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 | 3 + docs/dev/internals.rst | 3 + src/gafaelfawr/config.py | 65 +------- src/gafaelfawr/factory.py | 12 +- src/gafaelfawr/handlers/api.py | 73 ++++++++- src/gafaelfawr/models/quota.py | 124 ++++++++++++++++ src/gafaelfawr/models/userinfo.py | 48 +----- src/gafaelfawr/services/userinfo.py | 82 +++++++++- src/gafaelfawr/storage/quota.py | 73 +++++++++ tests/handlers/ingress_rate_test.py | 11 +- tests/handlers/quota_test.py | 156 ++++++++++++++++++++ 11 files changed, 526 insertions(+), 124 deletions(-) create mode 100644 changelog.d/20250116_152452_rra_DM_48432.md create mode 100644 src/gafaelfawr/models/quota.py create mode 100644 src/gafaelfawr/storage/quota.py diff --git a/changelog.d/20250116_152452_rra_DM_48432.md b/changelog.d/20250116_152452_rra_DM_48432.md new file mode 100644 index 00000000..1adeb7ea --- /dev/null +++ b/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. \ No newline at end of file diff --git a/docs/dev/internals.rst b/docs/dev/internals.rst index 94db04a0..27f4b98f 100644 --- a/docs/dev/internals.rst +++ b/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: diff --git a/src/gafaelfawr/config.py b/src/gafaelfawr/config.py index 2aae8e3a..ce30d7bc 100644 --- a/src/gafaelfawr/config.py +++ b/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.""" diff --git a/src/gafaelfawr/factory.py b/src/gafaelfawr/factory.py index 754161a9..35f15d12 100644 --- a/src/gafaelfawr/factory.py +++ b/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, ) diff --git a/src/gafaelfawr/handlers/api.py b/src/gafaelfawr/handlers/api.py index 520e9e6e..3b429389 100644 --- a/src/gafaelfawr/handlers/api.py +++ b/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", diff --git a/src/gafaelfawr/models/quota.py b/src/gafaelfawr/models/quota.py new file mode 100644 index 00000000..065956f8 --- /dev/null +++ b/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) diff --git a/src/gafaelfawr/models/userinfo.py b/src/gafaelfawr/models/userinfo.py index c3baeb8f..685c76db 100644 --- a/src/gafaelfawr/models/userinfo.py +++ b/src/gafaelfawr/models/userinfo.py @@ -5,15 +5,15 @@ from dataclasses import dataclass from datetime import datetime -from pydantic import BaseModel, ConfigDict, Field +from pydantic import BaseModel, Field from ..constants import GROUPNAME_REGEX from ..pydantic import Timestamp +from .quota import Quota __all__ = [ "CADCUserInfo", "Group", - "NotebookQuota", "Quota", "RateLimitStatus", "UserInfo", @@ -74,50 +74,6 @@ class Group(BaseModel): ) -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 UserInfo(BaseModel): """Metadata about a user. diff --git a/src/gafaelfawr/services/userinfo.py b/src/gafaelfawr/services/userinfo.py index b1040062..ddf18121 100644 --- a/src/gafaelfawr/services/userinfo.py +++ b/src/gafaelfawr/services/userinfo.py @@ -7,8 +7,10 @@ from ..config import Config from ..exceptions import FirestoreError from ..models.ldap import LDAPUserData +from ..models.quota import Quota, QuotaConfig from ..models.token import TokenData, TokenUserInfo from ..models.userinfo import Group, UserInfo +from ..storage.quota import QuotaOverridesStore from .firestore import FirestoreService from .ldap import LDAPService @@ -36,6 +38,8 @@ class UserInfoService: LDAP service for user metadata, if LDAP was configured. firestore Service for Firestore UID/GID lookups, if Firestore was configured. + quota_overrides_store + Storage for quota overrides. logger Logger to use. """ @@ -46,13 +50,35 @@ def __init__( config: Config, ldap: LDAPService | None, firestore: FirestoreService | None, + quota_overrides_store: QuotaOverridesStore, logger: BoundLogger, ) -> None: self._config = config self._ldap = ldap self._firestore = firestore + self._quota_overrides = quota_overrides_store self._logger = logger + async def delete_quota_overrides(self) -> bool: + """Delete any existing quota overrides. + + Returns + ------- + bool + `True` if quota overrides were deleted, `False` if none were set. + """ + return await self._quota_overrides.delete() + + async def get_quota_overrides(self) -> QuotaConfig | None: + """Get the current quota overrides, if any. + + Returns + ------- + QuotaConfig or None + Current quota overrides, or `None` if there are none. + """ + return await self._quota_overrides.get() + async def get_user_info_from_token( self, token_data: TokenData, *, uncached: bool = False ) -> UserInfo: @@ -112,12 +138,6 @@ async def get_user_info_from_token( if not gid and not ldap_data.gid and self._config.add_user_group: gid = uid or ldap_data.uid - # Calculate the quota. - quota = None - if self._config.quota: - group_names = {g.name for g in groups} - quota = self._config.quota.calculate_quota(group_names) - # Return the results. return UserInfo( username=username, @@ -126,7 +146,7 @@ async def get_user_info_from_token( gid=gid or ldap_data.gid, email=token_data.email or ldap_data.email, groups=sorted(groups, key=lambda g: g.name), - quota=quota, + quota=await self._calculate_quota(groups), ) async def get_scopes(self, user_info: TokenUserInfo) -> set[str] | None: @@ -216,6 +236,54 @@ async def invalidate_cache(self, username: str) -> None: if self._ldap: await self._ldap.invalidate_cache(username) + async def set_quota_overrides(self, overrides: QuotaConfig) -> None: + """Store quota overrides, overwriting any existing ones. + + Parameters + ---------- + overrides + New quota overrides to store. + """ + return await self._quota_overrides.store(overrides) + + async def _calculate_quota(self, groups: list[Group]) -> Quota | None: + """Calculate the quota for a user. + + Parameters + ---------- + groups + Group membership of the user. + + Returns + ------- + Quota or None + Quota information for that user, or `None` if no quotas apply. + """ + group_names = {g.name for g in groups} + quota = None + if self._config.quota: + quota = self._config.quota.calculate_quota(group_names) + + # Check if there are quota overrides. + overrides = await self.get_quota_overrides() + if not overrides: + return quota + + # Apply the override on top of the existing quota, if any. + override_quota = overrides.calculate_quota(group_names) + if not override_quota: + return quota + elif not quota: + return override_quota + elif overrides.bypass & group_names: + return Quota() + else: + api = quota.api + api.update(override_quota.api) + return Quota( + notebook=override_quota.notebook or quota.notebook, api=api + ) + async def _get_groups_from_ldap( self, username: str, diff --git a/src/gafaelfawr/storage/quota.py b/src/gafaelfawr/storage/quota.py new file mode 100644 index 00000000..19155685 --- /dev/null +++ b/src/gafaelfawr/storage/quota.py @@ -0,0 +1,73 @@ +"""Storage layer for quota overrides.""" + +from __future__ import annotations + +from safir.redis import DeserializeError, PydanticRedisStorage +from safir.slack.webhook import SlackWebhookClient +from structlog.stdlib import BoundLogger + +from ..models.quota import QuotaConfig + +__all__ = ["QuotaOverridesStore"] + + +class QuotaOverridesStore: + """Stores and retrieves quota overrides in Redis. + + Parameters + ---------- + storage + Underlying storage for quota overrides. + slack_client + If provided, Slack webhook client to report deserialization errors of + Redis data. + logger + Logger for diagnostics. + """ + + def __init__( + self, + storage: PydanticRedisStorage[QuotaConfig], + slack_client: SlackWebhookClient | None, + logger: BoundLogger, + ) -> None: + self._storage = storage + self._slack = slack_client + self._logger = logger + + async def delete(self) -> bool: + """Delete any stored quota overrides. + + Returns + ------- + bool + `True` if there were quota overrides to delete, `False` otherwise. + """ + return await self._storage.delete("quota-overrides") + + async def get(self) -> QuotaConfig | None: + """Retrieve quota overrides from Redis, if any. + + Returns + ------- + QuotaConfig or None + Quota overrides if any are set, or `None` if there are none. + """ + try: + return await self._storage.get("quota-overrides") + except DeserializeError as e: + msg = "Cannot retrieve quota overrides" + self._logger.exception(msg, error=str(e)) + if self._slack: + await self._slack.post_exception(e) + return None + + async def store(self, overrides: QuotaConfig) -> None: + """Store quota overrides in Redis. + + Parameters + ---------- + overrides + Overrides to store, replacing any existing overrides. + """ + await self._storage.store("quota-overrides", overrides) diff --git a/tests/handlers/ingress_rate_test.py b/tests/handlers/ingress_rate_test.py index 1c96442c..ff8c168c 100644 --- a/tests/handlers/ingress_rate_test.py +++ b/tests/handlers/ingress_rate_test.py @@ -20,6 +20,7 @@ async def test_rate_limit(client: AsyncClient, factory: Factory) -> None: token_data = await create_session_token( factory, group_names=["foo"], scopes={"read:all"} ) + headers = {"Authorization": f"bearer {token_data.token}"} now = datetime.now(tz=UTC) expected = now + timedelta(minutes=15) - timedelta(seconds=1) @@ -28,7 +29,7 @@ async def test_rate_limit(client: AsyncClient, factory: Factory) -> None: r = await client.get( "/ingress/auth", params={"scope": "read:all", "service": "test"}, - headers={"Authorization": f"Bearer {token_data.token}"}, + headers=headers, ) assert r.status_code == 200 assert r.headers["X-RateLimit-Limit"] == "2" @@ -40,7 +41,7 @@ async def test_rate_limit(client: AsyncClient, factory: Factory) -> None: r = await client.get( "/ingress/auth", params={"scope": "read:all", "service": "test"}, - headers={"Authorization": f"Bearer {token_data.token}"}, + headers=headers, ) assert r.status_code == 200 assert r.headers["X-RateLimit-Limit"] == "2" @@ -55,7 +56,7 @@ async def test_rate_limit(client: AsyncClient, factory: Factory) -> None: r = await client.get( "/ingress/auth", params={"scope": "read:all", "service": "test"}, - headers={"Authorization": f"Bearer {token_data.token}"}, + headers=headers, ) assert r.status_code == 429 retry_after = parsedate_to_datetime(r.headers["Retry-After"]) @@ -77,10 +78,12 @@ async def test_rate_limit_bypass( token_data = await create_session_token( factory, group_names=["admin"], scopes={"read:all"} ) + headers = {"Authorization": f"bearer {token_data.token}"} + r = await client.get( "/ingress/auth", params={"scope": "read:all", "service": "test"}, - headers={"Authorization": f"Bearer {token_data.token}"}, + headers=headers, ) assert r.status_code == 200 diff --git a/tests/handlers/quota_test.py b/tests/handlers/quota_test.py index 0b5d503d..980c4690 100644 --- a/tests/handlers/quota_test.py +++ b/tests/handlers/quota_test.py @@ -2,9 +2,12 @@ from __future__ import annotations +from typing import Any + import pytest from httpx import AsyncClient +from gafaelfawr.config import Config from gafaelfawr.factory import Factory from gafaelfawr.models.token import TokenUserInfo from gafaelfawr.models.userinfo import Group @@ -84,3 +87,156 @@ async def test_no_spawn(client: AsyncClient, factory: Factory) -> None: "notebook": {"cpu": 8.0, "memory": 4.0, "spawn": False}, }, } + + +@pytest.mark.asyncio +async def test_rate_limit_override( + client: AsyncClient, factory: Factory +) -> None: + config = await reconfigure("github-quota", factory) + assert config.quota + token_data = await create_session_token( + factory, + group_names=["foo"], + scopes={"admin:token", "read:all"}, + ) + assert token_data.groups + default_quota = config.quota.calculate_quota({"foo"}) + assert default_quota + headers = {"Authorization": f"bearer {token_data.token}"} + + overrides: dict[str, Any] = { + "bypass": [], + "default": {"api": {"test": 10}}, + "groups": {}, + } + r = await client.put( + "/auth/api/v1/quota-overrides", json=overrides, headers=headers + ) + assert r.status_code == 200 + r = await client.get( + "/ingress/auth", + params={"scope": "read:all", "service": "test"}, + headers=headers, + ) + assert r.status_code == 200 + assert r.headers["X-RateLimit-Limit"] == "10" + assert r.headers["X-RateLimit-Remaining"] == "9" + + r = await client.get("/auth/api/v1/user-info", headers=headers) + expected_user_info: dict[str, Any] = { + "username": token_data.username, + "name": token_data.name, + "email": token_data.email, + "uid": token_data.uid, + "gid": token_data.gid, + "groups": [ + g.model_dump(mode="json") + for g in sorted(token_data.groups, key=lambda g: g.name) + ], + "quota": { + "api": {"datalinker": 1000, "test": 10}, + "notebook": {"cpu": 8.0, "memory": 8.0, "spawn": True}, + }, + } + assert r.json() == expected_user_info + + overrides["default"]["notebook"] = {"cpu": 1, "memory": 1, "spawn": False} + r = await client.put( + "/auth/api/v1/quota-overrides", json=overrides, headers=headers + ) + assert r.status_code == 200 + expected_user_info["quota"]["notebook"] = overrides["default"]["notebook"] + r = await client.get("/auth/api/v1/user-info", headers=headers) + assert r.json() == expected_user_info + + overrides["bypass"] = ["foo"] + r = await client.put( + "/auth/api/v1/quota-overrides", json=overrides, headers=headers + ) + assert r.status_code == 200 + expected_user_info["quota"] = {"api": {}} + r = await client.get("/auth/api/v1/user-info", headers=headers) + assert r.json() == expected_user_info + + # Return to normal behavior by deleting the overrides. + r = await client.delete("/auth/api/v1/quota-overrides", headers=headers) + assert r.status_code == 204 + expected_user_info["quota"] = default_quota.model_dump(mode="json") + r = await client.get("/auth/api/v1/user-info", headers=headers) + assert r.json() == expected_user_info + + +@pytest.mark.asyncio +async def test_rate_limit_override_only( + client: AsyncClient, factory: Factory, config: Config +) -> None: + """Check behavior when there is an override and no base quota.""" + assert not config.quota + token_data = await create_session_token( + factory, group_names=["admin"], scopes={"admin:token", "read:all"} + ) + assert token_data.groups + headers = {"Authorization": f"bearer {token_data.token}"} + + r = await client.get( + "/ingress/auth", + params={"scope": "read:all", "service": "test"}, + headers=headers, + ) + assert r.status_code == 200 + assert "X-RateLimit-Limit" not in r.headers + + overrides: dict[str, Any] = { + "bypass": [], + "default": { + "notebook": {"cpu": 1.0, "memory": 4.0, "spawn": True}, + "api": {"test": 10}, + }, + "groups": {}, + } + r = await client.put( + "/auth/api/v1/quota-overrides", json=overrides, headers=headers + ) + assert r.status_code == 200 + assert r.json() == overrides + r = await client.get( + "/ingress/auth", + params={"scope": "read:all", "service": "test"}, + headers=headers, + ) + assert r.status_code == 200 + assert r.headers["X-RateLimit-Limit"] == "10" + assert r.headers["X-RateLimit-Remaining"] == "9" + assert r.headers["X-RateLimit-Used"] == "1" + assert r.headers["X-RateLimit-Resource"] == "test" + + r = await client.get("/auth/api/v1/user-info", headers=headers) + expected_user_info = { + "username": token_data.username, + "name": token_data.name, + "email": token_data.email, + "uid": token_data.uid, + "gid": token_data.gid, + "groups": [ + g.model_dump(mode="json") + for g in sorted(token_data.groups, key=lambda g: g.name) + ], + "quota": { + "api": {"test": 10}, + "notebook": {"cpu": 1.0, "memory": 4.0, "spawn": True}, + }, + } + assert r.json() == expected_user_info + + # Add the user's group to bypass. + overrides["bypass"] = ["admin"] + r = await client.put( + "/auth/api/v1/quota-overrides", json=overrides, headers=headers + ) + assert r.status_code == 200 + assert r.json() == overrides + expected_user_info["quota"] = {"api": {}} + r = await client.get("/auth/api/v1/user-info", headers=headers) + assert r.status_code == 200 + assert r.json() == expected_user_info