cognite-toolkit 0.7.54__py3-none-any.whl → 0.7.56__py3-none-any.whl

cognite_toolkit/_cdf_tk/storageio/_datapoints.py

@@ -13,15 +13,16 @@ from cognite.client._proto.data_points_pb2 import (
 from cognite.client.data_classes import TimeSeriesFilter
 from cognite.client.data_classes.filters import Exists
 from cognite.client.data_classes.time_series import TimeSeriesProperty
+from pydantic import ConfigDict
 
 from cognite_toolkit._cdf_tk.client import ToolkitClient
+from cognite_toolkit._cdf_tk.client._resource_base import Identifier, RequestResource
 from cognite_toolkit._cdf_tk.client.http_client import (
-    DataBodyRequest,
     HTTPClient,
-    HTTPMessage,
-    SimpleBodyRequest,
-    SuccessResponse,
+    RequestMessage2,
+    SuccessResponse2,
 )
+from cognite_toolkit._cdf_tk.client.http_client._item_classes import ItemsResultList
 from cognite_toolkit._cdf_tk.exceptions import ToolkitNotImplementedError
 from cognite_toolkit._cdf_tk.tk_warnings import HighSeverityWarning
 from cognite_toolkit._cdf_tk.utils import humanize_collection
@@ -39,9 +40,22 @@ from ._base import Page, TableStorageIO, TableUploadableStorageIO, UploadItem
 from .selectors import DataPointsDataSetSelector, DataPointsFileSelector, DataPointsSelector
 
 
+class DatapointsRequestAdapter(RequestResource):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+    datapoints: DataPointInsertionRequest
+
+    def dump(self, camel_case: bool = True, exclude_extra: bool = False) -> dict[str, Any]:
+        return {"datapoints": self.datapoints.SerializeToString()}
+
+    def as_id(self) -> Identifier:
+        raise NotImplementedError(
+            "DatapointsRequestAdapter does not have an identifier. - it wraps multiple timeseries"
+        )
+
+
 class DatapointsIO(
     TableStorageIO[DataPointsSelector, DataPointListResponse],
-    TableUploadableStorageIO[DataPointsSelector, DataPointListResponse, DataPointInsertionRequest],
+    TableUploadableStorageIO[DataPointsSelector, DataPointListResponse, DatapointsRequestAdapter],
 ):
     SUPPORTED_DOWNLOAD_FORMATS = frozenset({".csv"})
     SUPPORTED_COMPRESSIONS = frozenset({".gz"})
@@ -117,8 +131,8 @@ class DatapointsIO(
                 }
                 for ts in timeseries
             ]
