feat: add caching

Author: MatthiasHarzer

.gitignore

@@ -3,3 +3,4 @@ __pycache__
 .venv
 .vscode
 .idea
+*.db

requirements.txt

@@ -1,2 +1,5 @@
 fastapi~=0.115.0
 requests~=2.32.3
+mysql-connector~=2.2.9
+uvicorn~=0.21.1
+pydantic~=1.10.7

server.py

@@ -0,0 +1 @@
+from .cache_handler import get_cache_provider

server/__init__.py

@@ -0,0 +1,50 @@
+
+import os
+
+import requests
+from server.caching.cache_provider import CacheProvider
+from server.caching.cache_provider.in_memory_cache_provider import InMemoryCacheProvider
+from server.caching.cache_provider.mysql_cache_provider import MySQLCacheProvider
+from server.caching.cache_provider.no_cache_provider import NoCacheProvider
+from server.caching.cache_provider.sqlite_cache_provider import SQLiteCacheProvider
+from server.caching.cache_request import CacheRequest
+
+CACHE_DB_FILE = "./cache.db"
+
+def get_cache_provider() -> CacheProvider:
+    cache_type = os.environ.get("CACHE_MODE", "memory")
+
+    print("Using cache provider: " + cache_type)
+
+    match cache_type:
+        case "sqlite":
+            db_file = os.environ.get("SQLITE_FILE", CACHE_DB_FILE)
+            return SQLiteCacheProvider(db_file)
+        case "memory":
+            return InMemoryCacheProvider()
+        case "mysql":
+            host = os.environ.get("MYSQL_HOST")
+            user = os.environ.get("MYSQL_USER")
+            password = os.environ.get("MYSQL_PASSWORD")
+            database = os.environ.get("MYSQL_DATABASE")
+            port = os.environ.get("MYSQL_PORT", 3306)
+
+            if None in (host, user, password, database):
+                print("WARNING: Missing environment variables for MySQL cache provider. "
+                      "Required: MYSQL_HOST, MYSQL_USER, MYSQL_PASSWORD, MYSQL_DATABASE"
+                      "\nFalling back to NoCacheProvider.")
+                return NoCacheProvider()
+
+            return MySQLCacheProvider(
+                host=host,  # type: ignore
+                user=user,  # type: ignore
+                password=password,  # type: ignore
+                database=database,  # type: ignore
+                port=port # type: ignore
+            )
+        case "none":
+            return NoCacheProvider()
+
+    print(f"WARNING: No cache provider found for type {cache_type}. Using NoCacheProvider.")
+
+    return NoCacheProvider()

server/caching/__init__.py

@@ -0,0 +1 @@
+from .cache_provider import CacheProvider

server/caching/cache_handler.py

@@ -0,0 +1,48 @@
+from abc import ABC, abstractmethod
+from datetime import datetime
+
+from server.caching.cache_request import CacheRequest
+
+type TimedCache = tuple[bytes, datetime]
+
+
+class CacheProvider(ABC):
+    """
+    A cache provider that can be used to cache requests
+    """
+
+    @abstractmethod
+    def _get(self, request: CacheRequest) -> TimedCache | None:
+        """
+        Get the response and timestamp from the cache
+        :param request: The request to get the response for
+        :return:  A tuple containing the response and timestamp
+        """
+
+    def get(self, request: CacheRequest) -> bytes | None:
+        """
+        Get the response from the cache if it exists and is not stale
+        :param request: The request to get the response for
+        :return: The response if it exists and is not stale, None otherwise
+        """
+        cached = self._get(request)
+        if not cached:
+            return None
+
+        response, timestamp = cached
+
+        time_diff = datetime.now() - timestamp
+
+        if request.max_age <= 0 or 0 < time_diff.seconds < request.max_age:
+            return response
+
+        return None
+
+    @abstractmethod
+    def set(self, request: CacheRequest, response: bytes) -> None:
+        """
+        Save the response to the cache
+        :param request: The request to save the response for
+        :param response: The response to save
+        :return:
+        """

server/caching/cache_provider/__init__.py

