cognite-toolkit 0.5.97__py3-none-any.whl → 0.7.34__py3-none-any.whl

cognite_toolkit/_cdf_tk/utils/http_client/_data_classes.py

@@ -0,0 +1,412 @@
+from abc import ABC, abstractmethod
+from collections import UserList
+from collections.abc import Hashable, Sequence
+from dataclasses import dataclass, field
+from typing import Generic, Literal, Protocol, TypeAlias, TypeVar
+
+import httpx
+from cognite.client.utils import _json
+
+from cognite_toolkit._cdf_tk.utils.http_client._exception import ToolkitAPIError
+from cognite_toolkit._cdf_tk.utils.http_client._tracker import ItemsRequestTracker
+from cognite_toolkit._cdf_tk.utils.useful_types import JsonVal, PrimitiveType
+
+StatusCode: TypeAlias = int
+
+
+@dataclass
+class HTTPMessage(ABC):
+    """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 ErrorDetails:
+    """This is the expected structure of error details in the CDF API"""
+
+    code: int
+    message: str
+    missing: list[JsonVal] | None = None
+    duplicated: list[JsonVal] | None = None
+    is_auto_retryable: bool | None = None
+
+    @classmethod
+    def from_response(cls, response: httpx.Response) -> "ErrorDetails":
+        try:
+            error_data = response.json()["error"]
+            if not isinstance(error_data, dict):
+                return cls(code=response.status_code, message=str(error_data))
+            return cls(
+                code=error_data["code"],
+                message=error_data["message"],
+                missing=error_data.get("missing"),
+                duplicated=error_data.get("duplicated"),
+                is_auto_retryable=error_data.get("isAutoRetryable"),
+            )
+        except (ValueError, KeyError):
+            # Fallback if response is not JSON or does not have expected structure
+            return cls(code=response.status_code, message=response.text)
+
+    def dump(self) -> dict[str, JsonVal]:
+        output: dict[str, JsonVal] = {
+            "code": self.code,
+            "message": self.message,
+        }
+        if self.missing is not None:
+            output["missing"] = self.missing
+        if self.duplicated is not None:
+            output["duplicated"] = self.duplicated
+        if self.is_auto_retryable is not None:
+            output["isAutoRetryable"] = self.is_auto_retryable
+        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", "PUT"]
+    connect_attempt: int = 0
+    read_attempt: int = 0
+    status_attempt: int = 0
+    api_version: str | None = None
+    content_type: str = "application/json"
+    accept: str = "application/json"
+    content_length: int | None = None
+
+    @property
+    def total_attempts(self) -> int:
+        return self.connect_attempt + self.read_attempt + self.status_attempt
+
+    @abstractmethod
+    def create_success_response(self, response: httpx.Response) -> Sequence[HTTPMessage]:
+        raise NotImplementedError()
+
+    @abstractmethod
+    def create_failure_response(self, response: httpx.Response) -> Sequence[HTTPMessage]:
+        raise NotImplementedError()
+
+    @abstractmethod
+    def create_failed_request(self, error_message: str) -> Sequence[HTTPMessage]:
+        raise NotImplementedError()
+
+
+@dataclass
+class SuccessResponse(ResponseMessage):
+    body: str
+    content: bytes
+
+    def dump(self) -> dict[str, JsonVal]:
+        output = super().dump()
+        # We cannot serialize bytes, so we indicate its presence instead
+        output["content"] = "<bytes>" if self.content else None
+        return output
+
+
+@dataclass
+class FailedResponse(ResponseMessage):
+    body: str
+    error: ErrorDetails
+
+    def dump(self) -> dict[str, JsonVal]:
+        output = super().dump()
+        output["error"] = self.error.dump()
+        return output
+
+
+@dataclass
+class SimpleRequest(RequestMessage):
+    """Base class for requests with a simple success/fail response structure"""
+
+    @classmethod
+    def create_success_response(cls, response: httpx.Response) -> Sequence[ResponseMessage]:
+        return [SuccessResponse(status_code=response.status_code, body=response.text, content=response.content)]
+
+    @classmethod
+    def create_failure_response(cls, response: httpx.Response) -> Sequence[HTTPMessage]:
+        return [
+            FailedResponse(
+                status_code=response.status_code, error=ErrorDetails.from_response(response), body=response.text
+            )
+        ]
+
+    @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 data(self) -> str:
+        raise NotImplementedError()
+
+
+@dataclass
+class ParamRequest(SimpleRequest):
+    """Base class for HTTP request messages with query parameters"""
+
+    parameters: dict[str, PrimitiveType] | None = None
+
+
+@dataclass
+class SimpleBodyRequest(SimpleRequest, BodyRequest):
+    body_content: dict[str, JsonVal] = field(default_factory=dict)
+
+    def data(self) -> str:
+        return _dump_body(self.body_content)
+
+
+@dataclass
+class DataBodyRequest(SimpleRequest):
+    data_content: bytes = b""
+
+    def data(self) -> bytes:
+        return self.data_content
+
+    def dump(self) -> dict[str, JsonVal]:
+        output = super().dump()
+        # We cannot serialize bytes, so we indicate its presence instead
+        output["data_content"] = "<bytes>"
+        return output
+
+
+T_COVARIANT_ID = TypeVar("T_COVARIANT_ID", covariant=True)
+
+
+@dataclass
+class ItemMessage(Generic[T_COVARIANT_ID], ABC):
+    """Base class for message related to a specific item identified by an ID"""
+
+    ids: list[T_COVARIANT_ID] = field(default_factory=list)
+
+
+@dataclass
+class SuccessResponseItems(ItemMessage[T_COVARIANT_ID], SuccessResponse): ...
+
+
+@dataclass
+class FailedResponseItems(ItemMessage[T_COVARIANT_ID], FailedResponse): ...
+
+
+@dataclass
+class FailedRequestItems(ItemMessage[T_COVARIANT_ID], FailedRequestMessage): ...
+
+
+class RequestItem(Generic[T_COVARIANT_ID], Protocol):
+    def dump(self) -> JsonVal: ...
+
+    def as_id(self) -> T_COVARIANT_ID: ...
+
+
+@dataclass
+class ItemsRequest(Generic[T_COVARIANT_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[T_RequestItem]): 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
+        max_failures_before_abort (int): The maximum number of failed split requests before aborting further splits.
+
+    """
+
+    items: list[RequestItem[T_COVARIANT_ID]] = field(default_factory=list)
+    extra_body_fields: dict[str, JsonVal] = field(default_factory=dict)
+    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()
+        output["items"] = self.dump_items()
+        if self.tracker is not None:
+            # We cannot serialize the tracker
+            del output["tracker"]
+        return output
+
+    def dump_items(self) -> list[JsonVal]:
+        """Dumps the items to a list of JSON serializable dictionaries.
+
+        Returns:
+            list[JsonVal]: The items as a list of dictionaries.
+        """
+        return [item.dump() for item in self.items]
+
+    def body(self) -> dict[str, JsonVal]:
+        if self.extra_body_fields:
+            return {"items": self.dump_items(), **self.extra_body_fields}
+        return {"items": self.dump_items()}
+
+    def data(self) -> str:
+        return _dump_body(self.body())
+
+    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_COVARIANT_ID](
+            endpoint_url=self.endpoint_url,
+            method=self.method,
+            items=self.items[:mid],
+            extra_body_fields=self.extra_body_fields,
+            connect_attempt=self.connect_attempt,
+            read_attempt=self.read_attempt,
+            status_attempt=status_attempts,
+        )
+        first_half.tracker = tracker
+        second_half = ItemsRequest[T_COVARIANT_ID](
+            endpoint_url=self.endpoint_url,
+            method=self.method,
+            items=self.items[mid:],
+            extra_body_fields=self.extra_body_fields,
+            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_success_response(self, response: httpx.Response) -> Sequence[HTTPMessage]:
+        ids = [item.as_id() for item in self.items]
+        return [
+            SuccessResponseItems(
+                status_code=response.status_code, ids=ids, body=response.text, content=response.content
+            )
+        ]
+
+    def create_failure_response(self, response: httpx.Response) -> Sequence[HTTPMessage]:
+        error = ErrorDetails.from_response(response)
+        ids = [item.as_id() for item in self.items]
+        return [FailedResponseItems(status_code=response.status_code, ids=ids, error=error, body=response.text)]
+
+    def create_failed_request(self, error_message: str) -> Sequence[HTTPMessage]:
+        ids = [item.as_id() for item in self.items]
+        return [FailedRequestItems(ids=ids, error=error_message)]
+
+
+class ResponseList(UserList[ResponseMessage | FailedRequestMessage]):
+    def __init__(self, collection: Sequence[ResponseMessage | FailedRequestMessage] | None = None) -> None:
+        super().__init__(collection or [])
+
+    def raise_for_status(self) -> None:
+        """Raises an exception if any response in the list indicates a failure."""
+        failed_responses = [resp for resp in self.data if isinstance(resp, FailedResponse)]
+        failed_requests = [resp for resp in self.data if isinstance(resp, FailedRequestMessage)]
+        if not failed_responses and not failed_requests:
+            return
+        error_messages = "; ".join(f"Status {err.status_code}: {err.error}" for err in failed_responses)
+        if failed_requests:
+            if error_messages:
+                error_messages += "; "
+            error_messages += "; ".join(f"Request error: {err.error}" for err in failed_requests)
+        raise ToolkitAPIError(f"One or more requests failed: {error_messages}")
+
+    @property
+    def has_failed(self) -> bool:
+        """Indicates whether any response in the list indicates a failure.
+
+        Returns:
+            bool: True if there are any failed responses or requests, False otherwise.
+        """
+        for resp in self.data:
+            if isinstance(resp, FailedResponse | FailedRequestMessage):
+                return True
+        return False
+
+    def get_first_body(self) -> dict[str, JsonVal]:
+        """Returns the body of the first successful response in the list.
+
+        Raises:
+            ValueError: If there are no successful responses with a body.
+
+        Returns:
+            dict[str, JsonVal]: The body of the first successful response.
+        """
+        for resp in self.data:
+            if isinstance(resp, SuccessResponse) and resp.body is not None:
+                return _json.loads(resp.body)
+        raise ValueError("No successful responses with a body found.")
+
+    def as_item_responses(self, item_id: Hashable) -> list[ResponseMessage | FailedRequestMessage]:
+        # Convert the responses to per-item responses
+        results: list[ResponseMessage | FailedRequestMessage] = []
+        for message in self.data:
+            if isinstance(message, SuccessResponse):
+                results.append(
+                    SuccessResponseItems(
+                        status_code=message.status_code, content=message.content, ids=[item_id], body=message.body
+                    )
+                )
+            elif isinstance(message, FailedResponse):
+                results.append(
+                    FailedResponseItems(
+                        status_code=message.status_code, ids=[item_id], body=message.body, error=message.error
+                    )
+                )
+            elif isinstance(message, FailedRequestMessage):
+                results.append(FailedRequestItems(ids=[item_id], error=message.error))
+            else:
+                results.append(message)
+        return results
+
+
+def _dump_body(body: dict[str, JsonVal]) -> str:
+    try:
+        return _json.dumps(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

cognite_toolkit/_cdf_tk/utils/http_client/_data_classes2.py

@@ -0,0 +1,249 @@
+import gzip
+import sys
+from abc import ABC, abstractmethod
+from collections import UserList
+from collections.abc import Hashable, Sequence
+from typing import Any, Literal
+
+import httpx
+from cognite.client import global_config
+from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter, model_validator
+from pydantic.alias_generators import to_camel
+
+from cognite_toolkit._cdf_tk.utils.http_client._exception import ToolkitAPIError
+from cognite_toolkit._cdf_tk.utils.http_client._tracker import ItemsRequestTracker
+from cognite_toolkit._cdf_tk.utils.useful_types import PrimitiveType
+
+if sys.version_info >= (3, 11):
+    from typing import Self
+else:
+    from typing_extensions import Self
+
+
+class HTTPResult2(BaseModel):
+    def get_success_or_raise(self) -> "SuccessResponse2":
+        """Raises an exception if any response in the list indicates a failure."""
+        if isinstance(self, SuccessResponse2):
+            return self
+        elif isinstance(self, FailedResponse2):
+            raise ToolkitAPIError(
+                f"Request failed with status code {self.status_code}: {self.error.code} - {self.error.message}"
+            )
+        elif isinstance(self, FailedRequest2):
+            raise ToolkitAPIError(f"Request failed with error: {self.error}")
+        else:
+            raise ToolkitAPIError("Unknown HTTPResult2 type")
+
+
+class FailedRequest2(HTTPResult2):
+    error: str
+
+
+class SuccessResponse2(HTTPResult2):
+    status_code: int
+    body: str
+    content: bytes
+
+    @property
+    def body_json(self) -> dict[str, Any]:
+        """Parse the response body as JSON."""
+        return TypeAdapter(dict[str, JsonValue]).validate_json(self.body)
+
+
+class ErrorDetails2(BaseModel):
+    """This is the expected structure of error details in the CDF API"""
+
+    code: int
+    message: str
+    missing: list[JsonValue] | None = None
+    duplicated: list[JsonValue] | None = None
+    is_auto_retryable: bool | None = None
+
+    @classmethod
+    def from_response(cls, response: httpx.Response) -> "ErrorDetails2":
+        """Populate the error details from a httpx response."""
+        try:
+            res = TypeAdapter(dict[Literal["error"], ErrorDetails2]).validate_json(response.text)
+        except ValueError:
+            return cls(code=response.status_code, message=response.text)
+        return res["error"]
+
+
+class FailedResponse2(HTTPResult2):
+    status_code: int
+    body: str
+    error: ErrorDetails2
+
+
+class BaseRequestMessage(BaseModel, ABC):
+    endpoint_url: str
+    method: Literal["GET", "POST", "PATCH", "DELETE", "PUT"]
+    connect_attempt: int = 0
+    read_attempt: int = 0
+    status_attempt: int = 0
+    api_version: str | None = None
+    content_type: str = "application/json"
+    accept: str = "application/json"
+
+    parameters: dict[str, PrimitiveType] | None = None
+
+    @property
+    def total_attempts(self) -> int:
+        return self.connect_attempt + self.read_attempt + self.status_attempt
+
+    @property
+    @abstractmethod
+    def content(self) -> str | bytes | None: ...
+
+
+class RequestMessage2(BaseRequestMessage):
+    data_content: bytes | None = None
+    body_content: dict[str, JsonValue] | None = None
+
+    @model_validator(mode="before")
+    def check_data_or_body(cls, values: dict[str, Any]) -> dict[str, Any]:
+        if values.get("data_content") is not None and values.get("body_content") is not None:
+            raise ValueError("Only one of data_content or body_content can be set.")
+        return values
+
+    @property
+    def content(self) -> str | bytes | None:
+        data: str | bytes | None = None
+        if self.data_content is not None:
+            data = self.data_content
+            if not global_config.disable_gzip:
+                data = gzip.compress(data)
+        elif self.body_content is not None:
+            # We serialize using pydantic instead of json.dumps. This is because pydantic is faster
+            # and handles more complex types such as datetime, float('nan'), etc.
+            data = _BODY_SERIALIZER.dump_json(self.body_content)
+            if not global_config.disable_gzip and isinstance(data, bytes):
+                data = gzip.compress(data)
+        return data
+
+
+_BODY_SERIALIZER = TypeAdapter(dict[str, JsonValue])
+
+
+class ItemsResultMessage2(BaseModel):
+    ids: list[Hashable]
+
+
+class ItemsFailedRequest2(ItemsResultMessage2):
+    error_message: str
+
+
+class ItemsSuccessResponse2(ItemsResultMessage2):
+    status_code: int
+    body: str
+    content: bytes
+
+
+class ItemsFailedResponse2(ItemsResultMessage2):
+    status_code: int
+    error: ErrorDetails2
+    body: str
+
+
+class BaseModelObject(BaseModel):
+    """Base class for all object. This includes resources and nested objects."""
+
+    # We allow extra fields to support forward compatibility.
+    model_config = ConfigDict(alias_generator=to_camel, extra="allow")
+
+    def dump(self, camel_case: bool = True) -> dict[str, Any]:
+        """Dump the resource to a dictionary.
+
+        This is the default serialization method for request resources.
+        """
+        return self.model_dump(mode="json", by_alias=camel_case, exclude_unset=True)
+
+    @classmethod
+    def _load(cls, resource: dict[str, Any]) -> "Self":
+        """Load method to match CogniteResource signature."""
+        return cls.model_validate(resource)
+
+
+class RequestResource(BaseModelObject, ABC):
+    @abstractmethod
+    def as_id(self) -> Hashable: ...
+
+
+def _set_default_tracker(data: dict[str, Any]) -> ItemsRequestTracker:
+    if "tracker" not in data or data["tracker"] is None:
+        return ItemsRequestTracker(data.get("max_failures_before_abort", 50))
+    return data["tracker"]
+
+
+class ItemsRequest2(BaseRequestMessage):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+    items: Sequence[RequestResource]
+    extra_body_fields: dict[str, JsonValue] | None = None
+    max_failures_before_abort: int = 50
+    tracker: ItemsRequestTracker = Field(init=False, default_factory=_set_default_tracker, exclude=True)
+
+    @property
+    def content(self) -> str | bytes | None:
+        body: dict[str, JsonValue] = {"items": [item.dump() for item in self.items]}
+        if self.extra_body_fields:
+            body.update(self.extra_body_fields)
+        res = _BODY_SERIALIZER.dump_json(body)
+        if not global_config.disable_gzip and isinstance(res, bytes):
+            return gzip.compress(res)
+        return res
+
+    def split(self, status_attempts: int) -> list["ItemsRequest2"]:
+        """Split the request into multiple requests with a single item each."""
+        mid = len(self.items) // 2
+        if mid == 0:
+            return [self]
+        self.tracker.register_failure()
+        messages: list[ItemsRequest2] = []
+        for part in (self.items[:mid], self.items[mid:]):
+            new_request = self.model_copy(update={"items": part, "status_attempt": status_attempts})
+            new_request.tracker = self.tracker
+            messages.append(new_request)
+        return messages
+
+
+class ItemResponse(BaseModel):
+    items: list[dict[str, JsonValue]]
+
+
+class ItemsResultList(UserList[ItemsResultMessage2]):
+    def __init__(self, collection: Sequence[ItemsResultMessage2] | None = None) -> None:
+        super().__init__(collection or [])
+
+    def raise_for_status(self) -> None:
+        """Raises an exception if any response in the list indicates a failure."""
+        failed_responses = [resp for resp in self.data if isinstance(resp, ItemsFailedResponse2)]
+        failed_requests = [resp for resp in self.data if isinstance(resp, ItemsFailedRequest2)]
+        if not failed_responses and not failed_requests:
+            return
+        error_messages = "; ".join(f"Status {err.status_code}: {err.error.message}" for err in failed_responses)
+        if failed_requests:
+            if error_messages:
+                error_messages += "; "
+            error_messages += "; ".join(f"Request error: {err.error_message}" for err in failed_requests)
+        raise ToolkitAPIError(f"One or more requests failed: {error_messages}")
+
+    @property
+    def has_failed(self) -> bool:
+        """Indicates whether any response in the list indicates a failure.
+
+        Returns:
+            bool: True if there are any failed responses or requests, False otherwise.
+        """
+        for resp in self.data:
+            if isinstance(resp, ItemsFailedResponse2 | ItemsFailedRequest2):
+                return True
+        return False
+
+    def get_items(self) -> list[dict[str, JsonValue]]:
+        """Get the items from all successful responses."""
+        items: list[dict[str, JsonValue]] = []
+        for resp in self.data:
+            if isinstance(resp, ItemsSuccessResponse2):
+                body_json = ItemResponse.model_validate_json(resp.body)
+                items.extend(body_json.items)
+        return items

cognite_toolkit/_cdf_tk/utils/http_client/_exception.py

@@ -0,0 +1,4 @@
+class ToolkitAPIError(Exception):
+    """Base class for all exceptions raised by the Cognite Toolkit API client."""
+
+    pass

cognite_toolkit/_cdf_tk/utils/http_client/_tracker.py

@@ -0,0 +1,34 @@
+import threading
+from dataclasses import dataclass, field
+from typing import Any
+
+
+@dataclass
+class ItemsRequestTracker:
+    """Tracks the state of requests split from an original request.
+
+    Attributes:
+        max_failures_before_abort (int): Maximum number of allowed failed split requests before aborting
+            the entire operation. A value of -1 indicates no early abort.
+        lock (Any): A lock to ensure thread-safe updates to the failure count.
+        failed_split_count (int): The current count of failed split requests.
+
+    """
+
+    max_failures_before_abort: int = -1  # -1 means no early abort
+    # NOTE: `threading.Lock` is a factory function (backed by `_thread.allocate_lock`), not a type.
+    # Annotate as `Any` so Pydantic won't attempt to generate a schema for the lock field.
+    lock: Any = field(default_factory=threading.Lock, init=False)
+    failed_split_count: int = field(default=0, init=False)
+
+    def register_failure(self) -> None:
+        """Register a failed split request and return whether to continue splitting."""
+        with self.lock:
+            self.failed_split_count += 1
+
+    def limit_reached(self) -> bool:
+        """Check if the failure limit has been reached."""
+        with self.lock:
+            if self.max_failures_before_abort < 0:
+                return False
+            return self.failed_split_count >= self.max_failures_before_abort