Python 3.13 support: wrapping update

.github/workflows/benchmark.yml

@@ -1,7 +1,8 @@
-name: Benchmark
+name: 📊 Benchmarks
 on:
   push: {}
   workflow_call: {}
+
 jobs:
   benchmark_sql_algorithm:
     runs-on: ubuntu-latest
@@ -30,15 +31,24 @@ jobs:
         run: |
           poetry run python ./benchmarks/sql_benchmark/sql_benchmark_fw.py
           poetry run python ./benchmarks/sql_benchmark/sql_benchmark_no_fw.py
+
   benchmark_with_flask_mysql:
     runs-on: ubuntu-latest
     timeout-minutes: 10
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
     steps:
       - name: Checkout code
         uses: actions/checkout@v4
       - name: Start databases
         working-directory:  ./sample-apps/databases
         run: docker compose up --build -d
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v2
+        with:
+          python-version: ${{ matrix.python-version }}
       - name: Install dependencies & build
         run: |
           python -m pip install --upgrade pip
@@ -51,15 +61,24 @@ jobs:
       - name: Run flask-mysql k6 Benchmark
         run: |
           k6 run -q ./benchmarks/flask-mysql-benchmarks.js
+
   benchmark_with_starlette_app:
     runs-on: ubuntu-latest
     timeout-minutes: 10
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
     steps:
       - name: Checkout code
         uses: actions/checkout@v4
       - name: Start databases
         working-directory:  ./sample-apps/databases
         run: docker compose up --build -d
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v2
+        with:
+          python-version: ${{ matrix.python-version }}
       - name: Install dependencies & build
         run: |
           python -m pip install --upgrade pip
@@ -71,9 +90,5 @@ jobs:
         run: |
           sudo apt-get update
           sudo apt-get install -y wrk
-      - name: Set up Python 3.9
-        uses: actions/setup-python@v2
-        with:
-          python-version: 3.9
       - name: Run benchmark
         run: python ./benchmarks/starlette_benchmark.py

.github/workflows/end2end.yml

@@ -35,7 +35,7 @@ jobs:
           - { name: flask-postgres-xml, testfile: end2end/flask_postgres_xml_lxml_test.py }
           - { name: quart-postgres-uvicorn, testfile: end2end/quart_postgres_uvicorn_test.py }
           - { name: starlette-postgres-uvicorn, testfile: end2end/starlette_postgres_uvicorn_test.py }
-        python-version: ["3.10", "3.11", "3.12"]
+        python-version: ["3.10", "3.11", "3.12", "3.13"]
     steps:
       - name: Install packages
         run: sudo apt update && sudo apt install python3-dev libmysqlclient-dev

.github/workflows/unit-test.yml

@@ -10,7 +10,7 @@ jobs:
       # Don't cancel jobs if one fails
       fail-fast: false
       matrix:
-        python-version: ["3.8", "3.9", "3.10", "3.11", "3.12"]
+        python-version: ["3.8", "3.9", "3.10", "3.11", "3.12", "3.13"]
     steps:
     - name: Checkout code
       uses: actions/checkout@v4

aikido_zen/background_process/commands/__init__.py

@@ -3,7 +3,6 @@
 from aikido_zen.helpers.logging import logger
 from .attack import process_attack
 from .read_property import process_read_property
-from .initialize_route import process_initialize_route
 from .user import process_user
 from .should_ratelimit import process_should_ratelimit
 from .kill import process_kill
@@ -16,7 +15,6 @@
     # This maps to a tuple : (function, returns_data?)
     # Commands that don't return data :
     "ATTACK": (process_attack, False),
-    "INITIALIZE_ROUTE": (process_initialize_route, False),
     "USER": (process_user, False),
     "KILL": (process_kill, False),
     "STATISTICS": (process_statistics, False),

aikido_zen/background_process/routes/__init__.py

@@ -19,20 +19,17 @@ def __init__(self, max_size=1000):
 
     def initialize_route(self, route_metadata):
         """
-        Initializes a route for the first time.
+        Initializes a route for the first time. `hits_delta_since_sync` counts delta between syncs.
         """
         self.manage_routes_size()
         key = route_to_key(route_metadata)
-        if self.routes.get(key):
-            return
         self.routes[key] = {
             "method": route_metadata.get("method"),
             "path": route_metadata.get("route"),
             "hits": 0,
+            "hits_delta_since_sync": 0,
             "apispec": {},
         }
-        # This field counts the difference in hits in between synchronisation for threads :
-        self.routes[key]["hits_delta_since_sync"] = 0
 
     def increment_route(self, route_metadata):
         """
@@ -41,8 +38,8 @@ def increment_route(self, route_metadata):
         """
         key = route_to_key(route_metadata)
         if not self.routes.get(key):
-            return
-        # Add a hit to the route :
+            self.initialize_route(route_metadata)
+        # Add hits to
         route = self.routes.get(key)
         route["hits"] += 1
         route["hits_delta_since_sync"] += 1

aikido_zen/background_process/routes/init_test.py

@@ -102,13 +102,6 @@ def test_increment_route_twice():
     assert routes.routes["GET:/api/resource"]["hits"] == 2
 
 
-def test_increment_route_that_does_not_exist():
-    routes = Routes(max_size=3)
-    routes.increment_route(gen_route_metadata(route="/api/resource"))
-    routes.increment_route(gen_route_metadata(route="/api/resource"))
-    assert len(routes.routes) == 0
-
-
 def test_clear_routes():
     routes = Routes(max_size=3)
     routes.initialize_route(gen_route_metadata(route="/api/resource"))

aikido_zen/middleware/init_test.py

@@ -2,16 +2,18 @@
 
 import pytest
 from aikido_zen.context import current_context, Context, get_current_context
-from aikido_zen.thread.thread_cache import ThreadCache, threadlocal_storage
+from aikido_zen.thread.thread_cache import ThreadCache, get_cache
 from . import should_block_request
 
 
 @pytest.fixture(autouse=True)
 def run_around_tests():
+    get_cache().reset()
     yield
     # Make sure to reset context and cache after every test so it does not
     # interfere with other tests
     current_context.set(None)
+    get_cache().reset()
 
 
 def test_without_context():
@@ -39,20 +41,15 @@ def set_context(user=None, executed_middleware=False):
     ).set_as_current_context()
 
 
-class MyThreadCache(ThreadCache):
-    def renew_if_ttl_expired(self):
-        return
-
-
 def test_with_context_without_cache():
     set_context()
-    threadlocal_storage.cache = None
+    get_cache().cache = None
     assert should_block_request() == {"block": False}
 
 
 def test_with_context_with_cache():
     set_context(user={"id": "123"})
-    thread_cache = MyThreadCache()
+    thread_cache = get_cache()
 
     thread_cache.config.blocked_uids = ["123"]
     assert get_current_context().executed_middleware == False
@@ -76,7 +73,7 @@ def test_with_context_with_cache():
 
 def test_cache_comms_with_endpoints():
     set_context(user={"id": "456"})
-    thread_cache = MyThreadCache()
+    thread_cache = get_cache()
     thread_cache.config.blocked_uids = ["123"]
     thread_cache.config.endpoints = [
         {

aikido_zen/sinks/__init__.py

@@ -0,0 +1,70 @@
+from wrapt import wrap_object, FunctionWrapper, when_imported
+from aikido_zen.background_process.packages import ANY_VERSION, is_package_compatible
+from aikido_zen.errors import AikidoException
+from aikido_zen.helpers.logging import logger
+
+
+def on_import(name, package="", version_requirement=ANY_VERSION):
+    """
+    Decorator to register a function to be called when a package is imported.
+    It checks if the package is compatible with the specified version requirement.
+    """
+
+    def decorator(func):
+        if package and not is_package_compatible(package, version_requirement):
+            return
+
+        when_imported(name)(func)  # Register the function to be called on import
+
+    return decorator
+
+
+def patch_function(module, name, wrapper):
+    """
+    Patches a function in the specified module with a wrapper function.
+    """
+    try:
+        wrap_object(module, name, FunctionWrapper, (wrapper,))
+    except Exception as e:
+        logger.info("Failed to wrap %s:%s, due to: %s", module, name, e)
+
+
+def before(wrapper):
+    """
+    Surrounds a patch with try-except and calls the original function at the end
+    """
+
+    def decorator(func, instance, args, kwargs):
+        try:
+            wrapper(func, instance, args, kwargs)  # Call the patch
+        except AikidoException as e:
+            raise e  # Re-raise AikidoException
+        except Exception as e:
+            logger.debug(
+                "%s:%s wrapping-before error: %s", func.__module__, func.__name__, e
+            )
+
+        return func(*args, **kwargs)  # Call the original function
+
+    return decorator
+
+
+def after(wrapper):
+    """
+    Surrounds a patch with try-except, calls the original function and gives the return value to the patch
+    """
+
+    def decorator(func, instance, args, kwargs):
+        return_value = func(*args, **kwargs)  # Call the original function
+        try:
+            wrapper(func, instance, args, kwargs, return_value)  # Call the patch
+        except AikidoException as e:
+            raise e  # Re-raise AikidoException
+        except Exception as e:
+            logger.debug(
+                "%s:%s wrapping-after error: %s", func.__module__, func.__name__, e
+            )
+        finally:
+            return return_value
+
+    return decorator