Revoke quota ability

road-core · Feb 19, 2025 · d9ac144 · d9ac144
1 parent a231067
commit d9ac144
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 0 deletions.
diff --git a/ols/src/quota/quota_limiter.py b/ols/src/quota/quota_limiter.py
@@ -10,6 +10,10 @@ class QuotaLimiter(ABC):
     def available_quota(self, user_id: str) -> int:
         """Retrieve available quota for given user."""
 
+    @abstractmethod
+    def revoke_quota(self, user_id: str) -> None:
+        """Revoke quota for given user."""
+
     @abstractmethod
     def consume_tokens(
         self, user_id: str, input_tokens: int, output_tokens: int

diff --git a/ols/src/quota/user_quota_limiter.py b/ols/src/quota/user_quota_limiter.py
@@ -79,6 +79,18 @@ def available_quota(self, user_id: str) -> int:
                 return self.initial_quota
             return value[0]
 
+    def revoke_quota(self, user_id: str) -> None:
+        """Revoke quota for given user."""
+        # timestamp to be used
+        updated_at = datetime.now()
+
+        with self.connection.cursor() as cursor:
+            cursor.execute(
+                UserQuotaLimiter.UPDATE_AVAILABLE_QUOTA_FOR_USER,
+                (self.initial_quota, updated_at, user_id),
+            )
+            self.connection.commit()
+
     def consume_tokens(
         self, user_id: str, input_tokens: int, output_tokens: int
     ) -> None: