Sync routes improved, allowing new fields to be added

aikido_zen/background_process/cloud_connection_manager/__init__.py

@@ -1,7 +1,7 @@
 """ This file simply exports the CloudConnectionManager class"""
 
 from aikido_zen.background_process.heartbeats import send_heartbeats_every_x_secs
-from aikido_zen.background_process.routes import Routes
+from aikido_zen.storage.routes import Routes
 from aikido_zen.ratelimiting.rate_limiter import RateLimiter
 from aikido_zen.helpers.logging import logger
 from .update_firewall_lists import update_firewall_lists
@@ -32,7 +32,7 @@ def __init__(self, block, api, token, serverless):
         self.block = block
         self.api: ReportingApiHTTP = api
         self.token = token  # Should be instance of the Token class!
-        self.routes = Routes(200)
+        self.routes = Routes(max_size=1000)
         self.hostnames = Hostnames(200)
         self.conf = ServiceConfig(
             endpoints=[],
@@ -73,7 +73,7 @@ def report_initial_stats(self):
         """
         data_present = (
             not self.statistics.empty()
-            or len(self.routes.routes) > 0
+            or not self.routes.empty()
             or not self.ai_stats.empty()
         )
         should_report_initial_stats = data_present and not self.conf.received_any_stats

aikido_zen/background_process/cloud_connection_manager/init_test.py

@@ -25,7 +25,7 @@ def test_cloud_connection_manager_initialization(setup_cloud_connection_manager)
     assert manager.block is None
     assert isinstance(manager.api, ReportingApiHTTP)
     assert isinstance(manager.token, Token)
-    assert len(manager.routes) == 0
+    assert len(manager.routes.export()) == 0
     assert isinstance(manager.hostnames, Hostnames)
     assert isinstance(manager.conf, ServiceConfig)
     assert isinstance(manager.rate_limiter, RateLimiter)

aikido_zen/background_process/cloud_connection_manager/send_heartbeat.py

@@ -14,7 +14,7 @@
     logger.debug("Aikido CloudConnectionManager : Sending out heartbeat")
     stats = connection_manager.statistics.get_record()
     users = connection_manager.users.as_array()
-    routes = list(connection_manager.routes)
+    routes = list(connection_manager.routes.export().values())
     outgoing_domains = connection_manager.hostnames.as_array()
     ai_stats = connection_manager.ai_stats.get_stats()
     packages = PackagesStore.export()

aikido_zen/background_process/commands/sync_data.py

@@ -13,25 +13,13 @@ def process_sync_data(connection_manager, data, conn, queue=None):
     BG Process -> Thread : Routes and config
     """
 
-    # Sync routes
-    routes = connection_manager.routes
-    for route in data.get("current_routes", {}).values():
-        route_metadata = {"method": route["method"], "route": route["path"]}
-        if not routes.get(route_metadata):
-            routes.initialize_route(route_metadata)
-        existing_route = routes.get(route_metadata)
-
-        # Update hit count :
-        hits_delta_since_sync = int(route.get("hits_delta_since_sync", 0))
-        existing_route["hits"] += hits_delta_since_sync
-
-        # Update API Spec :
-        update_route_info(route["apispec"], existing_route)
-
     # Save middleware installed :
     if data.get("middleware_installed", False):
         connection_manager.middleware_installed = True
 
+    # Sync routes
+    connection_manager.routes.import_from_record(data.get("routes", {}))
+
     # Sync hostnames
     for hostnames_entry in data.get("hostnames", list()):
         connection_manager.hostnames.add(
@@ -56,7 +44,7 @@ def process_sync_data(connection_manager, data, conn, queue=None):
     if connection_manager.conf.last_updated_at > 0:
         # Only report data if the config has been fetched.
         return {
-            "routes": dict(connection_manager.routes.routes),
+            "routes": connection_manager.routes.export(include_apispecs=False),
             "config": connection_manager.conf,
         }
     return {}

aikido_zen/background_process/commands/sync_data_test.py

@@ -1,9 +1,7 @@
-from multiprocessing.forkserver import connect_to_new_process
-
 import pytest
 from unittest.mock import MagicMock
 from .sync_data import process_sync_data
-from aikido_zen.background_process.routes import Routes
+from aikido_zen.storage.routes import Routes
 from aikido_zen.helpers.iplist import IPList
 from ..packages import PackagesStore
 from ...storage.hostnames import Hostnames
@@ -35,17 +33,18 @@ def test_process_sync_data_initialization(setup_connection_manager):
     test_hostnames.add("bumblebee.com", 8080)
 
     data = {
-        "current_routes": {
-            "route1": {
+        "routes": {
+            "GET:/api/v1/resource": {
                 "method": "GET",
                 "path": "/api/v1/resource",
-                "hits_delta_since_sync": 5,
+                "hits": 5,
+                "hits_delta_since_sync": 2000,
                 "apispec": {"info": "API spec for resource"},
             },
-            "route2": {
+            "POST:/api/v1/resource": {
                 "method": "POST",
                 "path": "/api/v1/resource",
-                "hits_delta_since_sync": 3,
+                "hits": 3,
                 "apispec": {"info": "API spec for resource"},
             },
         },
@@ -75,19 +74,9 @@ def test_process_sync_data_initialization(setup_connection_manager):
     result = process_sync_data(connection_manager, data, None)
 
     # Check that routes were initialized correctly
-    assert len(connection_manager.routes) == 2
-    assert (
-        connection_manager.routes.get({"method": "GET", "route": "/api/v1/resource"})[
-            "hits"
-        ]
-        == 5
-    )
-    assert (
-        connection_manager.routes.get({"method": "POST", "route": "/api/v1/resource"})[
-            "hits"
-        ]
-        == 3
-    )
+    assert len(connection_manager.routes.export()) == 2
+    assert connection_manager.routes.get("GET", "/api/v1/resource")["hits"] == 5
+    assert connection_manager.routes.get("POST", "/api/v1/resource")["hits"] == 3
 
     # Check that the total requests were updated
     assert connection_manager.statistics.get_record()["requests"] == {
@@ -116,18 +105,18 @@ def test_process_sync_data_with_last_updated_at_below_zero(setup_connection_mana
     connection_manager = setup_connection_manager
     connection_manager.conf.last_updated_at = -1
     data = {
-        "current_routes": {
-            "route1": {
+        "routes": {
+            "GET:/api/v1/resource": {
                 "method": "GET",
                 "path": "/api/v1/resource",
-                "hits_delta_since_sync": 5,
+                "hits": 5,
                 "apispec": {"info": "API spec for resource"},
             },
-            "route2": {
+            "POST:/api/v1/resource": {
                 "method": "POST",
                 "path": "/api/v1/resource",
-                "hits_delta_since_sync": 3,
-                "apispec": {"info": "API spec for resource"},
+                "hits": 3,
+                "apispec": {"info": "API spec for another resource"},
             },
         },
         "stats": {
@@ -148,19 +137,19 @@ def test_process_sync_data_with_last_updated_at_below_zero(setup_connection_mana
     result = process_sync_data(connection_manager, data, None)
 
     # Check that routes were initialized correctly
-    assert len(connection_manager.routes) == 2
-    assert (
-        connection_manager.routes.get({"method": "GET", "route": "/api/v1/resource"})[
-            "hits"
-        ]
-        == 5
-    )
-    assert (
-        connection_manager.routes.get({"method": "POST", "route": "/api/v1/resource"})[
-            "hits"
-        ]
-        == 3
-    )
+    assert len(connection_manager.routes.export()) == 2
+    assert connection_manager.routes.get("GET", "/api/v1/resource") == {
+        "hits": 5,
+        "method": "GET",
+        "path": "/api/v1/resource",
+        "apispec": {"info": "API spec for resource"},
+    }
+    assert connection_manager.routes.get("POST", "/api/v1/resource") == {
+        "hits": 3,
+        "method": "POST",
+        "path": "/api/v1/resource",
+        "apispec": {"info": "API spec for another resource"},
+    }
 
     # Check that the total requests were updated
     assert connection_manager.statistics.get_record()["requests"] == {
@@ -186,11 +175,11 @@ def test_process_sync_data_existing_route_and_hostnames(setup_connection_manager
     hostnames_sync.add("c.com", 443)
 
     data = {
-        "current_routes": {
+        "routes": {
             "route1": {
                 "method": "GET",
                 "path": "/api/v1/resource",
-                "hits_delta_since_sync": 5,
+                "hits": 5,
                 "apispec": {"info": "API spec for resource"},
             }
         },
@@ -214,11 +203,11 @@ def test_process_sync_data_existing_route_and_hostnames(setup_connection_manager
 
     # Second call to update the existing route
     data_update = {
-        "current_routes": {
+        "routes": {
             "route1": {
                 "method": "GET",
                 "path": "/api/v1/resource",
-                "hits_delta_since_sync": 10,
+                "hits": 10,
                 "apispec": {"info": "Updated API spec for resource"},
             }
         },
@@ -239,12 +228,7 @@ def test_process_sync_data_existing_route_and_hostnames(setup_connection_manager
     result = process_sync_data(connection_manager, data_update, None)
 
     # Check that the hit count was updated correctly
-    assert (
-        connection_manager.routes.get({"method": "GET", "route": "/api/v1/resource"})[
-            "hits"
-        ]
-        == 15
-    )
+    assert connection_manager.routes.get("GET", "/api/v1/resource")["hits"] == 15
 
     # Check that the total requests were updated
     assert connection_manager.statistics.get_record()["requests"] == {
@@ -267,11 +251,24 @@ def test_process_sync_data_existing_route_and_hostnames(setup_connection_manager
 def test_process_sync_data_no_routes(setup_connection_manager):
     """Test behavior when no routes are provided."""
     connection_manager = setup_connection_manager
-    data = {"current_routes": {}, "reqs": 0}  # No requests to add
+    data = {"routes": {}, "reqs": 0}  # No requests to add
+
+    result = process_sync_data(connection_manager, data, None)
+
+    # Check that no routes were initialized
+    assert len(connection_manager.routes.export()) == 0
+
+    # Check that the total requests remain unchanged
+
+
+def test_process_sync_data_routes_undefined(setup_connection_manager):
+    """Test behavior when no routes are provided."""
+    connection_manager = setup_connection_manager
+    data = {"sdfsdsdfsdfds": {}, "reqs": 0}  # No requests to add
 
     result = process_sync_data(connection_manager, data, None)
 
     # Check that no routes were initialized
-    assert len(connection_manager.routes) == 0
+    assert len(connection_manager.routes.export()) == 0
 
     # Check that the total requests remain unchanged