Support for SQLite quota tables

Author: tisnik

src/runners/quota_scheduler.py

@@ -33,21 +33,38 @@
     """
 
 
-INCREASE_QUOTA_STATEMENT = """
+INCREASE_QUOTA_STATEMENT_PG = """
     UPDATE quota_limits
        SET available=available+%s, revoked_at=NOW()
      WHERE subject=%s
        AND revoked_at < NOW() - INTERVAL %s ;
     """
 
-RESET_QUOTA_STATEMENT = """
+
+INCREASE_QUOTA_STATEMENT_SQLITE = """
+    UPDATE quota_limits
+       SET available=available+?, revoked_at=datetime('now')
+     WHERE subject=?
+       AND revoked_at < datetime('now', ?);
+    """
+
+
+RESET_QUOTA_STATEMENT_PG = """
     UPDATE quota_limits
        SET available=%s, revoked_at=NOW()
      WHERE subject=%s
        AND revoked_at < NOW() - INTERVAL %s ;
     """
 
 
+RESET_QUOTA_STATEMENT_SQLITE = """
+    UPDATE quota_limits
+       SET available=?, revoked_at=datetime('now')
+     WHERE subject=?
+       AND revoked_at < datetime('now', ?);
+    """
+
+
 def quota_scheduler(config: QuotaHandlersConfiguration) -> bool:
     """Quota scheduler task."""
     if config is None:
@@ -70,6 +87,9 @@ def quota_scheduler(config: QuotaHandlersConfiguration) -> bool:
     init_tables(connection)
     period = config.scheduler.period
 
+    increase_quota_statement = get_increase_quota_statement(config)
+    reset_quota_statement = get_reset_quota_statement(config)
+
     logger.info(
         "Quota scheduler started in separated thread with period set to %d seconds",
         period,
@@ -79,8 +99,10 @@ def quota_scheduler(config: QuotaHandlersConfiguration) -> bool:
         logger.info("Quota scheduler sync started")
         for limiter in config.limiters:
             try:
-                quota_revocation(connection, limiter)
-            except Exception as e:  # pylint: disable=broad-exception-caught)
+                quota_revocation(
+                    connection, limiter, increase_quota_statement, reset_quota_statement
+                )
+            except Exception as e:  # pylint: disable=broad-exception-caught
                 logger.error("Quota revoke error: %s", e)
         logger.info("Quota scheduler sync finished")
         sleep(period)
@@ -89,7 +111,26 @@ def quota_scheduler(config: QuotaHandlersConfiguration) -> bool:
     return True
 
 
-def quota_revocation(connection: Any, quota_limiter: QuotaLimiterConfiguration) -> None:
+def get_increase_quota_statement(config: QuotaHandlersConfiguration) -> str:
+    """Get the SQL statement to increase quota."""
+    if config.sqlite is not None:
+        return INCREASE_QUOTA_STATEMENT_SQLITE
+    return INCREASE_QUOTA_STATEMENT_PG
+
+
+def get_reset_quota_statement(config: QuotaHandlersConfiguration) -> str:
+    """Get the SQL statement to reset quota."""
+    if config.sqlite is not None:
+        return RESET_QUOTA_STATEMENT_SQLITE
+    return RESET_QUOTA_STATEMENT_PG
+
+
+def quota_revocation(
+    connection: Any,
+    quota_limiter: QuotaLimiterConfiguration,
+    increase_quota_statement: str,
+    reset_quota_statement: str,
+) -> None:
     """Quota revocation mechanism."""
     logger.info(
         "Quota revocation mechanism for limiter '%s' of type '%s'",
@@ -107,17 +148,29 @@ def quota_revocation(connection: Any, quota_limiter: QuotaLimiterConfiguration)
 
     if quota_limiter.quota_increase is not None:
         increase_quota(
-            connection, subject_id, quota_limiter.quota_increase, quota_limiter.period
+            connection,
+            increase_quota_statement,
+            subject_id,
+            quota_limiter.quota_increase,
+            quota_limiter.period,
         )
 
     if quota_limiter.initial_quota is not None and quota_limiter.initial_quota > 0:
         reset_quota(
-            connection, subject_id, quota_limiter.initial_quota, quota_limiter.period
+            connection,
+            reset_quota_statement,
+            subject_id,
+            quota_limiter.initial_quota,
+            quota_limiter.period,
         )
 
 
 def increase_quota(
-    connection: Any, subject_id: str, increase_by: int, period: str
+    connection: Any,
+    update_statement: str,
+    subject_id: str,
+    increase_by: int,
+    period: str,
 ) -> None:
     """Increase quota by specified amount."""
     logger.info(
@@ -127,41 +180,51 @@ def increase_quota(
         period,
     )
 
-    update_statement = INCREASE_QUOTA_STATEMENT
-
-    with connection.cursor() as cursor:
-        cursor.execute(
-            update_statement,
-            (
-                increase_by,
-                subject_id,
-                period,
-            ),
-        )
-        logger.info("Changed %d rows in database", cursor.rowcount)
+    # for compatibility with SQLite it is not possible to use context manager
+    # there
+    cursor = connection.cursor()
+    cursor.execute(
+        update_statement,
+        (
+            increase_by,
+            subject_id,
+            period,
+        ),
+    )
+    cursor.close()
+    connection.commit()
+    logger.info("Changed %d rows in database", cursor.rowcount)
 
 
-def reset_quota(connection: Any, subject_id: str, reset_to: int, period: str) -> None:
+def reset_quota(
+    connection: Any,
+    update_statement: str,
+    subject_id: str,
+    reset_to: int,
+    period: str,
+) -> None:
     """Reset quota to specified amount."""
     logger.info(
-        "Reseting quota for subject '%s' to %d when period %s is reached",
+        "Resetting quota for subject '%s' to %d when period %s is reached",
         subject_id,
         reset_to,
         period,
     )
 
-    update_statement = RESET_QUOTA_STATEMENT
-
-    with connection.cursor() as cursor:
-        cursor.execute(
-            update_statement,
-            (
-                reset_to,
-                subject_id,
-                period,
-            ),
-        )
-        logger.info("Changed %d rows in database", cursor.rowcount)
+    # for compatibility with SQLite it is not possible to use context manager
+    # there
+    cursor = connection.cursor()
+    cursor.execute(
+        update_statement,
+        (
+            reset_to,
+            subject_id,
+            period,
+        ),
+    )
+    cursor.close()
+    connection.commit()
+    logger.info("Changed %d rows in database", cursor.rowcount)
 
 
 def get_subject_id(limiter_type: str) -> str: