Add exclude_callback

apitally/client/request_logging.py

@@ -52,6 +52,25 @@
 ]
 
 
+class RequestDict(TypedDict):
+    timestamp: float
+    method: str
+    path: Optional[str]
+    url: str
+    headers: List[Tuple[str, str]]
+    size: Optional[int]
+    consumer: Optional[str]
+    body: Optional[bytes]
+
+
+class ResponseDict(TypedDict):
+    status_code: int
+    response_time: float
+    headers: List[Tuple[str, str]]
+    size: Optional[int]
+    body: Optional[bytes]
+
+
 @dataclass
 class RequestLoggingConfig:
     """
@@ -69,6 +88,7 @@ class RequestLoggingConfig:
         mask_request_body_callback: Callback to mask the request body. Expects (method, path, body) and returns the masked body as bytes or None.
         mask_response_body_callback: Callback to mask the response body. Expects (method, path, body) and returns the masked body as bytes or None.
         exclude_paths: Paths to exclude from logging. Expects regular expressions.
+        exclude_callback: Callback to exclude requests from logging. Should expect two arguments, `request: RequestDict` and `response: ResponseDict`, and return True to exclude the request.
     """
 
     enabled: bool = False
@@ -82,25 +102,7 @@ class RequestLoggingConfig:
     mask_request_body_callback: Optional[Callable[[str, str, bytes], Optional[bytes]]] = None
     mask_response_body_callback: Optional[Callable[[str, str, bytes], Optional[bytes]]] = None
     exclude_paths: List[str] = field(default_factory=list)
-
-
-class RequestDict(TypedDict):
-    timestamp: float
-    method: str
-    path: Optional[str]
-    url: str
-    headers: List[Tuple[str, str]]
-    size: Optional[int]
-    consumer: Optional[str]
-    body: Optional[bytes]
-
-
-class ResponseDict(TypedDict):
-    status_code: int
-    response_time: float
-    headers: List[Tuple[str, str]]
-    size: Optional[int]
-    body: Optional[bytes]
+    exclude_callback: Optional[Callable[[RequestDict, ResponseDict], bool]] = None
 
 
 class TempGzipFile:
@@ -160,7 +162,7 @@ def log_request(self, request: RequestDict, response: ResponseDict) -> None:
         if not self.enabled or self.suspend_until is not None:
             return
         parsed_url = urlparse(request["url"])
-        if self._should_exclude_path(request["path"] or parsed_url.path):
+        if self._should_exclude_path(request["path"] or parsed_url.path) or self._should_exclude(request, response):
             return
 
         query = self._mask_query_params(parsed_url.query) if self.config.log_query_params else ""
@@ -263,6 +265,11 @@ def close(self) -> None:
         self.enabled = False
         self.clear()
 
+    def _should_exclude(self, request: RequestDict, response: ResponseDict) -> bool:
+        if self.config.exclude_callback is not None:
+            return self.config.exclude_callback(request, response)
+        return False
+
     @lru_cache(maxsize=1000)
     def _should_exclude_path(self, url_path: str) -> bool:
         patterns = self.config.exclude_paths + EXCLUDE_PATH_PATTERNS

tests/test_client_request_logging.py

@@ -107,6 +107,7 @@ def test_request_log_exclusion(request_logger: RequestLogger, request_dict: Requ
     request_logger.config.log_response_headers = False
     request_logger.config.log_response_body = False
     request_logger.config.exclude_paths = ["/excluded$"]
+    request_logger.config.exclude_callback = lambda _, response: response["status_code"] == 404
 
     request_logger.log_request(request_dict, response_dict)
     assert len(request_logger.write_deque) == 1
@@ -117,6 +118,11 @@ def test_request_log_exclusion(request_logger: RequestLogger, request_dict: Requ
     assert "headers" not in item["response"]
     assert "body" not in item["response"]
 
+    response_dict["status_code"] = 404
+    request_logger.log_request(request_dict, response_dict)
+    assert len(request_logger.write_deque) == 1
+    response_dict["status_code"] = 200
+
     request_dict["path"] = "/api/excluded"
     request_logger.log_request(request_dict, response_dict)
     assert len(request_logger.write_deque) == 1