cognite-neat 0.123.24__py3-none-any.whl → 1.0.22__py3-none-any.whl

cognite/neat/_utils/http_client/_client.py

@@ -0,0 +1,245 @@
+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.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 (
+    APIResponse,
+    BodyRequest,
+    ErrorDetails,
+    FailedRequestMessage,
+    HTTPMessage,
+    ItemsRequest,
+    ParametersRequest,
+    RequestMessage,
+    ResponseMessage,
+)
+from cognite.neat._utils.useful_types import PrimaryTypes
+
+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) -> APIResponse:
+        """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 = APIResponse()
+
+        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
+
+    def _make_request(self, item: RequestMessage) -> httpx.Response:
+        headers = self._create_headers(item.api_version)
+        params: dict[str, PrimaryTypes] | None = None
+        if isinstance(item, ParametersRequest):
+            params = item.parameters
+        data: str | bytes | None = None
+        if isinstance(item, BodyRequest):
+            data = item.data()
+            if not global_config.disable_gzip:
+                data = gzip.compress(data.encode("utf-8"))
+        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]:
+        if 200 <= response.status_code < 300:
+            return request.create_success_response(response)
+
+        if (
+            isinstance(request, ItemsRequest)
+            and len(request.body.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_failure_response(response)
+            return splits
+
+        error = ErrorDetails.from_response(response)
+
+        if request.status_attempt < self._max_retries and (
+            response.status_code in self._retry_status_codes or error.is_auto_retryable
+        ):
+            request.status_attempt += 1
+            time.sleep(self._backoff_time(request.total_attempts))
+            return [request]
+        else:
+            # Permanent failure
+            return request.create_failure_response(response)
+
+    @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,294 @@
+import sys
+from abc import ABC, abstractmethod
+from collections import UserList
+from collections.abc import MutableSequence, Sequence
+from typing import Generic, Literal, TypeAlias, TypeVar
+
+import httpx
+from pydantic import BaseModel, ConfigDict, Field, JsonValue, ValidationError, model_serializer
+
+from cognite.neat._exceptions import CDFAPIException
+from cognite.neat._utils.http_client._tracker import ItemsRequestTracker
+from cognite.neat._utils.useful_types import PrimaryTypes, ReferenceObject, T_Reference
+
+if sys.version_info >= (3, 11):
+    from typing import Self
+else:
+    from typing_extensions import Self
+
+StatusCode: TypeAlias = int
+
+
+class HTTPMessage(BaseModel):
+    """Base class for HTTP messages (requests and responses)"""
+
+
+class FailedRequestMessage(HTTPMessage):
+    message: str
+
+    def __str__(self) -> str:
+        return self.message
+
+
+class ResponseMessage(HTTPMessage):
+    code: StatusCode
+    body: str
+
+
+class SuccessResponse(ResponseMessage): ...
+
+
+class ErrorDetails(BaseModel):
+    """This is the structure of failure responses from CDF APIs"""
+
+    code: StatusCode
+    message: str
+    missing: list[JsonValue] | None = None
+    duplicated: list[JsonValue] | None = None
+    is_auto_retryable: bool | None = Field(None, alias="isAutoRetryable")
+
+    @classmethod
+    def from_response(cls, response: httpx.Response) -> "ErrorDetails":
+        try:
+            return _ErrorResponse.model_validate_json(response.text).error
+        except ValidationError:
+            return cls(code=response.status_code, message=response.text)
+
+
+class _ErrorResponse(BaseModel):
+    error: ErrorDetails
+
+
+class FailedResponse(ResponseMessage):
+    error: ErrorDetails
+
+    def __str__(self) -> str:
+        return f"HTTP {self.code} | {self.error.message}"
+
+
+class RequestMessage(HTTPMessage, ABC):
+    """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_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()
+
+
+class SimpleRequest(RequestMessage):
+    """Base class for requests with a simple success/fail response structure"""
+
+    def create_success_response(self, response: httpx.Response) -> Sequence[HTTPMessage]:
+        return [SuccessResponse(code=response.status_code, body=response.text)]
+
+    def create_failure_response(self, response: httpx.Response) -> Sequence[HTTPMessage]:
+        return [
+            FailedResponse(code=response.status_code, body=response.text, error=ErrorDetails.from_response(response))
+        ]
+
+    def create_failed_request(self, error_message: str) -> Sequence[HTTPMessage]:
+        return [FailedRequestMessage(message=error_message)]
+
+
+class ParametersRequest(SimpleRequest):
+    """Base class for HTTP request messages with query parameters"""
+
+    parameters: dict[str, PrimaryTypes] | None = None
+
+
+class BodyRequest(ParametersRequest, ABC):
+    """Base class for HTTP request messages with a body"""
+
+    @abstractmethod
+    def data(self) -> str:
+        raise NotImplementedError()
+
+
+class SimpleBodyRequest(BodyRequest):
+    body: str
+
+    def data(self) -> str:
+        return self.body
+
+
+class ItemMessage(BaseModel, Generic[T_Reference], ABC):
+    """Base class for message related to a specific item"""
+
+    ids: Sequence[T_Reference]
+
+
+class SuccessResponseItems(ItemMessage[T_Reference], SuccessResponse): ...
+
+
+class FailedResponseItems(ItemMessage[T_Reference], FailedResponse): ...
+
+
+class FailedRequestItems(ItemMessage[T_Reference], FailedRequestMessage): ...
+
+
+T_BaseModel = TypeVar("T_BaseModel", bound=BaseModel)
+
+
+class ItemBody(BaseModel, Generic[T_Reference, T_BaseModel], ABC):
+    items: Sequence[T_BaseModel]
+    extra_args: dict[str, JsonValue] | None = None
+
+    @model_serializer(mode="plain", return_type=dict)
+    def serialize(self) -> dict[str, JsonValue]:
+        data: dict[str, JsonValue] = {
+            "items": [item.model_dump(exclude_unset=False, by_alias=True, exclude_none=False) for item in self.items]
+        }
+        if isinstance(self.extra_args, dict):
+            data.update(self.extra_args)
+        return data
+
+    @abstractmethod
+    def as_ids(self) -> list[T_Reference]:
+        """Returns the list of item identifiers for the items in the body."""
+        raise NotImplementedError()
+
+    def split(self, mid: int) -> tuple[Self, Self]:
+        """Splits the body into two smaller bodies.
+
+        This is useful for retrying requests that fail due to size limits or timeouts.
+
+        Args:
+            mid: The index at which to split the items.
+        Returns:
+            A tuple containing two new ItemBody instances, each with half of the original items.
+
+        """
+
+        type_ = type(self)
+        return type_(items=self.items[:mid], extra_args=self.extra_args), type_(
+            items=self.items[mid:], extra_args=self.extra_args
+        )
+
+
+class ItemIDBody(ItemBody[ReferenceObject, ReferenceObject]):
+    def as_ids(self) -> list[ReferenceObject]:
+        return list(self.items)
+
+
+class ItemsRequest(BodyRequest, Generic[T_Reference, T_BaseModel]):
+    """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:
+        body (ItemBody): The body of the request containing the items to be processed.
+        max_failures_before_abort (int): The maximum number of failed split requests before aborting further splits.
+
+    """
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+    body: ItemBody[T_Reference, T_BaseModel]
+    max_failures_before_abort: int = 50
+    tracker: ItemsRequestTracker | None = None
+
+    def data(self) -> str:
+        return self.body.model_dump_json(exclude_unset=True, by_alias=True)
+
+    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.body.items) // 2
+        if mid == 0:
+            return [self]
+        tracker = self.tracker or ItemsRequestTracker(self.max_failures_before_abort)
+        tracker.register_failure()
+        messages: list[ItemsRequest] = []
+        for body in self.body.split(mid):
+            item_request = ItemsRequest(
+                endpoint_url=self.endpoint_url,
+                method=self.method,
+                body=body,
+                connect_attempt=self.connect_attempt,
+                read_attempt=self.read_attempt,
+                status_attempt=status_attempts,
+            )
+            item_request.tracker = tracker
+            messages.append(item_request)
+        return messages
+
+    def create_success_response(self, response: httpx.Response) -> Sequence[HTTPMessage]:
+        return [SuccessResponseItems(code=response.status_code, body=response.text, ids=self.body.as_ids())]
+
+    def create_failure_response(self, response: httpx.Response) -> Sequence[HTTPMessage]:
+        """Creates response messages based on the HTTP response and the original request.
+
+        Args:
+            response: The HTTP response received from the server.
+        Returns:
+            A sequence of HTTPMessage instances representing the outcome for each item in the request.
+        """
+        return [
+            FailedResponseItems(
+                code=response.status_code,
+                body=response.text,
+                error=ErrorDetails.from_response(response),
+                ids=self.body.as_ids(),
+            )
+        ]
+
+    def create_failed_request(self, error_message: str) -> Sequence[HTTPMessage]:
+        """Creates failed request messages for each item in the request.
+
+        Args:
+            error_message: The error message to include in the failed request messages.
+
+        Returns:
+            A sequence of HTTPMessage instances representing the failed request for each item.
+        """
+        return [FailedRequestItems(message=error_message, ids=self.body.as_ids())]
+
+
+class APIResponse(UserList, MutableSequence[ResponseMessage | FailedRequestMessage]):
+    def __init__(self, collection: Sequence[ResponseMessage | FailedRequestMessage] | None = None):
+        super().__init__(collection or [])
+
+    def raise_for_status(self) -> None:
+        error_messages = [message for message in self.data if not isinstance(message, SuccessResponse)]
+        if error_messages:
+            raise CDFAPIException(error_messages)
+
+    @property
+    def success_response(self) -> SuccessResponse:
+        success = [msg for msg in self.data if isinstance(msg, SuccessResponse)]
+        if len(success) == 1:
+            return success[0]
+        elif success:
+            raise ValueError("Multiple successful HTTP responses found in the messages.")
+        else:
+            raise ValueError("No successful HTTP response found in the messages.")