-            responses = self.client.http_client.request_with_retries(
-                SimpleBodyRequest(
+            response = self.client.http_client.request_single_retries(
+                RequestMessage2(
                     endpoint_url=config.create_api_url("/timeseries/data/list"),
                     method="POST",
                     accept="application/protobuf",
@@ -126,10 +140,9 @@ class DatapointsIO(
                     body_content={"items": items},  # type: ignore[dict-item]
                 )
             )
-            first_success = next((resp for resp in responses if isinstance(resp, SuccessResponse)), None)
-            if first_success is None:
+            if not isinstance(response, SuccessResponse2):
                 continue
-            aggregate_response: DataPointListResponse = DataPointListResponse.FromString(first_success.content)
+            aggregate_response: DataPointListResponse = DataPointListResponse.FromString(response.content)
             timeseries_ids_with_data: dict[int, int] = {}
             for dp in aggregate_response.items:
                 if dp.aggregateDatapoints.datapoints:
@@ -176,8 +189,8 @@ class DatapointsIO(
                 yield page
 
     def _fetch_datapoints_batch(self, batch: list[dict[str, Any]], config: Any) -> Page[DataPointListResponse] | None:
-        responses = self.client.http_client.request_with_retries(
-            SimpleBodyRequest(
+        response = self.client.http_client.request_single_retries(
+            RequestMessage2(
                 endpoint_url=config.create_api_url("/timeseries/data/list"),
                 method="POST",
                 accept="application/protobuf",
@@ -185,10 +198,9 @@ class DatapointsIO(
                 body_content={"items": batch},  # type: ignore[dict-item]
             )
         )
-        first_success = next((resp for resp in responses if isinstance(resp, SuccessResponse)), None)
-        if first_success is None:
+        if not isinstance(response, SuccessResponse2):
             return None
-        data_response: DataPointListResponse = DataPointListResponse.FromString(first_success.content)
+        data_response: DataPointListResponse = DataPointListResponse.FromString(response.content)
         return Page("Main", [data_response])
 
     def count(self, selector: DataPointsSelector) -> int | None:
@@ -238,26 +250,26 @@ class DatapointsIO(
 
     def upload_items(
         self,
-        data_chunk: Sequence[UploadItem[DataPointInsertionRequest]],
+        data_chunk: Sequence[UploadItem[DatapointsRequestAdapter]],
         http_client: HTTPClient,
         selector: DataPointsSelector | None = None,
-    ) -> Sequence[HTTPMessage]:
-        results: list[HTTPMessage] = []
+    ) -> ItemsResultList:
+        results = ItemsResultList()
         for item in data_chunk:
-            response = http_client.request_with_retries(
-                DataBodyRequest(
+            response = http_client.request_single_retries(
+                RequestMessage2(
                     endpoint_url=http_client.config.create_api_url(self.UPLOAD_ENDPOINT),
                     method="POST",
                     content_type="application/protobuf",
-                    data_content=item.item.SerializeToString(),
+                    data_content=item.item.datapoints.SerializeToString(),
                 )
             )
-            results.extend(response)
+            results.append(response.as_item_response(item.source_id))
         return results
 
     def row_to_resource(
         self, source_id: str, row: dict[str, JsonVal], selector: DataPointsSelector | None = None
-    ) -> DataPointInsertionRequest:
+    ) -> DatapointsRequestAdapter:
         if selector is None:
             raise ValueError("Selector must be provided to convert row to DataPointInsertionItem.")
         # We assume that the row was read using the read_chunks method.
@@ -270,7 +282,7 @@ class DatapointsIO(
             raise RuntimeError(
                 f"Unsupported selector type {type(selector).__name__} for {type(self).__name__}. Trying to transform {source_id!r} from rows to DataPointInsertionRequest."
             )
-        return DataPointInsertionRequest(items=datapoints_items)
+        return DatapointsRequestAdapter(datapoints=DataPointInsertionRequest(items=datapoints_items))
 
     def _rows_to_datapoint_items_file_selector(
         self, rows: dict[str, list[Any]], selector: DataPointsFileSelector, source_id: str
@@ -406,7 +418,7 @@ class DatapointsIO(
             ).print_warning(console=self.client.console)
             self._warned_columns.add(column)
 
-    def json_to_resource(self, item_json: dict[str, JsonVal]) -> DataPointInsertionRequest:
+    def json_to_resource(self, item_json: dict[str, JsonVal]) -> DatapointsRequestAdapter:
         raise ToolkitNotImplementedError(
             f"Upload of {type(DatapointsIO).__name__.removesuffix('IO')} does not support json format."
         )

cognite_toolkit/_cdf_tk/storageio/_file_content.py

@@ -1,25 +1,23 @@
 import json
 import mimetypes
-from collections.abc import Iterable, MutableSequence, Sequence
+from collections.abc import Iterable, Sequence
 from dataclasses import dataclass
 from pathlib import Path
-from typing import cast
+from typing import Any, cast
 
 import httpx
 from cognite.client.data_classes.data_modeling import ViewId
 
 from cognite_toolkit._cdf_tk.client import ToolkitClient
 from cognite_toolkit._cdf_tk.client.http_client import (
-    ErrorDetails,
+    ErrorDetails2,
     FailedResponse2,
-    FailedResponseItems,
     HTTPClient,
-    HTTPMessage,
     HTTPResult2,
     RequestMessage2,
-    SimpleBodyRequest,
     SuccessResponse2,
 )
+from cognite_toolkit._cdf_tk.client.http_client._item_classes import ItemsFailedResponse2, ItemsResultList
 from cognite_toolkit._cdf_tk.client.resource_classes.data_modeling import NodeReference
 from cognite_toolkit._cdf_tk.client.resource_classes.filemetadata import FileMetadataRequest, FileMetadataResponse
 from cognite_toolkit._cdf_tk.cruds import FileMetadataCRUD
@@ -46,13 +44,12 @@ from .selectors._file_content import NodeId as SelectorNodeId
 COGNITE_FILE_VIEW = ViewId("cdf_cdm", "CogniteFile", "v1")
 
 
-@dataclass
 class UploadFileContentItem(UploadItem[FileMetadataRequest]):
     file_path: Path
     mime_type: str
 
-    def dump(self) -> JsonVal:
-        return self.item.dump(camel_case=True, exclude_extra=True)
+    def dump(self, camel_case: bool = True, exclude_extra: bool = True) -> dict[str, Any]:
+        return self.item.dump(camel_case=camel_case, exclude_extra=exclude_extra)
 
 
 @dataclass
@@ -123,8 +120,8 @@ class FileContentIO(UploadableStorageIO[FileContentSelector, MetadataWithFilePat
 
     def _retrieve_metadata(self, identifiers: Sequence[FileIdentifier]) -> Sequence[FileMetadataResponse] | None:
         config = self.client.config
-        responses = self.client.http_client.request_with_retries(
-            message=SimpleBodyRequest(
+        response = self.client.http_client.request_single_retries(
+            message=RequestMessage2(
                 endpoint_url=config.create_api_url("/files/byids"),
                 method="POST",
                 body_content={
@@ -136,9 +133,13 @@ class FileContentIO(UploadableStorageIO[FileContentSelector, MetadataWithFilePat
                 },
             )
         )
-        if responses.has_failed:
+        if not isinstance(response, SuccessResponse2):
+            return None
+        try:
+            body = response.body_json
+        except ValueError:
             return None
-        body = responses.get_first_body()
+
         items_data = body.get("items", [])
         if not isinstance(items_data, list):
             return None
@@ -182,20 +183,24 @@ class FileContentIO(UploadableStorageIO[FileContentSelector, MetadataWithFilePat
 
     def _retrieve_download_url(self, identifier: FileIdentifier) -> str | None:
         config = self.client.config
-        responses = self.client.http_client.request_with_retries(
-            message=SimpleBodyRequest(
+        response = self.client.http_client.request_single_retries(
+            message=RequestMessage2(
                 endpoint_url=config.create_api_url("/files/downloadlink"),
                 method="POST",
                 body_content={"items": [identifier.model_dump(mode="json", by_alias=True, exclude={"id_type"})]},
             )
         )
+        if not isinstance(response, SuccessResponse2):
+            return None
+
         try:
-            body = responses.get_first_body()
+            body = response.body_json
         except ValueError:
             return None
+
         if "items" in body and isinstance(body["items"], list) and len(body["items"]) > 0:
             # The API responses is not following the API docs, this is a workaround
-            body = body["items"][0]  # type: ignore[assignment]
+            body = body["items"][0]
         try:
             return cast(str, body["downloadUrl"])
         except (KeyError, IndexError):
@@ -257,8 +262,8 @@ class FileContentIO(UploadableStorageIO[FileContentSelector, MetadataWithFilePat
         data_chunk: Sequence[UploadItem[FileMetadataRequest]],
         http_client: HTTPClient,
         selector: FileContentSelector | None = None,
-    ) -> Sequence[HTTPMessage]:
-        results: MutableSequence[HTTPMessage] = []
+    ) -> ItemsResultList:
+        results = ItemsResultList()
         if isinstance(selector, FileMetadataTemplateSelector | FileIdentifierSelector):
             upload_url_getter = self._upload_url_asset_centric
         elif isinstance(selector, FileDataModelingTemplateSelector):
@@ -284,18 +289,17 @@ class FileContentIO(UploadableStorageIO[FileContentSelector, MetadataWithFilePat
                     content_length=len(content_bytes),
                 )
             )
-            results.append(upload_response.as_item_response(item.as_id()))
+            results.append(upload_response.as_item_response(str(item)))
         return results
 
     def _upload_url_asset_centric(
-        self, item: UploadFileContentItem, http_client: HTTPClient, results: MutableSequence[HTTPMessage]
+        self, item: UploadFileContentItem, http_client: HTTPClient, results: ItemsResultList
     ) -> str | None:
         response = http_client.request_single_retries(
             message=RequestMessage2(
                 endpoint_url=http_client.config.create_api_url(self.UPLOAD_ENDPOINT),
                 method="POST",
-                # MyPy does not understand that .dump is valid json
-                body_content=item.dump(),  # type: ignore[arg-type]
+                body_content=item.dump(),
             )
         )
         return self._parse_upload_link_response(response, item, results)
@@ -304,7 +308,7 @@ class FileContentIO(UploadableStorageIO[FileContentSelector, MetadataWithFilePat
         self,
         item: UploadFileContentItem,
         http_client: HTTPClient,
-        results: MutableSequence[HTTPMessage],
+        results: ItemsResultList,
         created_node: bool = False,
     ) -> str | None:
         """Get upload URL for data modeling file upload.
@@ -333,7 +337,7 @@ class FileContentIO(UploadableStorageIO[FileContentSelector, MetadataWithFilePat
             )
         )
         if isinstance(response, FailedResponse2) and response.error.missing and not created_node:
-            if self._create_cognite_file_node(instance_id, http_client, item.as_id(), results):
+            if self._create_cognite_file_node(instance_id, http_client, item.source_id, results):
                 return self._upload_url_data_modeling(item, http_client, results, created_node=True)
             else:
                 return None
@@ -342,7 +346,7 @@ class FileContentIO(UploadableStorageIO[FileContentSelector, MetadataWithFilePat
 
     @classmethod
     def _create_cognite_file_node(
-        cls, instance_id: NodeReference, http_client: HTTPClient, upload_id: str, results: MutableSequence[HTTPMessage]
+        cls, instance_id: NodeReference, http_client: HTTPClient, upload_id: str, results: ItemsResultList
     ) -> bool:
         node_creation = http_client.request_single_retries(
             message=RequestMessage2(
@@ -368,25 +372,25 @@ class FileContentIO(UploadableStorageIO[FileContentSelector, MetadataWithFilePat
         if isinstance(node_creation, SuccessResponse2):
             # Node created successfully
             return True
-        results.append(node_creation.as_item_response(instance_id))
+        results.append(node_creation.as_item_response(upload_id))
         return False
 
     @classmethod
     def _parse_upload_link_response(
-        cls, response: HTTPResult2, item: UploadFileContentItem, results: MutableSequence[HTTPMessage]
+        cls, response: HTTPResult2, item: UploadFileContentItem, results: ItemsResultList
     ) -> str | None:
         if not isinstance(response, SuccessResponse2):
-            results.append(response.as_item_response(item.as_id()))
+            results.append(response.as_item_response(item.source_id))
             return None
         try:
             body = response.body_json
         except ValueError:
             results.append(
-                FailedResponseItems(
+                ItemsFailedResponse2(
                     status_code=response.status_code,
                     body=response.body,
-                    error=ErrorDetails(code=response.status_code, message="Invalid JSON response"),
-                    ids=[item.as_id()],
+                    error=ErrorDetails2(code=response.status_code, message="Invalid JSON response"),
+                    ids=[item.source_id],
                 )
             )
             return None
@@ -396,11 +400,11 @@ class FileContentIO(UploadableStorageIO[FileContentSelector, MetadataWithFilePat
             upload_url = cast(str, body["uploadUrl"])
         except (KeyError, IndexError):
             results.append(
-                FailedResponseItems(
+                ItemsFailedResponse2(
                     status_code=200,
                     body=json.dumps(body),
-                    error=ErrorDetails(code=200, message="Malformed response"),
-                    ids=[item.as_id()],
+                    error=ErrorDetails2(code=200, message="Malformed response"),
+                    ids=[item.source_id],
                 )
             )
             return None

cognite_toolkit/_cdf_tk/storageio/_raw.py

@@ -4,7 +4,8 @@ from uuid import uuid4
 
 from cognite.client.data_classes import Row, RowWrite
 
-from cognite_toolkit._cdf_tk.client.http_client import HTTPClient, HTTPMessage, ItemsRequest
+from cognite_toolkit._cdf_tk.client.http_client import HTTPClient
+from cognite_toolkit._cdf_tk.client.http_client._item_classes import ItemsRequest2, ItemsResultList
 from cognite_toolkit._cdf_tk.cruds import RawDatabaseCRUD, RawTableCRUD
 from cognite_toolkit._cdf_tk.exceptions import ToolkitValueError
 from cognite_toolkit._cdf_tk.utils import sanitize_filename
@@ -60,16 +61,16 @@ class RawIO(
         data_chunk: Sequence[UploadItem[RowWrite]],
         http_client: HTTPClient,
         selector: RawTableSelector | None = None,
-    ) -> Sequence[HTTPMessage]:
+    ) -> ItemsResultList:
         if selector is None:
             raise ToolkitValueError("Selector must be provided for RawIO upload_items")
         url = self.UPLOAD_ENDPOINT.format(dbName=selector.table.db_name, tableName=selector.table.table_name)
         config = http_client.config
-        return http_client.request_with_retries(
-            message=ItemsRequest(
+        return http_client.request_items_retries(
+            message=ItemsRequest2(
                 endpoint_url=config.create_api_url(url),
                 method="POST",
-                items=list(data_chunk),
+                items=data_chunk,
             )
         )
 

cognite_toolkit/_cdf_tk/utils/auth.py

@@ -450,7 +450,7 @@ def prompt_user_environment_variables(current: EnvironmentVariables | None = Non
             if provider != "cdf"
         ],
         default=current.PROVIDER if current else "entra_id",
-    ).ask()
+    ).unsafe_ask()
     exclude = set()
     if provider == "cdf":
         exclude = set(VALID_LOGIN_FLOWS) - {"client_credentials"}
@@ -467,10 +467,10 @@ def prompt_user_environment_variables(current: EnvironmentVariables | None = Non
             "Choose the login flow (How do you going to authenticate?)",
             choices=choices,
             default=current.LOGIN_FLOW if current else "client_credentials",
-        ).ask()
+        ).unsafe_ask()
 
-    cdf_cluster = questionary.text("Enter the CDF cluster", default=current.CDF_CLUSTER if current else "").ask()
-    cdf_project = questionary.text("Enter the CDF project", default=current.CDF_PROJECT if current else "").ask()
+    cdf_cluster = questionary.text("Enter the CDF cluster", default=current.CDF_CLUSTER if current else "").unsafe_ask()
+    cdf_project = questionary.text("Enter the CDF project", default=current.CDF_PROJECT if current else "").unsafe_ask()
     args: dict[str, Any] = (
         current.dump(include_os=False)
         if current and _is_unchanged(current, provider, login_flow, cdf_project, cdf_cluster)  # type: ignore[arg-type]
@@ -490,7 +490,7 @@ def prompt_user_environment_variables(current: EnvironmentVariables | None = Non
     optional_values = env_vars.get_optional_with_value()
     for field_, value in optional_values:
         print(f"  {field_.name}={value}")
-    if questionary.confirm("Do you want to change any of these variables?", default=False).ask():
+    if questionary.confirm("Do you want to change any of these variables?", default=False).unsafe_ask():
         for field_, value in optional_values:
             user_value = get_user_value(field_, value, provider, cdf_cluster, cdf_project, idp_tenant_id)
             setattr(env_vars, field_.name, user_value)
@@ -517,9 +517,9 @@ def get_user_value(
     elif value is not None and not isinstance(value, str):
         default = str(value)
     if is_secret:
-        user_value = questionary.password(f"Enter the {display_name}:", default=default).ask()
+        user_value = questionary.password(f"Enter the {display_name}:", default=default).unsafe_ask()
     else:
-        user_value = questionary.text(f"Enter the {display_name}:", default=default).ask()
+        user_value = questionary.text(f"Enter the {display_name}:", default=default).unsafe_ask()
     if user_value is None:
         raise typer.Exit(0)
     if field_.type is int: