cognite-toolkit 0.7.28__tar.gz → 0.7.29__tar.gz

{cognite_toolkit-0.7.28 → cognite_toolkit-0.7.29}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cognite_toolkit
-Version: 0.7.28
+Version: 0.7.29
 Summary: Official Cognite Data Fusion tool for project templates and configuration deployment
 Author: Cognite AS
 Author-email: Cognite AS <support@cognite.com>

{cognite_toolkit-0.7.28 → cognite_toolkit-0.7.29}/cognite_toolkit/_cdf_tk/client/_toolkit_client.py

@@ -35,19 +35,26 @@ class ToolAPI:
 
 
 class ToolkitClient(CogniteClient):
-    def __init__(self, config: ToolkitClientConfig | None = None, enable_set_pending_ids: bool = False) -> None:
+    def __init__(
+        self,
+        config: ToolkitClientConfig | None = None,
+        enable_set_pending_ids: bool = False,
+        console: Console | None = None,
+    ) -> None:
         super().__init__(config=config)
         http_client = HTTPClient(self.config)
         self.http_client = http_client
         toolkit_config = ToolkitClientConfig.from_client_config(self.config)
-        self.console = Console()
+        self.console = console or Console()
         self.tool = ToolAPI(http_client, self.console)
         self.search = SearchAPI(self._config, self._API_VERSION, self)
         self.robotics = RoboticsAPI(self._config, self._API_VERSION, self)
         self.dml = DMLAPI(self._config, self._API_VERSION, self)
         self.verify = VerifyAPI(self._config, self._API_VERSION, self)
         self.lookup = LookUpGroup(self._config, self._API_VERSION, self, self.console)
-        self.functions: ExtendedFunctionsAPI = ExtendedFunctionsAPI(toolkit_config, self._API_VERSION, self)
+        self.functions: ExtendedFunctionsAPI = ExtendedFunctionsAPI(
+            toolkit_config, self._API_VERSION, self, self.console
+        )
         self.data_modeling: ExtendedDataModelingAPI = ExtendedDataModelingAPI(self._config, self._API_VERSION, self)
         if enable_set_pending_ids:
             self.time_series: ExtendedTimeSeriesAPI = ExtendedTimeSeriesAPI(self._config, self._API_VERSION, self)

{cognite_toolkit-0.7.28 → cognite_toolkit-0.7.29}/cognite_toolkit/_cdf_tk/client/api/extended_functions.py

@@ -17,15 +17,16 @@ class ExtendedFunctionsAPI(FunctionsAPI):
         config: ToolkitClientConfig,
         api_version: str | None,
         cognite_client: CogniteClient,
+        console: Console | None = None,
     ) -> None:
         """
         Extended Functions API to include custom headers and payload preparation.
         """
         super().__init__(config, api_version, cognite_client)
         self._toolkit_config = config
-        self._toolkit_http_client = HTTPClient(config, max_retries=global_config.max_retries)
+        self._toolkit_http_client = HTTPClient(config, max_retries=global_config.max_retries, console=console)
 
-    def create_with_429_retry(self, function: FunctionWrite, console: Console | None = None) -> Function:
+    def create_with_429_retry(self, function: FunctionWrite) -> Function:
         """Create a function with manual retry handling for 429 Too Many Requests responses.
 
         This method is a workaround for scenarios where the function creation API is temporarily unavailable
@@ -43,16 +44,13 @@ class ExtendedFunctionsAPI(FunctionsAPI):
                 endpoint_url=self._toolkit_config.create_api_url("/functions"),
                 method="POST",
                 body_content={"items": [function.dump(camel_case=True)]},
-            ),
-            console=console,
+            )
         )
         result.raise_for_status()
         # We assume the API response is one item on a successful creation
         return Function._load(result.get_first_body()["items"][0], cognite_client=self._cognite_client)  # type: ignore[arg-type,index]
 
-    def delete_with_429_retry(
-        self, external_id: SequenceNotStr[str], ignore_unknown_ids: bool = False, console: Console | None = None
-    ) -> None:
+    def delete_with_429_retry(self, external_id: SequenceNotStr[str], ignore_unknown_ids: bool = False) -> None:
         """Delete one or more functions with retry handling for 429 Too Many Requests responses.
 
         This method is an improvement over the standard delete method in the FunctionsAPI.