cognite/neat/_utils/http_client/_tracker.py

@@ -0,0 +1,31 @@
+import threading
+from dataclasses import dataclass, field
+
+
+@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 (threading.Lock): 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
+    lock: threading.Lock = 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

cognite/neat/_utils/repo.py

@@ -0,0 +1,19 @@
+import subprocess
+from pathlib import Path
+
+
+def get_repo_root() -> Path:
+    """Get the root path of the git repository.
+
+    Raises:
+        RuntimeError: If git is not installed or the current directory is not in a git repository
+
+    """
+    try:
+        result = subprocess.run("git rev-parse --show-toplevel".split(), stdout=subprocess.PIPE)
+    except FileNotFoundError as e:
+        raise RuntimeError("Git is not installed or not found in PATH") from e
+    output = result.stdout.decode().strip()
+    if not output:
+        raise RuntimeError("Not in a git repository")
+    return Path(output)

cognite/neat/_utils/text.py

@@ -0,0 +1,71 @@
+import re
+from collections.abc import Collection
+from typing import Any
+
+NEWLINE = "\n"
+TAB = "\t"
+
+
+def humanize_collection(collection: Collection[Any], /, *, sort: bool = True, bind_word: str = "and") -> str:
+    """Convert a collection of items to a human-readable string.
+
+    Args:
+        collection: The collection of items to convert.
+        sort: Whether to sort the collection before converting. Default is True.
+        bind_word: The word to use to bind the last two items. Default is "and".
+
+    Returns:
+        A human-readable string of the collection.
+
+    Examples:
+        >>> humanize_collection(["b", "c", "a"])
+        'a, b and c'
+        >>> humanize_collection(["b", "c", "a"], sort=False)
+        'b, c and a'
+        >>> humanize_collection(["a", "b"])
+        'a and b'
+        >>> humanize_collection(["a"])
+        'a'
+        >>> humanize_collection([])
+        ''
+
+    """
+    if not collection:
+        return ""
+    elif len(collection) == 1:
+        return str(next(iter(collection)))
+
+    strings = (str(item) for item in collection)
+    if sort:
+        sequence = sorted(strings)
+    else:
+        sequence = list(strings)
+
+    return f"{', '.join(sequence[:-1])} {bind_word} {sequence[-1]}"
+
+
+def title_case(s: str) -> str:
+    """Convert a string to title case, handling underscores and hyphens.
+
+    Args:
+        s: The string to convert.
+    Returns:
+        The title-cased string.
+    Examples:
+        >>> title_case("hello world")
+        'Hello World'
+        >>> title_case("hello_world")
+        'Hello World'
+        >>> title_case("hello-world")
+        'Hello World'
+        >>> title_case("hello_world-and-universe")
+        'Hello World And Universe'
+        >>> title_case("HELLO WORLD")
+        'Hello World'
+    """
+    return " ".join(word.capitalize() for word in s.replace("_", " ").replace("-", " ").split())
+
+
+def split_on_capitals(text: str) -> list[str]:
+    """Split a string at capital letters."""
+    return re.findall(r"[A-Z][a-z]*", text)