cognite-neat 0.123.32__py3-none-any.whl → 0.123.34__py3-none-any.whl

cognite/neat/_utils/http_client/_client.py

@@ -0,0 +1,284 @@
+import gzip
+import random
+import sys
+import time
+from collections import deque
+from collections.abc import MutableMapping, Sequence, Set
+from typing import Literal
+
+import httpx
+from cognite.client import ClientConfig, global_config
+from cognite.client.utils import _json
+
+from cognite.neat._utils.auxiliary import get_current_neat_version
+from cognite.neat._utils.http_client._config import get_user_agent
+from cognite.neat._utils.http_client._data_classes import (
+    BodyRequest,
+    FailedRequestMessage,
+    HTTPMessage,
+    ItemsRequest,
+    ParamRequest,
+    RequestMessage,
+    ResponseMessage,
+)
+from cognite.neat._utils.useful_types import JsonVal
+
+if sys.version_info >= (3, 11):
+    from typing import Self
+else:
+    from typing_extensions import Self
+
+
+class HTTPClient:
+    """An HTTP client.
+
+    This class handles rate limiting, retries, and error handling for HTTP requests.
+
+    Args:
+        config (ClientConfig): Configuration for the client.
+        pool_connections (int): The number of connection pools to cache. Default is 10.
+        pool_maxsize (int): The maximum number of connections to save in the pool. Default
+            is 20.
+        max_retries (int): The maximum number of retries for a request. Default is 10.
+        retry_status_codes (frozenset[int]): HTTP status codes that should trigger a retry.
+            Default is {408, 429, 502, 503, 504}.
+        split_items_status_codes (frozenset[int]): In the case of ItemRequest with multiple
+            items, these status codes will trigger splitting the request into smaller batches.
+
+    """
+
+    def __init__(
+        self,
+        config: ClientConfig,
+        max_retries: int = 10,
+        pool_connections: int = 10,
+        pool_maxsize: int = 20,
+        retry_status_codes: Set[int] = frozenset({429, 502, 503, 504}),
+        split_items_status_codes: Set[int] = frozenset({400, 408, 409, 422, 502, 503, 504}),
+    ):
+        self.config = config
+        self._max_retries = max_retries
+        self._pool_connections = pool_connections
+        self._pool_maxsize = pool_maxsize
+        self._retry_status_codes = retry_status_codes
+        self._split_items_status_codes = split_items_status_codes
+
+        # Thread-safe session for connection pooling
+        self.session = self._create_thread_safe_session()
+
+    def __enter__(self) -> Self:
+        return self
+
+    def __exit__(
+        self, exc_type: type[BaseException] | None, exc_value: BaseException | None, traceback: object | None
+    ) -> Literal[False]:
+        """Close the session when exiting the context."""
+        self.session.close()
+        return False  # Do not suppress exceptions
+
+    def request(self, message: RequestMessage) -> Sequence[HTTPMessage]:
+        """Send an HTTP request and return the response.
+
+        Args:
+            message (RequestMessage): The request message to send.
+
+        Returns:
+            Sequence[HTTPMessage]: The response message(s). This can also
+                include RequestMessage(s) to be retried.
+        """
+        if isinstance(message, ItemsRequest) and message.tracker and message.tracker.limit_reached():
+            error_msg = (
+                f"Aborting further splitting of requests after {message.tracker.failed_split_count} failed attempts."
+            )
+            return message.create_failed_request(error_msg)
+        try:
+            response = self._make_request(message)
+            results = self._handle_response(response, message)
+        except Exception as e:
+            results = self._handle_error(e, message)
+        return results
+
+    def request_with_retries(self, message: RequestMessage) -> Sequence[ResponseMessage | FailedRequestMessage]:
+        """Send an HTTP request and handle retries.
+
+        This method will keep retrying the request until it either succeeds or
+        exhausts the maximum number of retries.
+
+        Note this method will use the current thread to process all request, thus
+        it is blocking.
+
+        Args:
+            message (RequestMessage): The request message to send.
+
+        Returns:
+            Sequence[ResponseMessage | FailedRequestMessage]: The final response
+                messages, which can be either successful responses or failed requests.
+        """
+        if message.total_attempts > 0:
+            raise RuntimeError(f"RequestMessage has already been attempted {message.total_attempts} times.")
+        pending_requests: deque[RequestMessage] = deque()
+        pending_requests.append(message)
+        final_responses: list[ResponseMessage | FailedRequestMessage] = []
+
+        while pending_requests:
+            current_request = pending_requests.popleft()
+            results = self.request(current_request)
+
+            for result in results:
+                if isinstance(result, RequestMessage):
+                    pending_requests.append(result)
+                elif isinstance(result, ResponseMessage | FailedRequestMessage):
+                    final_responses.append(result)
+                else:
+                    raise TypeError(f"Unexpected result type: {type(result)}")
+
+        return final_responses
+
+    def _create_thread_safe_session(self) -> httpx.Client:
+        return httpx.Client(
+            limits=httpx.Limits(
+                max_connections=self._pool_maxsize,
+                max_keepalive_connections=self._pool_connections,
+            ),
+            timeout=self.config.timeout,
+        )
+
+    def _create_headers(self, api_version: str | None = None) -> MutableMapping[str, str]:
+        headers: MutableMapping[str, str] = {}
+        headers["User-Agent"] = f"httpx/{httpx.__version__} {get_user_agent()}"
+        auth_name, auth_value = self.config.credentials.authorization_header()
+        headers[auth_name] = auth_value
+        headers["content-type"] = "application/json"
+        headers["accept"] = "application/json"
+        headers["x-cdp-sdk"] = f"CogniteNeat:{get_current_neat_version()}"
+        headers["x-cdp-app"] = self.config.client_name
+        headers["cdf-version"] = api_version or self.config.api_subversion
+        if not global_config.disable_gzip:
+            headers["Content-Encoding"] = "gzip"
+        return headers
+
+    @staticmethod
+    def _prepare_payload(item: BodyRequest) -> str | bytes:
+        """
+        Prepare the payload for the HTTP request.
+        This method should be overridden in subclasses to customize the payload format.
+        """
+        data: str | bytes
+        try:
+            data = _json.dumps(item.body(), allow_nan=False)
+        except ValueError as e:
+            # A lot of work to give a more human friendly error message when nans and infs are present:
+            msg = "Out of range float values are not JSON compliant"
+            if msg in str(e):  # exc. might e.g. contain an extra ": nan", depending on build (_json.make_encoder)
+                raise ValueError(f"{msg}. Make sure your data does not contain NaN(s) or +/- Inf!").with_traceback(
+                    e.__traceback__
+                ) from None
+            raise
+
+        if not global_config.disable_gzip:
+            data = gzip.compress(data.encode())
+        return data
+
+    def _make_request(self, item: RequestMessage) -> httpx.Response:
+        headers = self._create_headers(item.api_version)
+        params: dict[str, str] | None = None
+        if isinstance(item, ParamRequest):
+            params = item.parameters
+        data: str | bytes | None = None
+        if isinstance(item, BodyRequest):
+            data = self._prepare_payload(item)
+        return self.session.request(
+            method=item.method,
+            url=item.endpoint_url,
+            content=data,
+            headers=headers,
+            params=params,
+            timeout=self.config.timeout,
+            follow_redirects=False,
+        )
+
+    def _handle_response(
+        self,
+        response: httpx.Response,
+        request: RequestMessage,
+    ) -> Sequence[HTTPMessage]:
+        try:
+            body = response.json()
+        except ValueError as e:
+            return request.create_responses(response, error_message=f"Invalid JSON response: {e!s}")
+
+        error_obj = body.get("error", {})
+        is_auto_retryable = False
+        if isinstance(error_obj, dict):
+            is_auto_retryable = error_obj.get("isAutoRetryable", False)
+
+        if 200 <= response.status_code < 300:
+            return request.create_responses(response, body)
+        elif (
+            isinstance(request, ItemsRequest)
+            and len(request.items) > 1
+            and response.status_code in self._split_items_status_codes
+        ):
+            # 4XX: Status there is at least one item that is invalid, split the batch to get all valid items processed
+            # 5xx: Server error, split to reduce the number of items in each request, and count as a status attempt
+            status_attempts = request.status_attempt
+            if 500 <= response.status_code < 600:
+                status_attempts += 1
+            splits = request.split(status_attempts=status_attempts)
+            if splits[0].tracker and splits[0].tracker.limit_reached():
+                return request.create_responses(response, body, self._get_error_message(body, response.text))
+            return splits
+        elif request.status_attempt < self._max_retries and (
+            response.status_code in self._retry_status_codes or is_auto_retryable
+        ):
+            request.status_attempt += 1
+            time.sleep(self._backoff_time(request.total_attempts))
+            return [request]
+        else:
+            # Permanent failure
+            return request.create_responses(response, body, self._get_error_message(body, response.text))
+
+    @staticmethod
+    def _get_error_message(body: JsonVal, default: str) -> str:
+        error = default
+        if not isinstance(body, dict):
+            return error
+        if "error" not in body:
+            return error
+        error_nested = body["error"]
+        if isinstance(error_nested, str):
+            return error_nested
+        if isinstance(error_nested, dict) and "message" in error_nested and isinstance(error_nested["message"], str):
+            return error_nested["message"]
+        return error
+
+    @staticmethod
+    def _backoff_time(attempts: int) -> float:
+        backoff_time = 0.5 * (2**attempts)
+        return min(backoff_time, global_config.max_retry_backoff) * random.uniform(0, 1.0)
+
+    def _handle_error(
+        self,
+        e: Exception,
+        request: RequestMessage,
+    ) -> Sequence[HTTPMessage]:
+        if isinstance(e, httpx.ReadTimeout | httpx.TimeoutException):
+            error_type = "read"
+            request.read_attempt += 1
+            attempts = request.read_attempt
+        elif isinstance(e, ConnectionError | httpx.ConnectError | httpx.ConnectTimeout):
+            error_type = "connect"
+            request.connect_attempt += 1
+            attempts = request.connect_attempt
+        else:
+            error_msg = f"Unexpected exception: {e!s}"
+            return request.create_failed_request(error_msg)
+
+        if attempts <= self._max_retries:
+            time.sleep(self._backoff_time(request.total_attempts))
+            return [request]
+        else:
+            # We have already incremented the attempt count, so we subtract 1 here
+            error_msg = f"RequestException after {request.total_attempts - 1} attempts ({error_type} error): {e!s}"
+
+            return request.create_failed_request(error_msg)