@@ -77,7 +75,6 @@ class ExtendedFunctionsAPI(FunctionsAPI):
                     endpoint_url=self._toolkit_config.create_api_url("/functions/delete"),
                     method="POST",
                     body_content=body_content,
-                ),
-                console=console,
+                )
             ).raise_for_status()
         return None

{cognite_toolkit-0.7.28 → cognite_toolkit-0.7.29}/cognite_toolkit/_cdf_tk/commands/_migrate/migration_io.py

@@ -231,6 +231,8 @@ class AnnotationMigrationIO(
     CHUNK_SIZE = 1000
     UPLOAD_ENDPOINT = InstanceIO.UPLOAD_ENDPOINT
 
+    SUPPORTED_ANNOTATION_TYPES = frozenset({"diagrams.AssetLink", "diagrams.FileLink"})
+
     def __init__(
         self,
         client: ToolkitClient,
@@ -272,6 +274,10 @@ class AnnotationMigrationIO(
         for data_chunk in self.annotation_io.stream_data(asset_centric_selector, limit):
             mapping_list = AssetCentricMappingList[Annotation]([])
             for resource in data_chunk.items:
+                if resource.annotation_type not in self.SUPPORTED_ANNOTATION_TYPES:
+                    # This should not happen, as the annotation_io should already filter these out.
+                    # This is just in case.
+                    continue
                 mapping = AnnotationMapping(
                     instance_id=EdgeId(space=self.instance_space, external_id=f"annotation_{resource.id!r}"),
                     id=resource.id,
@@ -294,11 +300,15 @@ class AnnotationMigrationIO(
             resources = self.client.annotations.retrieve_multiple(current_batch.get_ids())
             resources_by_id = {resource.id: resource for resource in resources}
             not_found = 0
+            incorrect_type_count = 0
             for mapping in current_batch:
                 resource = resources_by_id.get(mapping.id)
                 if resource is None:
                     not_found += 1
                     continue
+                if resource.annotation_type not in self.SUPPORTED_ANNOTATION_TYPES:
+                    incorrect_type_count += 1
+                    continue
                 mapping.ingestion_view = self._get_mapping(mapping.ingestion_view, resource)
                 chunk.append(AssetCentricMapping(mapping=mapping, resource=resource))
             if chunk:
@@ -308,6 +318,11 @@ class AnnotationMigrationIO(
                 MediumSeverityWarning(
                     f"Could not find {not_found} annotations referenced in the CSV file. They will be skipped during migration."
                 ).print_warning(include_timestamp=True, console=self.client.console)
+            if incorrect_type_count:
+                MediumSeverityWarning(
+                    f"Found {incorrect_type_count} annotations with unsupported types. Only 'diagrams.AssetLink' and "
+                    "'diagrams.FileLink' are supported. These annotations will be skipped during migration."
+                ).print_warning(include_timestamp=True, console=self.client.console)
 
     def _get_mapping(self, current_mapping: str | None, resource: Annotation) -> str:
         try:

{cognite_toolkit-0.7.28 → cognite_toolkit-0.7.29}/cognite_toolkit/_cdf_tk/cruds/_resource_cruds/function.py

@@ -320,7 +320,7 @@ class FunctionCRUD(ResourceCRUD[str, FunctionWrite, Function]):
                     " problem persists, please contact Cognite support."
                 )
             item.file_id = file_id
-            created_item = self.client.functions.create_with_429_retry(item, console=self.console)
+            created_item = self.client.functions.create_with_429_retry(item)
             self._warn_if_cpu_or_memory_changed(created_item, item)
             created.append(created_item)
         return created
@@ -404,7 +404,7 @@ class FunctionCRUD(ResourceCRUD[str, FunctionWrite, Function]):
     def delete(self, ids: SequenceNotStr[str]) -> int:
         functions = self.retrieve(ids)
 
-        self.client.functions.delete_with_429_retry(external_id=ids, ignore_unknown_ids=True, console=self.console)
+        self.client.functions.delete_with_429_retry(external_id=ids, ignore_unknown_ids=True)
         file_ids = {func.file_id for func in functions if func.file_id}
         self.client.files.delete(id=list(file_ids), ignore_unknown_ids=True)
         return len(ids)

{cognite_toolkit-0.7.28 → cognite_toolkit-0.7.29}/cognite_toolkit/_cdf_tk/storageio/_annotations.py

@@ -33,7 +33,8 @@ class AnnotationIO(StorageIO[AssetCentricSelector, Annotation]):
                         annotated_resource_type="file",
                         annotated_resource_ids=[{"id": file_metadata.id} for file_metadata in file_chunk.items],
                         annotation_type=annotation_type,
-                    )
+                    ),
+                    limit=-1,
                 )
                 if limit is not None and total + len(results) > limit:
                     results = results[: limit - total]

{cognite_toolkit-0.7.28 → cognite_toolkit-0.7.29}/cognite_toolkit/_cdf_tk/utils/http_client/_client.py

@@ -23,6 +23,15 @@ from cognite_toolkit._cdf_tk.utils.http_client._data_classes import (
     ResponseList,
     ResponseMessage,
 )
+from cognite_toolkit._cdf_tk.utils.http_client._data_classes2 import (
+    BaseRequestMessage,
+    ErrorDetails2,
+    FailedRequest2,
+    FailedResponse2,
+    HTTPResult2,
+    RequestMessage2,
+    SuccessResponse2,
+)
 from cognite_toolkit._cdf_tk.utils.useful_types import PrimitiveType
 
 if sys.version_info >= (3, 11):
@@ -48,6 +57,7 @@ class HTTPClient:
             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.
+        console (Console | None): Optional Rich Console for printing warnings.
 
     """
 
@@ -59,6 +69,7 @@ class HTTPClient:
         pool_maxsize: int = 20,
         retry_status_codes: Set[int] = frozenset({408, 429, 502, 503, 504}),
         split_items_status_codes: Set[int] = frozenset({400, 408, 409, 422, 502, 503, 504}),
+        console: Console | None = None,
     ):
         self.config = config
         self._max_retries = max_retries
@@ -66,6 +77,7 @@ class HTTPClient:
         self._pool_maxsize = pool_maxsize
         self._retry_status_codes = retry_status_codes
         self._split_items_status_codes = split_items_status_codes
+        self._console = console
 
         # Thread-safe session for connection pooling
         self.session = self._create_thread_safe_session()
@@ -80,12 +92,11 @@ class HTTPClient:
         self.session.close()
         return False  # Do not suppress exceptions
 
-    def request(self, message: RequestMessage, console: Console | None = None) -> Sequence[HTTPMessage]:
+    def request(self, message: RequestMessage) -> Sequence[HTTPMessage]:
         """Send an HTTP request and return the response.
 
         Args:
             message (RequestMessage): The request message to send.
-            console (Console | None): The rich console to use for printing warnings.
 
         Returns:
             Sequence[HTTPMessage]: The response message(s). This can also
@@ -98,12 +109,12 @@ class HTTPClient:
             return message.create_failed_request(error_msg)
         try:
             response = self._make_request(message)
-            results = self._handle_response(response, message, console)
+            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, console: Console | None = None) -> ResponseList:
+    def request_with_retries(self, message: RequestMessage) -> ResponseList:
         """Send an HTTP request and handle retries.
 
         This method will keep retrying the request until it either succeeds or
@@ -114,7 +125,6 @@ class HTTPClient:
 
         Args:
             message (RequestMessage): The request message to send.
-            console (Console | None): The rich console to use for printing warnings.
 
         Returns:
             Sequence[ResponseMessage | FailedRequestMessage]: The final response
@@ -127,7 +137,7 @@ class HTTPClient:
         final_responses = ResponseList([])
         while pending_requests:
             current_request = pending_requests.popleft()
-            results = self.request(current_request, console)
+            results = self.request(current_request)
 
             for result in results:
                 if isinstance(result, RequestMessage):
@@ -194,9 +204,7 @@ class HTTPClient:
             follow_redirects=False,
         )
 
-    def _handle_response(
-        self, response: httpx.Response, request: RequestMessage, console: Console | None = None
-    ) -> Sequence[HTTPMessage]:
+    def _handle_response(self, response: httpx.Response, request: RequestMessage) -> Sequence[HTTPMessage]:
         if 200 <= response.status_code < 300:
             return request.create_success_response(response)
         elif (
@@ -216,11 +224,11 @@ class HTTPClient:
 
         retry_after = self._get_retry_after_in_header(response)
         if retry_after is not None and response.status_code == 429 and request.status_attempt < self._max_retries:
-            if console is not None:
+            if self._console is not None:
                 short_url = request.endpoint_url.removeprefix(self.config.base_api_url)
                 HighSeverityWarning(
                     f"Rate limit exceeded for the {short_url!r} endpoint. Retrying after {retry_after} seconds."
-                ).print_warning(console=console)
+                ).print_warning(console=self._console)
             request.status_attempt += 1
             time.sleep(retry_after)
             return [request]
@@ -273,3 +281,110 @@ class HTTPClient:
             error_msg = f"RequestException after {request.total_attempts - 1} attempts ({error_type} error): {e!s}"
 
             return request.create_failed_request(error_msg)
+
+    def request_single(self, message: RequestMessage2) -> RequestMessage2 | HTTPResult2:
+        """Send an HTTP request and return the response.
+
+        Args:
+            message (RequestMessage2): The request message to send.
+        Returns:
+            HTTPMessage: The response message.
+        """
+        try:
+            response = self._make_request2(message)
+            result = self._handle_response_single(response, message)
+        except Exception as e:
+            result = self._handle_error_single(e, message)
+        return result
+
+    def request_single_retries(self, message: RequestMessage2) -> HTTPResult2:
+        """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 (RequestMessage2): The request message to send.
+        Returns:
+            HTTPMessage2: The final response message, which can be either successful response or failed request.
+        """
+        if message.total_attempts > 0:
+            raise RuntimeError(f"RequestMessage has already been attempted {message.total_attempts} times.")
+        current_request = message
+        while True:
+            result = self.request_single(current_request)
+            if isinstance(result, RequestMessage2):
+                current_request = result
+            elif isinstance(result, HTTPResult2):
+                return result
+            else:
+                raise TypeError(f"Unexpected result type: {type(result)}")
+
+    def _make_request2(self, message: BaseRequestMessage) -> httpx.Response:
+        headers = self._create_headers(message.api_version, message.content_type, message.accept)
+        return self.session.request(
+            method=message.method,
+            url=message.endpoint_url,
+            content=message.content,
+            headers=headers,
+            params=message.parameters,
+            timeout=self.config.timeout,
+            follow_redirects=False,
+        )
+
+    def _handle_response_single(
+        self, response: httpx.Response, request: RequestMessage2
+    ) -> RequestMessage2 | HTTPResult2:
+        if 200 <= response.status_code < 300:
+            return SuccessResponse2(
+                status_code=response.status_code,
+                body=response.text,
+                content=response.content,
+            )
+        retry_after = self._get_retry_after_in_header(response)
+        if retry_after is not None and response.status_code == 429 and request.status_attempt < self._max_retries:
+            if self._console is not None:
+                short_url = request.endpoint_url.removeprefix(self.config.base_api_url)
+                HighSeverityWarning(
+                    f"Rate limit exceeded for the {short_url!r} endpoint. Retrying after {retry_after} seconds."
+                ).print_warning(console=self._console)
+            request.status_attempt += 1
+            time.sleep(retry_after)
+            return request
+
+        if request.status_attempt < self._max_retries and response.status_code in self._retry_status_codes:
+            request.status_attempt += 1
+            time.sleep(self._backoff_time(request.total_attempts))
+            return request
+        else:
+            # Permanent failure
+            return FailedResponse2(
+                status_code=response.status_code,
+                body=response.text,
+                error=ErrorDetails2.from_response(response),
+            )
+
+    def _handle_error_single(self, e: Exception, request: RequestMessage2) -> RequestMessage2 | HTTPResult2:
+        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 FailedRequest2(error=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 FailedRequest2(error=error_msg)

cognite_toolkit-0.7.29/cognite_toolkit/_cdf_tk/utils/http_client/_data_classes2.py

@@ -0,0 +1,97 @@
+import gzip
+from abc import ABC, abstractmethod
+from typing import Any, Literal
+
+import httpx
+from cognite.client import global_config
+from pydantic import BaseModel, JsonValue, TypeAdapter, model_validator
+
+from cognite_toolkit._cdf_tk.utils.useful_types import PrimitiveType
+
+
+class HTTPResult2(BaseModel): ...
+
+
+class FailedRequest2(HTTPResult2):
+    error: str
+
+
+class SuccessResponse2(HTTPResult2):
+    status_code: int
+    body: str
+    content: bytes
+
+
+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, str):
+                data = gzip.compress(data.encode("utf-8"))
+        return data
+
+
+_BODY_SERIALIZER = TypeAdapter(dict[str, JsonValue])

{cognite_toolkit-0.7.28 → cognite_toolkit-0.7.29}/cognite_toolkit/_repo_files/GitHub/.github/workflows/deploy.yaml

@@ -12,7 +12,7 @@ jobs:
     environment: dev
     name: Deploy
     container:
-      image: cognite/toolkit:0.7.28
+      image: cognite/toolkit:0.7.29
       env:
         CDF_CLUSTER: ${{ vars.CDF_CLUSTER }}
         CDF_PROJECT: ${{ vars.CDF_PROJECT }}

{cognite_toolkit-0.7.28 → cognite_toolkit-0.7.29}/cognite_toolkit/_repo_files/GitHub/.github/workflows/dry-run.yaml

@@ -10,7 +10,7 @@ jobs:
     environment: dev
     name: Deploy Dry Run
     container:
-      image: cognite/toolkit:0.7.28
+      image: cognite/toolkit:0.7.29
       env:
         CDF_CLUSTER: ${{ vars.CDF_CLUSTER }}
         CDF_PROJECT: ${{ vars.CDF_PROJECT }}

{cognite_toolkit-0.7.28 → cognite_toolkit-0.7.29}/cognite_toolkit/_resources/cdf.toml

@@ -4,7 +4,7 @@ default_env = "<DEFAULT_ENV_PLACEHOLDER>"
 [modules]
 # This is the version of the modules. It should not be changed manually.
 # It will be updated by the 'cdf modules upgrade' command.
-version = "0.7.28"
+version = "0.7.29"
 
 
 [plugins]

cognite_toolkit-0.7.29/cognite_toolkit/_version.py

@@ -0,0 +1 @@
+__version__ = "0.7.29"

{cognite_toolkit-0.7.28 → cognite_toolkit-0.7.29}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "cognite_toolkit"
-version = "0.7.28"
+version = "0.7.29"
 description = "Official Cognite Data Fusion tool for project templates and configuration deployment"
 authors = [
     {name ="Cognite AS", email="support@cognite.com"}

cognite_toolkit-0.7.28/cognite_toolkit/_version.py

@@ -1 +0,0 @@
-__version__ = "0.7.28"