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

cognite_toolkit/_cdf_tk/storageio/_base.py

@@ -1,20 +1,17 @@
 from abc import ABC, abstractmethod
-from collections.abc import Iterable, Mapping, Sequence
+from collections.abc import Iterable, Mapping, Sequence, Sized
 from dataclasses import dataclass
-from pathlib import Path
-from typing import ClassVar, Generic
-
-from cognite.client.data_classes._base import (
-    T_CogniteResourceList,
-)
-from rich.console import Console
+from typing import ClassVar, Generic, Literal, TypeVar
 
 from cognite_toolkit._cdf_tk.client import ToolkitClient
 from cognite_toolkit._cdf_tk.exceptions import ToolkitNotImplementedError
-from cognite_toolkit._cdf_tk.utils.collection import chunker_sequence
-from cognite_toolkit._cdf_tk.utils.fileio import SchemaColumn
+from cognite_toolkit._cdf_tk.protocols import T_ResourceRequest, T_ResourceResponse
+from cognite_toolkit._cdf_tk.utils.collection import chunker
+from cognite_toolkit._cdf_tk.utils.fileio import MultiFileReader, SchemaColumn
 from cognite_toolkit._cdf_tk.utils.http_client import HTTPClient, HTTPMessage, ItemsRequest
-from cognite_toolkit._cdf_tk.utils.useful_types import T_ID, JsonVal, T_Selector, T_WritableCogniteResourceList
+from cognite_toolkit._cdf_tk.utils.useful_types import JsonVal
+
+from .selectors import DataSelector
 
 
 @dataclass
@@ -22,41 +19,66 @@ class StorageIOConfig:
     kind: str
     folder_name: str
     value: JsonVal
+    filename: str | None = None
 
 
-class StorageIO(ABC, Generic[T_ID, T_Selector, T_CogniteResourceList, T_WritableCogniteResourceList]):
-    """This is a base class for all storage classes in Cognite Toolkit
+T_Selector = TypeVar("T_Selector", bound=DataSelector)
+
+
+@dataclass
+class Page(Generic[T_ResourceResponse], Sized):
+    worker_id: str
+    items: Sequence[T_ResourceResponse]
+    next_cursor: str | None = None
 
-    It defines the interface for interacting with storage items in CDF, such as downloading,
-    uploading, and managing configurations. Each storage type (e.g., RAW, asset-centric)
-    should implement this interface to provide specific functionality.
+    def __len__(self) -> int:
+        return len(self.items)
+
+
+@dataclass
+class UploadItem(Generic[T_ResourceRequest]):
+    """An item to be uploaded to CDF, consisting of a source ID and the writable Cognite resource.
 
     Attributes:
-        FOLDER_NAME: The name of the folder in which the storage item is located.
-        KIND: The type of storage (e.g., 'raw', 'assets').
-        DISPLAY_NAME: A human-readable name for the storage item.
+        source_id: The source identifier for the item. For example, the line number in a CSV file.
+        item: The writable Cognite resource to be uploaded.
+    """
+
+    source_id: str
+    item: T_ResourceRequest
+
+    def as_id(self) -> str:
+        return self.source_id
+
+    def dump(self) -> JsonVal:
+        return self.item.dump(camel_case=True)
+
+
+class StorageIO(ABC, Generic[T_Selector, T_ResourceResponse]):
+    """This is a base class for all storage classes in Cognite Toolkit
+
+    It defines the interface for downloading data from CDF. Note this can also be used for multiple
+    types of resources, for example, a hierarchy of assets/files/events/time series.
+
+    Attributes
         SUPPORTED_DOWNLOAD_FORMATS: A set of formats that the storage item supports for downloading.
         SUPPORTED_COMPRESSIONS: A set of compression formats that the storage item supports.
-        SUPPORTED_READ_FORMATS: A set of formats that the storage item supports for reading.
         CHUNK_SIZE: The size of the data chunks to be processed during download and upload operations.
+
+    Args:
         client: An instance of ToolkitClient to interact with the CDF API.
     """
 
-    FOLDER_NAME: str
-    KIND: str
-    DISPLAY_NAME: str
-    SUPPORTED_DOWNLOAD_FORMATS: frozenset[str]
-    SUPPORTED_COMPRESSIONS: frozenset[str]
-    SUPPORTED_READ_FORMATS: frozenset[str]
-    CHUNK_SIZE: int
-    UPLOAD_ENDPOINT: ClassVar[str]
-    UPLOAD_EXTRA_ARGS: ClassVar[Mapping[str, JsonVal] | None] = None
+    SUPPORTED_DOWNLOAD_FORMATS: ClassVar[frozenset[str]]
+    SUPPORTED_COMPRESSIONS: ClassVar[frozenset[str]]
+    CHUNK_SIZE: ClassVar[int]
+    BASE_SELECTOR: ClassVar[type[DataSelector]]
 
     def __init__(self, client: ToolkitClient) -> None:
         self.client = client
 
     @abstractmethod
-    def as_id(self, item: dict[str, JsonVal] | object) -> T_ID:
+    def as_id(self, item: T_ResourceResponse) -> str:
         """Convert an item to its corresponding ID.
         Args:
             item: The item to convert.
@@ -66,7 +88,7 @@ class StorageIO(ABC, Generic[T_ID, T_Selector, T_CogniteResourceList, T_Writable
         raise NotImplementedError()
 
     @abstractmethod
-    def stream_data(self, selector: T_Selector, limit: int | None = None) -> Iterable[T_WritableCogniteResourceList]:
+    def stream_data(self, selector: T_Selector, limit: int | None = None) -> Iterable[Page]:
         """Download items from the storage given the selection criteria.
 
         Args:
@@ -91,21 +113,52 @@ class StorageIO(ABC, Generic[T_ID, T_Selector, T_CogniteResourceList, T_Writable
         raise NotImplementedError()
 
     @abstractmethod
-    def upload_items(self, data_chunk: T_CogniteResourceList, selector: T_Selector) -> None:
-        """Upload a chunk of data to the storage.
+    def data_to_json_chunk(
+        self, data_chunk: Sequence[T_ResourceResponse], selector: T_Selector | None = None
+    ) -> list[dict[str, JsonVal]]:
+        """Convert a chunk of data to a JSON-compatible format.
 
         Args:
-            data_chunk: The chunk of data to upload, which should be a list of writable Cognite resources.
-            selector: The selection criteria to identify where to upload the data.
+            data_chunk: The chunk of data to convert, which should be a writable Cognite resource list.
+            selector: Optional selection criteria to identify the data. This is required for some storage types.
+
+        Returns:
+            A list of dictionaries representing the data in a JSON-compatible format.
+
         """
         raise NotImplementedError()
 
-    def upload_items_force(
-        self, data_chunk: T_CogniteResourceList, http_client: HTTPClient, selector: T_Selector | None = None
+
+class UploadableStorageIO(
+    Generic[T_Selector, T_ResourceResponse, T_ResourceRequest], StorageIO[T_Selector, T_ResourceResponse], ABC
+):
+    """A base class for storage items that support uploading data to CDF.
+
+    Attributes:
+        KIND: The kind of storage item (e.g., "RAW", "AssetCentric").
+        SUPPORTED_READ_FORMATS: A set of formats that the storage item supports for reading.
+        UPLOAD_ENDPOINT: The API endpoint for uploading data to the storage item.
+        UPLOAD_EXTRA_ARGS: Additional arguments to include in the upload request.
+    """
+
+    KIND: ClassVar[str]
+    SUPPORTED_READ_FORMATS: ClassVar[frozenset[str]]
+    UPLOAD_ENDPOINT_TYPE: Literal["app", "api"] = "api"
+    UPLOAD_ENDPOINT_METHOD: Literal["GET", "POST", "PATCH", "DELETE", "PUT"] = "POST"
+    UPLOAD_ENDPOINT: ClassVar[str]
+    UPLOAD_EXTRA_ARGS: ClassVar[Mapping[str, JsonVal] | None] = None
+
+    def upload_items(
+        self,
+        data_chunk: Sequence[UploadItem[T_ResourceRequest]],
+        http_client: HTTPClient,
+        selector: T_Selector | None = None,
     ) -> Sequence[HTTPMessage]:
         """Upload a chunk of data to the storage using a custom HTTP client.
         This ensures that even if one item in the chunk fails, the rest will still be uploaded.
 
+        This assumes that the data_chunk is respecting the CHUNK_SIZE of the storage.
+
         Args:
             data_chunk: The chunk of data to upload, which should be a list of writable Cognite resources.
             http_client: The custom HTTP client to use for the upload.
@@ -113,83 +166,163 @@ class StorageIO(ABC, Generic[T_ID, T_Selector, T_CogniteResourceList, T_Writable
         """
         if not hasattr(self, "UPLOAD_ENDPOINT"):
             raise ToolkitNotImplementedError(f"Upload not implemented for {self.KIND} storage.")
+        if len(data_chunk) > self.CHUNK_SIZE:
+            raise ValueError(f"Data chunk size {len(data_chunk)} exceeds the maximum CHUNK_SIZE of {self.CHUNK_SIZE}.")
 
         config = http_client.config
-        results: list[HTTPMessage] = []
-        for batch in chunker_sequence(data_chunk, self.CHUNK_SIZE):
-            batch_results = http_client.request_with_retries(
-                message=ItemsRequest(
-                    endpoint_url=config.create_api_url(self.UPLOAD_ENDPOINT),
-                    method="POST",
-                    # The dump method from the PySDK always returns JsonVal, but mypy cannot infer that
-                    items=batch.dump(camel_case=True),  # type: ignore[arg-type]
-                    extra_body_fields=dict(self.UPLOAD_EXTRA_ARGS or {}),
-                    as_id=self.as_id,
-                )
+        if self.UPLOAD_ENDPOINT_TYPE == "api":
+            url = config.create_api_url(self.UPLOAD_ENDPOINT)
+        elif self.UPLOAD_ENDPOINT_TYPE == "app":
+            url = config.create_app_url(self.UPLOAD_ENDPOINT)
+        else:
+            raise ToolkitNotImplementedError(f"Unsupported UPLOAD_ENDPOINT_TYPE {self.UPLOAD_ENDPOINT_TYPE!r}.")
+
+        return http_client.request_with_retries(
+            message=ItemsRequest(
+                endpoint_url=url,
+                method=self.UPLOAD_ENDPOINT_METHOD,
+                items=list(data_chunk),
+                extra_body_fields=dict(self.UPLOAD_EXTRA_ARGS or {}),
             )
-            results.extend(batch_results)
-        return results
+        )
+
+    def json_chunk_to_data(
+        self, data_chunk: list[tuple[str, dict[str, JsonVal]]]
+    ) -> Sequence[UploadItem[T_ResourceRequest]]:
+        """Convert a JSON-compatible chunk of data back to a writable Cognite resource list.
+
+        Args:
+            data_chunk: A list of tuples, each containing a source ID and a dictionary representing
+                the data in a JSON-compatible format.
+        Returns:
+            A writable Cognite resource list representing the data.
+        """
+        result: list[UploadItem[T_ResourceRequest]] = []
+        for source_id, item_json in data_chunk:
+            item = self.json_to_resource(item_json)
+            result.append(UploadItem(source_id=source_id, item=item))
+        return result
 
     @abstractmethod
-    def data_to_json_chunk(self, data_chunk: T_WritableCogniteResourceList) -> list[dict[str, JsonVal]]:
-        """Convert a chunk of data to a JSON-compatible format.
+    def json_to_resource(self, item_json: dict[str, JsonVal]) -> T_ResourceRequest:
+        """Convert a JSON-compatible dictionary back to a writable Cognite resource.
 
         Args:
-            data_chunk: The chunk of data to convert, which should be a writable Cognite resource list.
+            item_json: A dictionary representing the data in a JSON-compatible format.
+        Returns:
+            A writable Cognite resource representing the data.
+        """
+        raise NotImplementedError()
+
+    @classmethod
+    def read_chunks(
+        cls, reader: MultiFileReader, selector: T_Selector
+    ) -> Iterable[list[tuple[str, dict[str, JsonVal]]]]:
+        """Read data from a MultiFileReader in chunks.
+
+        This method yields chunks of data, where each chunk is a list of tuples. Each tuple contains a source ID
+        (e.g., line number or row identifier) and a dictionary representing the data in a JSON-compatible format.
+
+        This method can be overridden by subclasses to customize how data is read and chunked.
+        Args:
+            reader: An instance of MultiFileReader to read data from.
+            selector: The selection criteria to identify the data.
+        """
+        data_name = "row" if reader.is_table else "line"
+        # Include name of line for better error messages
+        iterable = ((f"{data_name} {line_no}", item) for line_no, item in reader.read_chunks_with_line_numbers())
+
+        yield from chunker(iterable, cls.CHUNK_SIZE)
+
+    @classmethod
+    def count_chunks(cls, reader: MultiFileReader) -> int:
+        """Count the number of items in a MultiFileReader.
 
+        This method can be overridden by subclasses to customize how items are counted.
+
+        Args:
+            reader: An instance of MultiFileReader to count items from.
         Returns:
-            A list of dictionaries representing the data in a JSON-compatible format.
+            The number of items in the reader.
+        """
+        return reader.count()
+
+
+class TableUploadableStorageIO(UploadableStorageIO[T_Selector, T_ResourceResponse, T_ResourceRequest], ABC):
+    """A base class for storage items that support uploading data with table schemas."""
 
+    def rows_to_data(
+        self, rows: list[tuple[str, dict[str, JsonVal]]], selector: T_Selector | None = None
+    ) -> Sequence[UploadItem[T_ResourceRequest]]:
+        """Convert a row-based JSON-compatible chunk of data back to a writable Cognite resource list.
+
+        Args:
+            rows: A list of tuples, each containing a source ID and a dictionary representing
+                the data in a JSON-compatible format.
+            selector: Optional selection criteria to identify where to upload the data. This is required for some storage types.
+
+        Returns:
+            A writable Cognite resource list representing the data.
         """
-        raise NotImplementedError()
+        result: list[UploadItem[T_ResourceRequest]] = []
+        for source_id, row in rows:
+            item = self.row_to_resource(source_id, row, selector=selector)
+            result.append(UploadItem(source_id=source_id, item=item))
+        return result
 
     @abstractmethod
-    def json_chunk_to_data(self, data_chunk: list[dict[str, JsonVal]]) -> T_CogniteResourceList:
-        """Convert a JSON-compatible chunk of data back to a writable Cognite resource list.
+    def row_to_resource(
+        self, source_id: str, row: dict[str, JsonVal], selector: T_Selector | None = None
+    ) -> T_ResourceRequest:
+        """Convert a row-based JSON-compatible dictionary back to a writable Cognite resource.
 
         Args:
-            data_chunk: A list of dictionaries representing the data in a JSON-compatible format.
+            row: A dictionary representing the data in a JSON-compatible format.
+            source_id: The source identifier for the item. For example, the line number in a CSV file.
+            selector: Optional selection criteria to identify where to upload the data. This is required for some storage types.
         Returns:
-            A writable Cognite resource list representing the data.
+            A writable Cognite resource representing the data.
         """
         raise NotImplementedError()
 
+
+class ConfigurableStorageIO(StorageIO[T_Selector, T_ResourceResponse], ABC):
+    """A base class for storage items that support configurations for different storage items."""
+
     @abstractmethod
     def configurations(self, selector: T_Selector) -> Iterable[StorageIOConfig]:
         """Return configurations for the storage item."""
         raise NotImplementedError()
 
-    @abstractmethod
-    def load_selector(self, datafile: Path) -> T_Selector:
-        """Load the selector from adjacent filepath."""
-        raise NotImplementedError()
-
-    @abstractmethod
-    def ensure_configurations(self, selector: T_Selector, console: Console | None = None) -> None:
-        """Ensure that the necessary configurations for the storage item are in place.
 
-        This method should create the necessary configurations in CDF if they do not exist.
-        For example, for RAW tables, this will create the RAW database and table.
+class TableStorageIO(StorageIO[T_Selector, T_ResourceResponse], ABC):
+    """A base class for storage items that support table schemas."""
 
-        For asset-centric storage, this will create labels and data sets.
+    @abstractmethod
+    def get_schema(self, selector: T_Selector) -> list[SchemaColumn]:
+        """Get the schema of the table associated with the given selector.
 
         Args:
-            selector: The selection criteria to find the data.
-            console: An optional console for outputting messages during the configuration process.
+            selector: The selection criteria to identify the data.
+
+        Returns:
+            A list of SchemaColumn objects representing the schema of the table.
 
         """
         raise NotImplementedError()
 
-
-class TableStorageIO(StorageIO[T_ID, T_Selector, T_CogniteResourceList, T_WritableCogniteResourceList], ABC):
     @abstractmethod
-    def get_schema(self, selector: T_Selector) -> list[SchemaColumn]:
-        """Get the schema of the table associated with the given selector.
+    def data_to_row(
+        self, data_chunk: Sequence[T_ResourceResponse], selector: T_Selector | None = None
+    ) -> list[dict[str, JsonVal]]:
+        """Convert a chunk of data to a row-based JSON-compatible format.
 
         Args:
-            selector: The selection criteria to identify the data.
+            data_chunk: The chunk of data to convert, which should be a writable Cognite resource list.
+            selector: Optional selection criteria to identify the data. This is required for some storage types.
 
         Returns:
-            A list of SchemaColumn objects representing the schema of the table.
+            A list of dictionaries representing the data in a JSON-compatible format.
 
         """
+        raise NotImplementedError()