cognite/neat/_utils/http_client/_config.py

@@ -0,0 +1,19 @@
+import functools
+import platform
+
+from cognite.neat._utils.auxiliary import get_current_neat_version
+
+
+@functools.lru_cache(maxsize=1)
+def get_user_agent() -> str:
+    neat_version = f"CogniteNeat/{get_current_neat_version()}"
+    python_version = (
+        f"{platform.python_implementation()}/{platform.python_version()} "
+        f"({platform.python_build()};{platform.python_compiler()})"
+    )
+    os_version_info = [platform.release(), platform.machine(), platform.architecture()[0]]
+    os_version_info = [s for s in os_version_info if s]  # Ignore empty strings
+    os_version_info_str = "-".join(os_version_info)
+    operating_system = f"{platform.system()}/{os_version_info_str}"
+
+    return f"{neat_version} {python_version} {operating_system}"

cognite/neat/_utils/http_client/_data_classes.py

@@ -0,0 +1,389 @@
+from abc import ABC, abstractmethod
+from collections.abc import Callable, Sequence
+from dataclasses import dataclass, field
+from typing import Generic, Literal, TypeAlias
+
+import httpx
+
+from cognite.neat._utils.http_client._tracker import ItemsRequestTracker
+from cognite.neat._utils.useful_types import T_ID, JsonVal
+
+StatusCode: TypeAlias = int
+
+
+@dataclass
+class HTTPMessage:
+    """Base class for HTTP messages (requests and responses)"""
+
+    def dump(self) -> dict[str, JsonVal]:
+        """Dumps the message to a JSON serializable dictionary.
+
+        Returns:
+            dict[str, JsonVal]: The message as a dictionary.
+        """
+        # We avoid using the asdict function as we know we have a shallow structure,
+        # and this roughly ~10x faster.
+        output = self.__dict__.copy()
+        output["type"] = type(self).__name__
+        return output
+
+
+@dataclass
+class FailedRequestMessage(HTTPMessage):
+    error: str
+
+
+@dataclass
+class ResponseMessage(HTTPMessage):
+    status_code: StatusCode
+
+
+@dataclass
+class RequestMessage(HTTPMessage):
+    """Base class for HTTP request messages"""
+
+    endpoint_url: str
+    method: Literal["GET", "POST", "PATCH", "DELETE"]
+    connect_attempt: int = 0
+    read_attempt: int = 0
+    status_attempt: int = 0
+    api_version: str | None = None
+
+    @property
+    def total_attempts(self) -> int:
+        return self.connect_attempt + self.read_attempt + self.status_attempt
+
+    @abstractmethod
+    def create_responses(
+        self,
+        response: httpx.Response,
+        response_body: dict[str, JsonVal] | None = None,
+        error_message: str | None = None,
+    ) -> Sequence[HTTPMessage]:
+        raise NotImplementedError()
+
+    @abstractmethod
+    def create_failed_request(self, error_message: str) -> Sequence[HTTPMessage]:
+        raise NotImplementedError()
+
+
+@dataclass
+class SuccessResponse(ResponseMessage):
+    body: dict[str, JsonVal] | None = None
+
+
+@dataclass
+class FailedResponse(ResponseMessage):
+    error: str
+    body: dict[str, JsonVal] | None = None
+
+
+@dataclass
+class SimpleRequest(RequestMessage):
+    """Base class for requests with a simple success/fail response structure"""
+
+    @classmethod
+    def create_responses(
+        cls,
+        response: httpx.Response,
+        response_body: dict[str, JsonVal] | None = None,
+        error_message: str | None = None,
+    ) -> Sequence[ResponseMessage]:
+        if 200 <= response.status_code < 300 and error_message is None:
+            return [SuccessResponse(status_code=response.status_code, body=response_body)]
+        if error_message is None:
+            error_message = f"Request failed with status code {response.status_code}"
+        return [FailedResponse(status_code=response.status_code, error=error_message, body=response_body)]
+
+    @classmethod
+    def create_failed_request(cls, error_message: str) -> Sequence[HTTPMessage]:
+        return [FailedRequestMessage(error=error_message)]
+
+
+@dataclass
+class BodyRequest(RequestMessage, ABC):
+    """Base class for HTTP request messages with a body"""
+
+    @abstractmethod
+    def body(self) -> dict[str, JsonVal]:
+        raise NotImplementedError()
+
+
+@dataclass
+class ParamRequest(SimpleRequest):
+    """Base class for HTTP request messages with query parameters"""
+
+    parameters: dict[str, str] | None = None
+
+
+@dataclass
+class SimpleBodyRequest(SimpleRequest, BodyRequest):
+    body_content: dict[str, JsonVal] = field(default_factory=dict)
+
+    def body(self) -> dict[str, JsonVal]:
+        return self.body_content
+
+
+@dataclass
+class ItemMessage:
+    """Base class for message related to a specific item"""
+
+    ...
+
+
+@dataclass
+class ItemIDMessage(Generic[T_ID], ItemMessage, ABC):
+    """Base class for message related to a specific item identified by an ID"""
+
+    id: T_ID
+
+
+@dataclass
+class ItemResponse(ItemIDMessage, ResponseMessage, ABC): ...
+
+
+@dataclass
+class SuccessItem(ItemResponse):
+    item: JsonVal | None = None
+
+
+@dataclass
+class FailedItem(ItemResponse):
+    error: str
+
+
+@dataclass
+class MissingItem(ItemResponse): ...
+
+
+@dataclass
+class UnexpectedItem(ItemResponse):
+    item: JsonVal | None = None
+
+
+@dataclass
+class FailedRequestItem(ItemIDMessage, FailedRequestMessage): ...
+
+
+@dataclass
+class UnknownRequestItem(ItemMessage, FailedRequestMessage):
+    item: JsonVal | None = None
+
+
+@dataclass
+class UnknownResponseItem(ItemMessage, ResponseMessage):
+    error: str
+    item: JsonVal | None = None
+
+
+@dataclass
+class ItemsRequest(Generic[T_ID], BodyRequest):
+    """Requests message for endpoints that accept multiple items in a single request.
+
+    This class provides functionality to split large requests into smaller ones, handle responses for each item,
+    and manage errors effectively.
+
+    Attributes:
+        items (list[JsonVal]): The list of items to be sent in the request body.
+        extra_body_fields (dict[str, JsonVal]): Additional fields to include in the request body
+        as_id (Callable[[JsonVal], T_ID] | None): A function to extract the ID from each item. If None,
+            IDs are not used.
+        max_failures_before_abort (int): The maximum number of failed split requests before aborting further splits.
+
+    """
+
+    items: list[JsonVal] = field(default_factory=list)
+    extra_body_fields: dict[str, JsonVal] = field(default_factory=dict)
+    as_id: Callable[[JsonVal], T_ID] | None = None
+    max_failures_before_abort: int = 50
+    tracker: ItemsRequestTracker | None = field(default=None, init=False)
+
+    def dump(self) -> dict[str, JsonVal]:
+        """Dumps the message to a JSON serializable dictionary.
+
+        This override removes the 'as_id' attribute as it is not serializable.
+
+        Returns:
+            dict[str, JsonVal]: The message as a dictionary.
+        """
+        output = super().dump()
+        if self.as_id is not None:
+            # We cannot serialize functions
+            del output["as_id"]
+        if self.tracker is not None:
+            # We cannot serialize the tracker
+            del output["tracker"]
+        return output
+
+    def body(self) -> dict[str, JsonVal]:
+        if self.extra_body_fields:
+            return {"items": self.items, **self.extra_body_fields}
+        return {"items": self.items}
+
+    def split(self, status_attempts: int) -> "list[ItemsRequest]":
+        """Splits the request into two smaller requests.
+
+        This is useful for retrying requests that fail due to size limits or timeouts.
+
+        Args:
+            status_attempts: The number of status attempts to set for the new requests. This is used when the
+                request failed with a 5xx status code and we want to track the number of attempts. For 4xx errors,
+                there is at least one item causing the error, so we do not increment the status attempts, but
+                instead essentially do a binary search to find the problematic item(s).
+
+        Returns:
+            A list containing two new ItemsRequest instances, each with half of the original items.
+
+        """
+        mid = len(self.items) // 2
+        if mid == 0:
+            return [self]
+        tracker = self.tracker or ItemsRequestTracker(self.max_failures_before_abort)
+        tracker.register_failure()
+        first_half = ItemsRequest[T_ID](
+            endpoint_url=self.endpoint_url,
+            method=self.method,
+            items=self.items[:mid],
+            extra_body_fields=self.extra_body_fields,
+            as_id=self.as_id,
+            connect_attempt=self.connect_attempt,
+            read_attempt=self.read_attempt,
+            status_attempt=status_attempts,
+        )
+        first_half.tracker = tracker
+        second_half = ItemsRequest[T_ID](
+            endpoint_url=self.endpoint_url,
+            method=self.method,
+            items=self.items[mid:],
+            extra_body_fields=self.extra_body_fields,
+            as_id=self.as_id,
+            connect_attempt=self.connect_attempt,
+            read_attempt=self.read_attempt,
+            status_attempt=status_attempts,
+        )
+        second_half.tracker = tracker
+        return [first_half, second_half]
+
+    def create_responses(
+        self,
+        response: httpx.Response,
+        response_body: dict[str, JsonVal] | None = None,
+        error_message: str | None = None,
+    ) -> Sequence[HTTPMessage]:
+        """Creates response messages based on the HTTP response and the original request.
+
+        Args:
+            response: The HTTP response received from the server.
+            response_body: The parsed JSON body of the response, if available.
+            error_message: An optional error message to use if the response indicates a failure.
+
+        Returns:
+            A sequence of HTTPMessage instances representing the outcome for each item in the request.
+        """
+        if self.as_id is None:
+            return SimpleBodyRequest.create_responses(response, response_body, error_message)
+        request_items_by_id, errors = self._create_items_by_id()
+        responses: list[HTTPMessage] = list(errors)
+        error_message = error_message or "Unknown error"
+
+        if not self._is_items_response(response_body):
+            return self._handle_non_items_response(responses, response, error_message, request_items_by_id)
+
+        # Process items from response
+        if response_body is not None:
+            self._process_response_items(responses, response, response_body, error_message, request_items_by_id)
+
+        # Handle missing items
+        self._handle_missing_items(responses, response, request_items_by_id)
+
+        return responses
+
+    @staticmethod
+    def _handle_non_items_response(
+        responses: list[HTTPMessage],
+        response: httpx.Response,
+        error_message: str,
+        request_items_by_id: dict[T_ID, JsonVal],
+    ) -> list[HTTPMessage]:
+        """Handles responses that do not contain an 'items' field in the body."""
+        if 200 <= response.status_code < 300:
+            responses.extend(
+                SuccessItem(status_code=response.status_code, id=id_) for id_ in request_items_by_id.keys()
+            )
+        else:
+            responses.extend(
+                FailedItem(status_code=response.status_code, error=error_message, id=id_)
+                for id_ in request_items_by_id.keys()
+            )
+        return responses
+
+    def _process_response_items(
+        self,
+        responses: list[HTTPMessage],
+        response: httpx.Response,
+        response_body: dict[str, JsonVal],
+        error_message: str,
+        request_items_by_id: dict[T_ID, JsonVal],
+    ) -> None:
+        """Processes each item in the response body and categorizes them based on their status."""
+        for response_item in response_body["items"]:  # type: ignore[union-attr]
+            try:
+                item_id = self.as_id(response_item)  # type: ignore[misc]
+            except Exception as e:
+                responses.append(
+                    UnknownResponseItem(
+                        status_code=response.status_code, item=response_item, error=f"Error extracting ID: {e!s}"
+                    )
+                )
+                continue
+            request_item = request_items_by_id.pop(item_id, None)
+            if request_item is None:
+                responses.append(UnexpectedItem(status_code=response.status_code, id=item_id, item=response_item))
+            elif 200 <= response.status_code < 300:
+                responses.append(SuccessItem(status_code=response.status_code, id=item_id, item=response_item))
+            else:
+                responses.append(FailedItem(status_code=response.status_code, id=item_id, error=error_message))
+
+    @staticmethod
+    def _handle_missing_items(
+        responses: list[HTTPMessage], response: httpx.Response, request_items_by_id: dict[T_ID, JsonVal]
+    ) -> None:
+        """Handles items that were in the request but not present in the response."""
+        for item_id in request_items_by_id.keys():
+            responses.append(MissingItem(status_code=response.status_code, id=item_id))
+
+    def create_failed_request(self, error_message: str) -> Sequence[HTTPMessage]:
+        if self.as_id is None:
+            return SimpleBodyRequest.create_failed_request(error_message)
+        items_by_id, errors = self._create_items_by_id()
+        results: list[HTTPMessage] = []
+        results.extend(errors)
+        results.extend(FailedRequestItem(id=item_id, error=error_message) for item_id in items_by_id.keys())
+        return results
+
+    def _create_items_by_id(self) -> tuple[dict[T_ID, JsonVal], list[FailedRequestItem | UnknownRequestItem]]:
+        if self.as_id is None:
+            raise ValueError("as_id function must be provided to create items by ID")
+        items_by_id: dict[T_ID, JsonVal] = {}
+        errors: list[FailedRequestItem | UnknownRequestItem] = []
+        for item in self.items:
+            try:
+                item_id = self.as_id(item)
+            except Exception as e:
+                errors.append(UnknownRequestItem(error=f"Error extracting ID: {e!s}", item=item))
+                continue
+            if item_id in items_by_id:
+                errors.append(FailedRequestItem(id=item_id, error=f"Duplicate item ID: {item_id!r}"))
+            else:
+                items_by_id[item_id] = item
+        return items_by_id, errors
+
+    @staticmethod
+    def _is_items_response(body: dict[str, JsonVal] | None = None) -> bool:
+        if body is None:
+            return False
+        if "items" not in body:
+            return False
+        if not isinstance(body["items"], list):
+            return False
+        return True