cognite-toolkit 0.6.39__py3-none-any.whl → 0.7.30__py3-none-any.whl

cognite_toolkit/_cdf_tk/utils/http_client/_data_classes.py

@@ -1,12 +1,15 @@
 from abc import ABC, abstractmethod
-from collections.abc import Callable, Sequence
+from collections import UserList
+from collections.abc import Hashable, Sequence
 from dataclasses import dataclass, field
-from typing import Generic, Literal, TypeAlias
+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 T_ID, JsonVal
+from cognite_toolkit._cdf_tk.utils.useful_types import JsonVal, PrimitiveType
 
 StatusCode: TypeAlias = int
 
@@ -28,6 +31,47 @@ class HTTPMessage(ABC):
         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
@@ -43,23 +87,25 @@ class RequestMessage(HTTPMessage):
     """Base class for HTTP request messages"""
 
     endpoint_url: str
-    method: Literal["GET", "POST", "PATCH", "DELETE"]
+    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_responses(
-        self,
-        response: httpx.Response,
-        response_body: dict[str, JsonVal] | None = None,
-        error_message: str | None = None,
-    ) -> Sequence[HTTPMessage]:
+    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
@@ -69,13 +115,25 @@ class RequestMessage(HTTPMessage):
 
 @dataclass
 class SuccessResponse(ResponseMessage):
-    body: dict[str, JsonVal] | None = None
+    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):
-    error: str
-    body: dict[str, JsonVal] | None = None
+    body: str
+    error: ErrorDetails
+
+    def dump(self) -> dict[str, JsonVal]:
+        output = super().dump()
+        output["error"] = self.error.dump()
+        return output
 
 
 @dataclass
@@ -83,17 +141,16 @@ 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)]
+    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]:
@@ -105,7 +162,7 @@ class BodyRequest(RequestMessage, ABC):
     """Base class for HTTP request messages with a body"""
 
     @abstractmethod
-    def body(self) -> dict[str, JsonVal]:
+    def data(self) -> str:
         raise NotImplementedError()
 
 
@@ -113,87 +170,75 @@ class BodyRequest(RequestMessage, ABC):
 class ParamRequest(SimpleRequest):
     """Base class for HTTP request messages with query parameters"""
 
-    parameters: dict[str, str] | None = None
+    parameters: dict[str, PrimitiveType] | 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
+    def data(self) -> str:
+        return _dump_body(self.body_content)
 
 
 @dataclass
-class ItemMessage(ABC):
-    """Base class for message related to a specific item"""
-
-    ...
+class DataBodyRequest(SimpleRequest):
+    data_content: bytes = b""
 
+    def data(self) -> bytes:
+        return self.data_content
 
-@dataclass
-class ItemIDMessage(Generic[T_ID], ItemMessage, ABC):
-    """Base class for message related to a specific item identified by an ID"""
-
-    id: T_ID
+    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
 
 
-@dataclass
-class ItemResponse(ItemIDMessage, ResponseMessage, ABC): ...
+T_COVARIANT_ID = TypeVar("T_COVARIANT_ID", covariant=True)
 
 
 @dataclass
-class SuccessItem(ItemResponse):
-    item: JsonVal | None = None
-
+class ItemMessage(Generic[T_COVARIANT_ID], ABC):
+    """Base class for message related to a specific item identified by an ID"""
 
-@dataclass
-class FailedItem(ItemResponse):
-    error: str
+    ids: list[T_COVARIANT_ID] = field(default_factory=list)
 
 
 @dataclass
-class MissingItem(ItemResponse): ...
+class SuccessResponseItems(ItemMessage[T_COVARIANT_ID], SuccessResponse): ...
 
 
 @dataclass
-class UnexpectedItem(ItemResponse):
-    item: JsonVal | None = None
+class FailedResponseItems(ItemMessage[T_COVARIANT_ID], FailedResponse): ...
 
 
 @dataclass
-class FailedRequestItem(ItemIDMessage, FailedRequestMessage): ...
+class FailedRequestItems(ItemMessage[T_COVARIANT_ID], FailedRequestMessage): ...
 
 
-@dataclass
-class UnknownRequestItem(ItemMessage, FailedRequestMessage):
-    item: JsonVal | None = None
-
+class RequestItem(Generic[T_COVARIANT_ID], Protocol):
+    def dump(self) -> JsonVal: ...
 
-@dataclass
-class UnknownResponseItem(ItemMessage, ResponseMessage):
-    error: str
-    item: JsonVal | None = None
+    def as_id(self) -> T_COVARIANT_ID: ...
 
 
 @dataclass
-class ItemsRequest(Generic[T_ID], BodyRequest):
+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[JsonVal]): The list of items to be sent in the request body.
+        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
-        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)
+    items: list[RequestItem[T_COVARIANT_ID]] = 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)
 
@@ -206,18 +251,27 @@ class ItemsRequest(Generic[T_ID], BodyRequest):
             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"]
+        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.items, **self.extra_body_fields}
-        return {"items": self.items}
+            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.
@@ -239,23 +293,21 @@ class ItemsRequest(Generic[T_ID], BodyRequest):
             return [self]
         tracker = self.tracker or ItemsRequestTracker(self.max_failures_before_abort)
         tracker.register_failure()
-        first_half = ItemsRequest[T_ID](
+        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,
-            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](
+        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,
-            as_id=self.as_id,
             connect_attempt=self.connect_attempt,
             read_attempt=self.read_attempt,
             status_attempt=status_attempts,
@@ -263,126 +315,98 @@ class ItemsRequest(Generic[T_ID], BodyRequest):
         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.
+    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)]
 
-        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.
+    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:
-            A sequence of HTTPMessage instances representing the outcome for each item in the request.
+            bool: True if there are any failed responses or requests, False otherwise.
         """
-        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}"
+        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
                     )
                 )
-                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))
+            elif isinstance(message, FailedRequestMessage):
+                results.append(FailedRequestItems(ids=[item_id], error=message.error))
             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())
+                results.append(message)
         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
+
+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