Add RequestDetails class to be used in usage tracking callbacks

cognite/client/config.py

@@ -3,10 +3,14 @@
 import getpass
 import pprint
 from contextlib import suppress
+from typing import TYPE_CHECKING, Callable
 
 from cognite.client._version import __api_subversion__
 from cognite.client.credentials import CredentialProvider
 
+if TYPE_CHECKING:
+    from cognite.client.utils._api_usage import RequestDetails
+
 
 class GlobalConfig:
     """Global configuration object
@@ -39,6 +43,7 @@ def __init__(self) -> None:
         self.max_connection_pool_size: int = 50
         self.disable_ssl: bool = False
         self.proxies: dict[str, str] | None = {}
+        self.usage_tracking_callback: Callable[[RequestDetails], None] | None = None
 
 
 global_config = GlobalConfig()

cognite/client/utils/_api_usage.py

@@ -0,0 +1,65 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from datetime import timedelta
+from typing import TYPE_CHECKING
+
+from typing_extensions import Self
+
+if TYPE_CHECKING:
+    from requests import Response
+
+
+@dataclass
+class RequestDetails:
+    """
+    SDK users wanting to track their own API usage (with the SDK) - for metrics or surveilance, may set
+    a callback on the global_config object that will then receive instances of this class, one per
+    actual request.
+
+    Args:
+        url (str): The API endpoint that was called.
+        status_code (int): The status code of the API response.
+        content_length (int | None): The size of the response if available.
+        time_elapsed (timedelta): The amount of time elapsed between sending the request and the arrival of the response.
+
+    Example:
+
+        Store info on the last 1000 requests made:
+
+            >>> from cognite.client.config import global_config
+            >>> from collections import deque
+            >>> usage_info = deque(maxlen=1000)
+            >>> global_config.usage_tracking_callback = usage_info.append
+
+        Store the time elapsed per request, grouped per API endpoint, for all requests:
+
+            >>> from collections import defaultdict
+            >>> usage_info = defaultdict(list)
+            >>> def callback(details):
+            ...     usage_info[details.url].append(details.time_elapsed)
+            >>> global_config.usage_tracking_callback = callback
+
+    Note:
+        Due to concurrency, the sum of time_elapsed is much greater than the actual request waiting time.
+
+    Tip:
+        Ensure the provided callback is fast to execute, or it might negatively influence the overall performance.
+    """
+
+    url: str
+    status_code: int
+    content_length: int | None
+    time_elapsed: timedelta
+
+    @classmethod
+    def from_response(cls, resp: Response) -> Self:
+        # If header not set, we don't report the size. We could do len(resp.content), but
+        # for streaming requests this would fetch everything into memory...
+        content_length = int(resp.headers.get("Content-length", 0)) or None
+        return cls(
+            url=resp.url,
+            status_code=resp.status_code,
+            content_length=content_length,
+            time_elapsed=resp.elapsed,
+        )