mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-13 18:15:02 +00:00
feat(quotas): limit resolution and the quota service
docsgpt/quotas resolves a user's effective token and cost limits from the policy rows that apply to them: their own override, then the most generous allowance among their teams, then the instance default, then any defaults a registered provider supplies. Each budget resolves on its own, and a team counts once however many memberships the user holds in it. QuotaService compares those limits with the user's token_usage totals over the current QUOTA_PERIOD window (calendar-aligned, UTC, computed at read time). It fails open, and skips the usage query for unlimited users.
This commit is contained in:
Alex committed
1 parent
6c42139224
commit
6826313b60
12 files changed
+1175
No files matched your search
@@ -0,0 +1,24 @@
|
||||
"""Admin-set usage quotas.
|
||||
|
||||
Limits live in ``quota_policies`` at three layers (instance default, team
|
||||
per-member allowance, user override), each with a token budget and a USD
|
||||
budget. ``QuotaService`` resolves a user's effective limits and compares them
|
||||
with their ``token_usage`` totals over the current ``QUOTA_PERIOD`` window.
|
||||
"""
|
||||
|
||||
from docsgpt.quotas.providers import QuotaDefaultsProvider, register_defaults_provider
|
||||
from docsgpt.quotas.resolver import ResolvedLimit, ResolvedLimits, resolve_limits
|
||||
from docsgpt.quotas.service import BucketStatus, QuotaExceeded, QuotaService
|
||||
from docsgpt.quotas.windows import window_bounds
|
||||
|
||||
__all__ = [
|
||||
"BucketStatus",
|
||||
"QuotaDefaultsProvider",
|
||||
"QuotaExceeded",
|
||||
"QuotaService",
|
||||
"ResolvedLimit",
|
||||
"ResolvedLimits",
|
||||
"register_defaults_provider",
|
||||
"resolve_limits",
|
||||
"window_bounds",
|
||||
]
|
||||
@@ -0,0 +1,51 @@
|
||||
"""Extension point for limits that do not come from ``quota_policies``.
|
||||
|
||||
A deployment can register a provider that supplies per-user default policies
|
||||
(for example from a subscription plan) and adjusts the quota error payload.
|
||||
Defaults sit below every stored layer: a stored instance, team or user row
|
||||
with an opinion always wins.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Mapping
|
||||
|
||||
|
||||
class QuotaDefaultsProvider:
|
||||
"""Base provider: no defaults, error payload unchanged."""
|
||||
|
||||
def default_policies(self, user_id: str) -> list[dict]:
|
||||
"""Return default policy rows for ``user_id``.
|
||||
|
||||
Each row uses the ``quota_policies`` field names (``bucket``,
|
||||
``token_limit``, ``token_unlimited``, ``cost_limit_usd``,
|
||||
``cost_unlimited``); ``scope`` is set by the caller.
|
||||
"""
|
||||
return []
|
||||
|
||||
def error_payload(self, payload: dict, user_id: str) -> dict:
|
||||
"""Return the payload sent to a client whose quota is exhausted."""
|
||||
return payload
|
||||
|
||||
|
||||
_provider: QuotaDefaultsProvider = QuotaDefaultsProvider()
|
||||
|
||||
|
||||
def register_defaults_provider(provider: QuotaDefaultsProvider) -> None:
|
||||
"""Replace the process-wide defaults provider."""
|
||||
global _provider
|
||||
_provider = provider
|
||||
|
||||
|
||||
def get_defaults_provider() -> QuotaDefaultsProvider:
|
||||
"""Return the registered defaults provider."""
|
||||
return _provider
|
||||
|
||||
|
||||
def default_rows(user_id: str) -> list[dict]:
|
||||
"""Return the provider's defaults for ``user_id`` as ``default``-layer rows."""
|
||||
rows: list[dict] = []
|
||||
for row in get_defaults_provider().default_policies(user_id) or []:
|
||||
if isinstance(row, Mapping):
|
||||
rows.append({"bucket": "all", "enabled": True, **row, "scope": "default", "subject_id": None})
|
||||
return rows
|
||||
@@ -0,0 +1,99 @@
|
||||
"""Resolve the policy rows that apply to a user into effective limits.
|
||||
|
||||
Each budget (tokens, cost) resolves on its own: the user's row wins, then the
|
||||
most generous of the user's team rows, then the instance row, then the
|
||||
registered defaults. A row with neither a limit nor the unlimited flag for a
|
||||
budget has no opinion on it and is skipped.
|
||||
|
||||
Teams resolve to the most generous allowance because team membership is not
|
||||
controlled by the instance admin: under "most restrictive", any team admin
|
||||
could throttle a user by adding them to a low-allowance team.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Iterable, Mapping, Optional
|
||||
|
||||
LAYERS = ("user", "team", "instance", "default")
|
||||
|
||||
_FIELDS = {"tokens": ("token_limit", "token_unlimited"), "cost": ("cost_limit_usd", "cost_unlimited")}
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ResolvedLimit:
|
||||
"""One budget's effective limit and the layer it came from.
|
||||
|
||||
``limit`` is ``None`` when the budget is unlimited. ``source`` is ``None``
|
||||
when no layer had an opinion, otherwise one of ``LAYERS``; ``source_id`` is
|
||||
the team id for a team-sourced limit.
|
||||
"""
|
||||
|
||||
limit: Optional[float] = None
|
||||
source: Optional[str] = None
|
||||
source_id: Optional[str] = None
|
||||
|
||||
@property
|
||||
def unlimited(self) -> bool:
|
||||
return self.limit is None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ResolvedLimits:
|
||||
"""A user's effective token and cost limits for one bucket."""
|
||||
|
||||
tokens: ResolvedLimit
|
||||
cost: ResolvedLimit
|
||||
|
||||
@property
|
||||
def unlimited(self) -> bool:
|
||||
return self.tokens.unlimited and self.cost.unlimited
|
||||
|
||||
|
||||
def _opinion(row: Mapping, budget: str) -> Optional[tuple[bool, Optional[float]]]:
|
||||
"""Return ``(unlimited, limit)`` for a row's budget, or ``None`` if it defers."""
|
||||
limit_field, unlimited_field = _FIELDS[budget]
|
||||
if row.get(unlimited_field):
|
||||
return True, None
|
||||
value = row.get(limit_field)
|
||||
if value is None:
|
||||
return None
|
||||
return False, float(value)
|
||||
|
||||
|
||||
def _resolve_budget(rows: Iterable[Mapping], budget: str) -> ResolvedLimit:
|
||||
by_layer: dict[str, list[tuple[Mapping, tuple[bool, Optional[float]]]]] = {}
|
||||
for row in rows:
|
||||
opinion = _opinion(row, budget)
|
||||
if opinion is not None:
|
||||
by_layer.setdefault(row["scope"], []).append((row, opinion))
|
||||
for layer in LAYERS:
|
||||
candidates = by_layer.get(layer)
|
||||
if not candidates:
|
||||
continue
|
||||
# Most generous first: unlimited, then the larger limit. Only the team
|
||||
# layer can hold more than one candidate. Ties break on subject id so
|
||||
# the reported source is stable.
|
||||
row, (unlimited, limit) = min(
|
||||
candidates,
|
||||
key=lambda c: (not c[1][0], -(c[1][1] or 0.0), str(c[0].get("subject_id") or "")),
|
||||
)
|
||||
source_id = str(row["subject_id"]) if layer == "team" else None
|
||||
return ResolvedLimit(limit=None if unlimited else limit, source=layer, source_id=source_id)
|
||||
return ResolvedLimit()
|
||||
|
||||
|
||||
def resolve_limits(rows: Iterable[Mapping], bucket: str = "all") -> ResolvedLimits:
|
||||
"""Return the effective limits for ``bucket`` from a user's applicable rows.
|
||||
|
||||
Args:
|
||||
rows: Policy rows that apply to the user (their own, their teams', the
|
||||
instance's and any provider defaults). Disabled rows and rows for
|
||||
other buckets are ignored.
|
||||
bucket: The policy bucket to resolve.
|
||||
"""
|
||||
applicable = [r for r in rows if r.get("bucket", "all") == bucket and r.get("enabled", True)]
|
||||
return ResolvedLimits(
|
||||
tokens=_resolve_budget(applicable, "tokens"),
|
||||
cost=_resolve_budget(applicable, "cost"),
|
||||
)
|
||||
@@ -0,0 +1,174 @@
|
||||
"""Compare a user's usage with their effective limits."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from typing import Optional
|
||||
|
||||
from docsgpt.core.settings import settings
|
||||
from docsgpt.quotas.providers import default_rows, get_defaults_provider
|
||||
from docsgpt.quotas.resolver import ResolvedLimit, ResolvedLimits, resolve_limits
|
||||
from docsgpt.quotas.windows import window_bounds
|
||||
from docsgpt.storage.db.repositories.quota_policies import QuotaPoliciesRepository
|
||||
from docsgpt.storage.db.repositories.token_usage import TokenUsageRepository
|
||||
from docsgpt.storage.db.session import db_readonly
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
REQUEST_BUCKETS = ("direct", "agent")
|
||||
|
||||
_UNITS = {"tokens": "tokens", "cost": "USD"}
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class BucketStatus:
|
||||
"""A user's limits and usage for one policy bucket in the current window."""
|
||||
|
||||
bucket: str
|
||||
limits: ResolvedLimits
|
||||
tokens_used: int
|
||||
cost_used: float
|
||||
resets_at: datetime
|
||||
|
||||
def exceeded_budget(self) -> Optional[str]:
|
||||
"""Return ``tokens`` or ``cost`` when that budget is used up, else ``None``."""
|
||||
if not self.limits.tokens.unlimited and self.tokens_used >= self.limits.tokens.limit:
|
||||
return "tokens"
|
||||
if not self.limits.cost.unlimited and self.cost_used >= self.limits.cost.limit:
|
||||
return "cost"
|
||||
return None
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
"""Return the JSON shape shared by the admin and user quota endpoints."""
|
||||
|
||||
def budget(limit: ResolvedLimit, used: float) -> dict:
|
||||
return {
|
||||
"limit": limit.limit,
|
||||
"used": used,
|
||||
"source": limit.source,
|
||||
"source_id": limit.source_id,
|
||||
}
|
||||
|
||||
return {
|
||||
"bucket": self.bucket,
|
||||
"tokens": budget(self.limits.tokens, self.tokens_used),
|
||||
"cost": budget(self.limits.cost, round(self.cost_used, 6)),
|
||||
"resets_at": self.resets_at.isoformat(),
|
||||
}
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class QuotaExceeded:
|
||||
"""The exhausted budget that blocks a request."""
|
||||
|
||||
user_id: str
|
||||
bucket: str
|
||||
budget: str
|
||||
usage: float
|
||||
limit: float
|
||||
source: Optional[str]
|
||||
source_id: Optional[str]
|
||||
resets_at: datetime
|
||||
|
||||
@property
|
||||
def retry_after_seconds(self) -> int:
|
||||
now = datetime.now(self.resets_at.tzinfo)
|
||||
return max(int((self.resets_at - now).total_seconds()), 1)
|
||||
|
||||
def to_payload(self) -> dict:
|
||||
"""Return the client-facing error body, after the provider's adjustments."""
|
||||
unit = _UNITS[self.budget]
|
||||
if self.budget == "cost":
|
||||
amounts = f"${self.usage:.2f} of ${self.limit:.2f}"
|
||||
else:
|
||||
amounts = f"{int(self.usage):,} of {int(self.limit):,} tokens"
|
||||
payload = {
|
||||
"success": False,
|
||||
"error_code": "quota-exceeded",
|
||||
"message": f"Usage quota reached ({amounts}). It resets at {self.resets_at.isoformat()}.",
|
||||
"limit_scope": "user_quota",
|
||||
"dimension": self.budget,
|
||||
"unit": unit,
|
||||
"usage": round(self.usage, 6) if self.budget == "cost" else int(self.usage),
|
||||
"limit": self.limit if self.budget == "cost" else int(self.limit),
|
||||
"bucket": self.bucket,
|
||||
"source": self.source,
|
||||
"resets_at": self.resets_at.isoformat(),
|
||||
}
|
||||
try:
|
||||
return get_defaults_provider().error_payload(payload, self.user_id) or payload
|
||||
except Exception:
|
||||
logger.exception("quota defaults provider failed to build the error payload")
|
||||
return payload
|
||||
|
||||
|
||||
class QuotaService:
|
||||
"""Resolve limits and measure usage for the current ``QUOTA_PERIOD`` window."""
|
||||
|
||||
@staticmethod
|
||||
def status(
|
||||
user_id: str,
|
||||
buckets: tuple[str, ...] = ("all",),
|
||||
now: Optional[datetime] = None,
|
||||
) -> list[BucketStatus]:
|
||||
"""Return the user's status for each of ``buckets``.
|
||||
|
||||
Usage is only summed for buckets that carry a limit, so a user with no
|
||||
applicable policy costs one policy lookup and no usage query.
|
||||
"""
|
||||
start, resets_at = window_bounds(settings.QUOTA_PERIOD, now)
|
||||
statuses: list[BucketStatus] = []
|
||||
with db_readonly() as conn:
|
||||
rows = QuotaPoliciesRepository(conn).policies_for_user(user_id) + default_rows(user_id)
|
||||
usage_repo = TokenUsageRepository(conn)
|
||||
for bucket in buckets:
|
||||
limits = resolve_limits(rows, bucket)
|
||||
tokens_used, cost_used = (0, 0.0)
|
||||
if not limits.unlimited:
|
||||
tokens_used, cost_used = usage_repo.usage_totals(user_id=user_id, start=start, bucket=bucket)
|
||||
statuses.append(BucketStatus(bucket, limits, tokens_used, cost_used, resets_at))
|
||||
return statuses
|
||||
|
||||
@classmethod
|
||||
def check(
|
||||
cls, user_id: Optional[str], bucket: str = "direct", now: Optional[datetime] = None
|
||||
) -> Optional[QuotaExceeded]:
|
||||
"""Return why ``user_id`` may not start a request, or ``None`` if they may.
|
||||
|
||||
Both the ``all`` policies and the request's own bucket must have room.
|
||||
The check runs before the request, so the call that crosses a limit
|
||||
completes and the next one is refused. Any failure here allows the
|
||||
request: a quota outage must not take chat down.
|
||||
|
||||
Args:
|
||||
user_id: The billable user. Requests with no user are not limited.
|
||||
bucket: ``direct`` for chat without an agent, ``agent`` for traffic through one.
|
||||
now: Reference instant, for tests.
|
||||
"""
|
||||
if not user_id:
|
||||
return None
|
||||
if bucket not in REQUEST_BUCKETS:
|
||||
raise ValueError(f"unknown request bucket: {bucket!r}")
|
||||
try:
|
||||
statuses = cls.status(user_id, ("all", bucket), now)
|
||||
except Exception:
|
||||
logger.exception("quota check failed; allowing the request", extra={"user_id": user_id})
|
||||
return None
|
||||
for status in statuses:
|
||||
budget = status.exceeded_budget()
|
||||
if budget is None:
|
||||
continue
|
||||
limit = status.limits.tokens if budget == "tokens" else status.limits.cost
|
||||
return QuotaExceeded(
|
||||
user_id=user_id,
|
||||
bucket=status.bucket,
|
||||
budget=budget,
|
||||
usage=status.tokens_used if budget == "tokens" else status.cost_used,
|
||||
limit=limit.limit,
|
||||
source=limit.source,
|
||||
source_id=limit.source_id,
|
||||
resets_at=status.resets_at,
|
||||
)
|
||||
return None
|
||||
@@ -0,0 +1,37 @@
|
||||
"""Calendar-aligned UTC quota windows, computed at read time (no reset job)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Optional
|
||||
|
||||
PERIODS = ("day", "week", "month")
|
||||
|
||||
|
||||
def window_bounds(period: str, now: Optional[datetime] = None) -> tuple[datetime, datetime]:
|
||||
"""Return ``(start, resets_at)`` of the window containing ``now``.
|
||||
|
||||
Args:
|
||||
period: ``day`` (from 00:00), ``week`` (from Monday) or ``month`` (from the 1st).
|
||||
now: Reference instant; defaults to the current time. Naive values are read as UTC.
|
||||
|
||||
Raises:
|
||||
ValueError: If ``period`` is not one of ``PERIODS``.
|
||||
"""
|
||||
if now is None:
|
||||
now = datetime.now(timezone.utc)
|
||||
elif now.tzinfo is None:
|
||||
now = now.replace(tzinfo=timezone.utc)
|
||||
else:
|
||||
now = now.astimezone(timezone.utc)
|
||||
midnight = now.replace(hour=0, minute=0, second=0, microsecond=0)
|
||||
if period == "day":
|
||||
return midnight, midnight + timedelta(days=1)
|
||||
if period == "week":
|
||||
start = midnight - timedelta(days=midnight.weekday())
|
||||
return start, start + timedelta(days=7)
|
||||
if period == "month":
|
||||
start = midnight.replace(day=1)
|
||||
next_month = (start + timedelta(days=32)).replace(day=1)
|
||||
return start, next_month
|
||||
raise ValueError(f"unknown quota period: {period!r}")
|
||||
@@ -0,0 +1,184 @@
|
||||
"""Repository for the ``quota_policies`` table.
|
||||
|
||||
One row per ``(scope, subject, bucket)``: the instance default (no subject), a
|
||||
team's per-member allowance (``teams.id``) or a user's override (auth ``sub``).
|
||||
All methods take a ``Connection`` and do not manage their own transactions.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Optional
|
||||
|
||||
from sqlalchemy import Connection, text
|
||||
|
||||
from docsgpt.storage.db.base_repository import row_to_dict
|
||||
|
||||
SCOPES = ("instance", "team", "user")
|
||||
BUCKETS = ("all", "direct", "agent")
|
||||
|
||||
_COLUMNS = (
|
||||
"id, scope, subject_id, bucket, token_limit, token_unlimited, cost_limit_usd, "
|
||||
"cost_unlimited, enabled, note, created_by, updated_by, created_at, updated_at"
|
||||
)
|
||||
|
||||
|
||||
def _validate(scope: str, subject_id: Optional[str], bucket: str) -> None:
|
||||
if scope not in SCOPES:
|
||||
raise ValueError(f"unknown quota scope: {scope!r}")
|
||||
if bucket not in BUCKETS:
|
||||
raise ValueError(f"unknown quota bucket: {bucket!r}")
|
||||
if (scope == "instance") != (subject_id is None):
|
||||
raise ValueError("subject_id is required for team and user policies and must be omitted for instance")
|
||||
|
||||
|
||||
class QuotaPoliciesRepository:
|
||||
"""Admin-set usage limits."""
|
||||
|
||||
def __init__(self, conn: Connection) -> None:
|
||||
self._conn = conn
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Reads
|
||||
# ------------------------------------------------------------------
|
||||
def policies_for_user(self, user_id: str) -> list[dict]:
|
||||
"""Return every enabled row that applies to ``user_id``.
|
||||
|
||||
That is the instance rows, the user's own rows, and the rows of each
|
||||
team the user belongs to (once per team, whatever roles or sources
|
||||
the membership has).
|
||||
"""
|
||||
result = self._conn.execute(
|
||||
text(
|
||||
f"""
|
||||
SELECT {_COLUMNS} FROM quota_policies
|
||||
WHERE enabled AND (
|
||||
scope = 'instance'
|
||||
OR (scope = 'user' AND subject_id = :user_id)
|
||||
OR (scope = 'team' AND subject_id IN (
|
||||
SELECT DISTINCT team_id::text FROM team_members WHERE user_id = :user_id
|
||||
))
|
||||
)
|
||||
"""
|
||||
),
|
||||
{"user_id": user_id},
|
||||
)
|
||||
return [row_to_dict(row) for row in result.fetchall()]
|
||||
|
||||
def get(self, scope: str, subject_id: Optional[str], bucket: str = "all") -> Optional[dict]:
|
||||
"""Return one policy row, or ``None``."""
|
||||
_validate(scope, subject_id, bucket)
|
||||
row = self._conn.execute(
|
||||
text(
|
||||
f"SELECT {_COLUMNS} FROM quota_policies "
|
||||
"WHERE scope = :scope AND COALESCE(subject_id, '') = :subject AND bucket = :bucket"
|
||||
),
|
||||
{"scope": scope, "subject": subject_id or "", "bucket": bucket},
|
||||
).fetchone()
|
||||
return row_to_dict(row) if row is not None else None
|
||||
|
||||
def list_for_subject(self, scope: str, subject_id: Optional[str]) -> list[dict]:
|
||||
"""Return a subject's rows across buckets, ``all`` first."""
|
||||
_validate(scope, subject_id, "all")
|
||||
result = self._conn.execute(
|
||||
text(
|
||||
f"SELECT {_COLUMNS} FROM quota_policies "
|
||||
"WHERE scope = :scope AND COALESCE(subject_id, '') = :subject "
|
||||
"ORDER BY array_position(ARRAY['all', 'direct', 'agent'], bucket)"
|
||||
),
|
||||
{"scope": scope, "subject": subject_id or ""},
|
||||
)
|
||||
return [row_to_dict(row) for row in result.fetchall()]
|
||||
|
||||
def list_by_scope(self, scope: str) -> list[dict]:
|
||||
"""Return every row of a scope, ordered by subject then bucket."""
|
||||
if scope not in SCOPES:
|
||||
raise ValueError(f"unknown quota scope: {scope!r}")
|
||||
result = self._conn.execute(
|
||||
text(
|
||||
f"SELECT {_COLUMNS} FROM quota_policies WHERE scope = :scope "
|
||||
"ORDER BY subject_id NULLS FIRST, array_position(ARRAY['all', 'direct', 'agent'], bucket)"
|
||||
),
|
||||
{"scope": scope},
|
||||
)
|
||||
return [row_to_dict(row) for row in result.fetchall()]
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Writes
|
||||
# ------------------------------------------------------------------
|
||||
def upsert(
|
||||
self,
|
||||
*,
|
||||
scope: str,
|
||||
subject_id: Optional[str],
|
||||
bucket: str = "all",
|
||||
token_limit: Optional[int] = None,
|
||||
token_unlimited: bool = False,
|
||||
cost_limit_usd: Optional[float] = None,
|
||||
cost_unlimited: bool = False,
|
||||
enabled: bool = True,
|
||||
note: Optional[str] = None,
|
||||
actor: Optional[str] = None,
|
||||
) -> dict:
|
||||
"""Create or replace the policy for ``(scope, subject_id, bucket)``.
|
||||
|
||||
Raises:
|
||||
ValueError: On an unknown scope or bucket, a subject that does not
|
||||
match the scope, a negative limit, or a budget that is both
|
||||
limited and unlimited.
|
||||
"""
|
||||
_validate(scope, subject_id, bucket)
|
||||
if token_unlimited and token_limit is not None:
|
||||
raise ValueError("token budget cannot be both limited and unlimited")
|
||||
if cost_unlimited and cost_limit_usd is not None:
|
||||
raise ValueError("cost budget cannot be both limited and unlimited")
|
||||
if (token_limit is not None and token_limit < 0) or (cost_limit_usd is not None and cost_limit_usd < 0):
|
||||
raise ValueError("limits must not be negative")
|
||||
row = self._conn.execute(
|
||||
text(
|
||||
f"""
|
||||
INSERT INTO quota_policies (
|
||||
scope, subject_id, bucket, token_limit, token_unlimited,
|
||||
cost_limit_usd, cost_unlimited, enabled, note, created_by, updated_by
|
||||
)
|
||||
VALUES (
|
||||
:scope, :subject_id, :bucket, :token_limit, :token_unlimited,
|
||||
:cost_limit_usd, :cost_unlimited, :enabled, :note, :actor, :actor
|
||||
)
|
||||
ON CONFLICT (scope, COALESCE(subject_id, ''), bucket) DO UPDATE SET
|
||||
token_limit = EXCLUDED.token_limit,
|
||||
token_unlimited = EXCLUDED.token_unlimited,
|
||||
cost_limit_usd = EXCLUDED.cost_limit_usd,
|
||||
cost_unlimited = EXCLUDED.cost_unlimited,
|
||||
enabled = EXCLUDED.enabled,
|
||||
note = EXCLUDED.note,
|
||||
updated_by = EXCLUDED.updated_by
|
||||
RETURNING {_COLUMNS}
|
||||
"""
|
||||
),
|
||||
{
|
||||
"scope": scope,
|
||||
"subject_id": subject_id,
|
||||
"bucket": bucket,
|
||||
"token_limit": token_limit,
|
||||
"token_unlimited": token_unlimited,
|
||||
"cost_limit_usd": cost_limit_usd,
|
||||
"cost_unlimited": cost_unlimited,
|
||||
"enabled": enabled,
|
||||
"note": note,
|
||||
"actor": actor,
|
||||
},
|
||||
).one()
|
||||
return row_to_dict(row)
|
||||
|
||||
def delete(self, scope: str, subject_id: Optional[str], bucket: Optional[str] = None) -> int:
|
||||
"""Delete a subject's policy for ``bucket``, or all of them; return the count."""
|
||||
_validate(scope, subject_id, bucket or "all")
|
||||
clauses = ["scope = :scope", "COALESCE(subject_id, '') = :subject"]
|
||||
params = {"scope": scope, "subject": subject_id or ""}
|
||||
if bucket is not None:
|
||||
clauses.append("bucket = :bucket")
|
||||
params["bucket"] = bucket
|
||||
result = self._conn.execute(
|
||||
text(f"DELETE FROM quota_policies WHERE {' AND '.join(clauses)}"), params
|
||||
)
|
||||
return result.rowcount
|
||||
Reference in new issue
Block a user