diff --git a/cognite/client/config.py b/cognite/client/config.py index 3bda116368..aa73ca223c 100644 --- a/cognite/client/config.py +++ b/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() diff --git a/cognite/client/utils/_api_usage.py b/cognite/client/utils/_api_usage.py new file mode 100644 index 0000000000..8241e536cb --- /dev/null +++ b/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, + )