diff --git a/changelog.d/20250114_135624_rra_DM_48390.md b/changelog.d/20250114_135624_rra_DM_48390.md new file mode 100644 index 00000000..742f9138 --- /dev/null +++ b/changelog.d/20250114_135624_rra_DM_48390.md @@ -0,0 +1,3 @@ +### New features + +- Add a `bypass` key to the quota configuration containing a group list. Any member of one of those groups ignores all quota restrictions. \ No newline at end of file diff --git a/src/gafaelfawr/config.py b/src/gafaelfawr/config.py index e0496bc3..4136aaea 100644 --- a/src/gafaelfawr/config.py +++ b/src/gafaelfawr/config.py @@ -712,6 +712,12 @@ class QuotaConfig(BaseModel): description="Additional quota grants by group name", ) + bypass: set[str] = Field( + set(), + title="Groups without quotas", + description="Groups whose members bypass all quota restrictions", + ) + class GitHubGroupTeam(BaseModel): """Specification for a GitHub team.""" diff --git a/src/gafaelfawr/services/userinfo.py b/src/gafaelfawr/services/userinfo.py index 760aed11..7d1e23e8 100644 --- a/src/gafaelfawr/services/userinfo.py +++ b/src/gafaelfawr/services/userinfo.py @@ -225,6 +225,11 @@ def _calculate_quota(self, groups: list[Group]) -> Quota | None: """ if not self._config.quota: return None + group_names = {g.name for g in groups} + if group_names & self._config.quota.bypass: + return Quota() + + # Start with the defaults. api = dict(self._config.quota.default.api) notebook = None if self._config.quota.default.notebook: @@ -232,23 +237,28 @@ def _calculate_quota(self, groups: list[Group]) -> Quota | None: cpu=self._config.quota.default.notebook.cpu, memory=self._config.quota.default.notebook.memory, ) - for group in groups or []: - if group.name in self._config.quota.groups: - extra = self._config.quota.groups[group.name] - if extra.notebook: - if notebook: - notebook.cpu += extra.notebook.cpu - notebook.memory += extra.notebook.memory - else: - notebook = NotebookQuota( - cpu=extra.notebook.cpu, - memory=extra.notebook.memory, - ) - for service in extra.api: - if service in api: - api[service] += extra.api[service] - else: - api[service] = extra.api[service] + + # Look for group-specific rules. + for group in group_names: + if group not in self._config.quota.groups: + continue + extra = self._config.quota.groups[group] + if extra.notebook: + if notebook: + notebook.cpu += extra.notebook.cpu + notebook.memory += extra.notebook.memory + else: + notebook = NotebookQuota( + cpu=extra.notebook.cpu, + memory=extra.notebook.memory, + ) + for service in extra.api: + if service in api: + api[service] += extra.api[service] + else: + api[service] = extra.api[service] + + # Return the results. return Quota(api=api, notebook=notebook) async def _get_groups_from_ldap( diff --git a/tests/data/config/github-quota.yaml b/tests/data/config/github-quota.yaml index 8877e803..ce13c345 100644 --- a/tests/data/config/github-quota.yaml +++ b/tests/data/config/github-quota.yaml @@ -32,6 +32,8 @@ quota: notebook: cpu: 0.0 memory: 4.0 + bypass: + - "admin" metrics: enabled: false application: "gafaelfawr" diff --git a/tests/handlers/ingress_rate_test.py b/tests/handlers/ingress_rate_test.py index bb2b5e95..1c96442c 100644 --- a/tests/handlers/ingress_rate_test.py +++ b/tests/handlers/ingress_rate_test.py @@ -67,3 +67,26 @@ async def test_rate_limit(client: AsyncClient, factory: Factory) -> None: assert r.headers["X-RateLimit-Resource"] == "test" reset = int(r.headers["X-RateLimit-Reset"]) assert expected.timestamp() <= reset <= expected.timestamp() + 5 + + +@pytest.mark.asyncio +async def test_rate_limit_bypass( + client: AsyncClient, factory: Factory +) -> None: + await reconfigure("github-quota", factory) + token_data = await create_session_token( + factory, group_names=["admin"], scopes={"read:all"} + ) + r = await client.get( + "/ingress/auth", + params={"scope": "read:all", "service": "test"}, + headers={"Authorization": f"Bearer {token_data.token}"}, + ) + assert r.status_code == 200 + + # There should be no quota set since the user is in the bypass group. + assert "X-RateLimit-Limit" not in r.headers + assert "X-RateLimit-Remaining" not in r.headers + assert "X-RateLimit-Used" not in r.headers + assert "X-RateLimit-Resource" not in r.headers + assert "X-RateLimit-Reset" not in r.headers