@@ -0,0 +1,28 @@
+import json
+from datetime import datetime
+
+from server.caching.cache_provider import CacheProvider
+from server.caching.cache_provider.cache_provider import TimedCache
+from server.caching.cache_request import CacheRequest
+
+
+def _hash_request(request: CacheRequest) -> int:
+    """
+    Hash a request
+    :param request:
+    :return:
+    """
+    return hash((request.method, request.url, json.dumps(request.body), json.dumps(request.headers)))
+
+
+class InMemoryCacheProvider(CacheProvider):
+    def __init__(self):
+        self.cache: dict[int, TimedCache] = {}
+
+    def _get(self, request: CacheRequest) -> TimedCache | None:
+        hashed = _hash_request(request)
+        return self.cache.get(hashed)
+
+    def set(self, request: CacheRequest, response: bytes) -> None:
+        hashed = _hash_request(request)
+        self.cache[hashed] = (response, datetime.now())

server/caching/cache_provider/cache_provider.py

@@ -0,0 +1,106 @@
+import json
+from datetime import datetime
+
+from server.caching.cache_provider import CacheProvider
+from server.caching.cache_provider.cache_provider import TimedCache
+from server.caching.cache_request import CacheRequest
+import mysql.connector
+import mysql.connector.cursor
+
+_TABLE_SQL = """
+        CREATE TABLE IF NOT EXISTS cache (
+    id INTEGER PRIMARY KEY AUTO_INCREMENT,
+    method TEXT NOT NULL,
+    url TEXT NOT NULL,
+    response MEDIUMBLOB NOT NULL,
+    body TEXT NULL,
+    headers TEXT NULL,
+    timestamp DATETIME DEFAULT CURRENT_TIMESTAMP
+    );"""
+
+
+def _get_current_timestamp() -> str:
+    """
+    Get the current timestamp as a string
+    :return:
+    """
+    return datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+
+
+class MySQLCacheProvider(CacheProvider):
+    def __init__(self, host: str, user: str, password: str, database: str, port: int = 3306):
+        self.host = host
+        self.user = user
+        self.password = password
+        self.database = database
+        self.port = port
+        self.initialized = False
+        try:
+            self._create_table()
+            self.initialized = True
+        except mysql.connector.errors.Error as e:
+            print("Error while initializing cache: " + str(e))
+
+    def _connect(self) -> mysql.connector.MySQLConnection:
+        return mysql.connector.connect(
+            host=self.host,
+            user=self.user,
+            password=self.password,
+            database=self.database,
+            port=self.port
+        )  # type: ignore
+
+    def _create_table(self):
+        db = self._connect()
+        try:
+            c = db.cursor()
+            c.execute(_TABLE_SQL)
+            db.commit()
+        except mysql.connector.errors.Error as e:
+            print(e)
+
+    def _get(self, request: CacheRequest) -> TimedCache | None:
+        """
+        Get the response and timestamp from the cache
+        :param request: The request to get the response for
+        :return:  A tuple containing the response and timestamp
+        """
+
+        headers = json.dumps(request.headers) if request.headers else ""
+
+        db = self._connect()
+        c = db.cursor()
+        c.execute(
+            "SELECT response, timestamp FROM cache WHERE method = %s AND url = %s AND body <=> %s AND headers <=> %s",
+            (request.method, request.url, request.body, headers))
+
+        result = c.fetchone()
+        if result:
+            return result[0], result[1]
+        return None
+
+    def set(self, request: CacheRequest, response: bytes) -> None:
+        """
+        Save the response to the cache
+        :param request: The request to save the response for
+        :param response: The response to save
+        :return:
+        """
+        headers = json.dumps(request.headers) if request.headers else ""
+
+        db = self._connect()
+        c = db.cursor()
+        c.execute(
+            "SELECT id FROM cache WHERE method = %s AND url = %s AND body <=> %s AND headers <=> %s",
+            (request.method, request.url, request.body, headers))
+        existing = c.fetchone()
+
+        if existing:
+            c.execute("UPDATE cache SET response = %s, timestamp = %s WHERE id = %s",
+                      (response, _get_current_timestamp(), existing[0]))
+        else:
+            c.execute(
+                "INSERT INTO cache (method, url, response, body, headers, timestamp) VALUES (%s, %s, %s, %s, %s, %s, %s)",
+                (request.method, request.url, response, request.body, headers, _get_current_timestamp()))
+
+        db.commit()

server/caching/cache_provider/in_memory_cache_provider.py

@@ -0,0 +1,11 @@
+from server.caching.cache_provider import CacheProvider
+from server.caching.cache_provider.cache_provider import TimedCache
+from server.caching.cache_request import CacheRequest
+
+
+class NoCacheProvider(CacheProvider):
+    def _get(self, request: CacheRequest) -> TimedCache | None:
+        return None
+
+    def set(self, request: CacheRequest, response: bytes) -> None:
+        pass