Add optional callback into HTTPClient.request

cognitedata · Nov 4, 2023 · ed91068 · ed91068
1 parent 69b1b0b
commit ed91068
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/cognite/client/_http_client.py b/cognite/client/_http_client.py
@@ -13,6 +13,7 @@
 
 from cognite.client.config import global_config
 from cognite.client.exceptions import CogniteConnectionError, CogniteConnectionRefused, CogniteReadTimeout
+from cognite.client.utils._api_usage import RequestDetails
 
 
 class BlockAll(cookiejar.CookiePolicy):
@@ -114,6 +115,10 @@ def request(self, method: str, url: str, **kwargs: Any) -> requests.Response:
         while True:
             try:
                 res = self._do_request(method=method, url=url, **kwargs)
+                if global_config.usage_tracking_callback:
+                    # We use the RequestDetails as an indirection to avoid the user mutating the request:
+                    global_config.usage_tracking_callback(RequestDetails.from_response(res))
+
                 last_status = res.status_code
                 retry_tracker.status += 1
                 if not retry_tracker.should_retry(status_code=last_status):