DEV-19025: add configure method

huntflow · Feb 12, 2025 · 9014c7e · 9014c7e
1 parent d272feb
commit 9014c7e
Show file tree

Hide file tree

Showing 4 changed files with 13 additions and 14 deletions.
diff --git a/src/huntflow_base_metrics/web_frameworks/_middleware.py b/src/huntflow_base_metrics/web_frameworks/_middleware.py
@@ -1,7 +1,7 @@
 import abc
 import time
 from dataclasses import dataclass
-from typing import Generic, Optional, Set, TypeVar
+from typing import Generic, Iterable, Optional, Set, TypeVar
 
 from huntflow_base_metrics import apply_labels
 from huntflow_base_metrics._context import METRIC_CONTEXT
@@ -40,6 +40,15 @@ class PrometheusMiddleware(abc.ABC, Generic[RequestType]):
     include_routes: Optional[Set[str]] = None
     exclude_routes: Optional[Set[str]] = None
 
+    @classmethod
+    def configure(
+        cls,
+        include_routes: Optional[Iterable[str]] = None,
+        exclude_routes: Optional[Iterable[str]] = None,
+    ) -> None:
+        cls.include_routes = set(include_routes) if include_routes is not None else None
+        cls.exclude_routes = set(exclude_routes) if exclude_routes is not None else None
+
     @staticmethod
     @abc.abstractmethod
     def get_method(request: RequestType) -> str:

diff --git a/src/huntflow_base_metrics/web_frameworks/aiohttp.py b/src/huntflow_base_metrics/web_frameworks/aiohttp.py
@@ -63,10 +63,7 @@ def add_middleware(
         If it's not empty (and include_routes is not specified), then the
         specified routes will not be observed.
     """
-    include_routes = set(include_routes) if include_routes is not None else include_routes
-    exclude_routes = set(exclude_routes) if exclude_routes is not None else exclude_routes
-    _PrometheusMiddleware.include_routes = include_routes
-    _PrometheusMiddleware.exclude_routes = exclude_routes
+    _PrometheusMiddleware.configure(include_routes, exclude_routes)
     app.middlewares.append(_PrometheusMiddleware.dispatch)
 
 

diff --git a/src/huntflow_base_metrics/web_frameworks/fastapi.py b/src/huntflow_base_metrics/web_frameworks/fastapi.py
@@ -66,10 +66,7 @@ def add_middleware(
         If it's not empty (and include_routes is not specified), then the
         specified routes will not be observed.
     """
-    include_routes = set(include_routes) if include_routes is not None else include_routes
-    exclude_routes = set(exclude_routes) if exclude_routes is not None else exclude_routes
-    _PrometheusMiddleware.include_routes = include_routes
-    _PrometheusMiddleware.exclude_routes = exclude_routes
+    _PrometheusMiddleware.configure(include_routes, exclude_routes)
     app.add_middleware(_PrometheusMiddleware)
 
 

diff --git a/src/huntflow_base_metrics/web_frameworks/litestar.py b/src/huntflow_base_metrics/web_frameworks/litestar.py
@@ -97,11 +97,7 @@ def get_middleware(
         specified routes will not be observed.
     """
 
-    include_routes = set(include_routes) if include_routes is not None else include_routes
-    exclude_routes = set(exclude_routes) if exclude_routes is not None else exclude_routes
-    _PrometheusMiddleware.include_routes = include_routes
-    _PrometheusMiddleware.exclude_routes = exclude_routes
-
+    _PrometheusMiddleware.configure(include_routes, exclude_routes)
     return _PrometheusMiddleware