cognite-toolkit 0.6.97__py3-none-any.whl → 0.7.39__py3-none-any.whl

cognite_toolkit/_cdf_tk/storageio/_asset_centric.py

@@ -1,6 +1,6 @@
 from abc import ABC, abstractmethod
 from collections import defaultdict
-from collections.abc import Iterable, MutableSequence, Sequence
+from collections.abc import Iterable, Sequence
 from typing import Any, ClassVar, Generic
 
 from cognite.client.data_classes import (
@@ -8,26 +8,16 @@ from cognite.client.data_classes import (
     Asset,
     AssetList,
     AssetWrite,
-    AssetWriteList,
     Event,
     EventList,
     EventWrite,
-    EventWriteList,
     FileMetadata,
     FileMetadataList,
-    FileMetadataWrite,
-    FileMetadataWriteList,
     Label,
     LabelDefinition,
     TimeSeries,
     TimeSeriesList,
     TimeSeriesWrite,
-    TimeSeriesWriteList,
-)
-from cognite.client.data_classes._base import (
-    T_CogniteResourceList,
-    T_WritableCogniteResource,
-    T_WriteClass,
 )
 
 from cognite_toolkit._cdf_tk.client import ToolkitClient
@@ -40,6 +30,7 @@ from cognite_toolkit._cdf_tk.cruds import (
     TimeSeriesCRUD,
 )
 from cognite_toolkit._cdf_tk.exceptions import ToolkitMissingResourceError, ToolkitNotImplementedError
+from cognite_toolkit._cdf_tk.protocols import T_ResourceRequest, T_ResourceResponse
 from cognite_toolkit._cdf_tk.utils.aggregators import (
     AssetAggregator,
     AssetCentricAggregator,
@@ -50,23 +41,10 @@ from cognite_toolkit._cdf_tk.utils.aggregators import (
 from cognite_toolkit._cdf_tk.utils.cdf import metadata_key_counts
 from cognite_toolkit._cdf_tk.utils.fileio import FileReader, SchemaColumn
 from cognite_toolkit._cdf_tk.utils.fileio._readers import TableReader
-from cognite_toolkit._cdf_tk.utils.http_client import (
-    FailedRequestItems,
-    FailedRequestMessage,
-    FailedResponse,
-    FailedResponseItems,
-    HTTPClient,
-    HTTPMessage,
-    SimpleBodyRequest,
-    SuccessResponse,
-    SuccessResponseItems,
-)
 from cognite_toolkit._cdf_tk.utils.useful_types import (
-    T_ID,
     AssetCentricResource,
     AssetCentricType,
     JsonVal,
-    T_WritableCogniteResourceList,
 )
 
 from ._base import (
@@ -80,11 +58,10 @@ from ._base import (
 from .selectors import AssetCentricSelector, AssetSubtreeSelector, DataSetSelector
 
 
-class BaseAssetCentricIO(
-    Generic[T_ID, T_WriteClass, T_WritableCogniteResource, T_CogniteResourceList, T_WritableCogniteResourceList],
-    TableStorageIO[AssetCentricSelector, T_WritableCogniteResource],
-    ConfigurableStorageIO[AssetCentricSelector, T_WritableCogniteResource],
-    TableUploadableStorageIO[AssetCentricSelector, T_WritableCogniteResource, T_WriteClass],
+class AssetCentricIO(
+    Generic[T_ResourceResponse],
+    TableStorageIO[AssetCentricSelector, T_ResourceResponse],
+    ConfigurableStorageIO[AssetCentricSelector, T_ResourceResponse],
     ABC,
 ):
     RESOURCE_TYPE: ClassVar[AssetCentricType]
@@ -102,7 +79,7 @@ class BaseAssetCentricIO(
         raise NotImplementedError()
 
     @abstractmethod
-    def retrieve(self, ids: Sequence[int]) -> T_WritableCogniteResourceList:
+    def retrieve(self, ids: Sequence[int]) -> Sequence[T_ResourceResponse]:
         raise NotImplementedError()
 
     def count(self, selector: AssetCentricSelector) -> int | None:
@@ -188,6 +165,26 @@ class BaseAssetCentricIO(
             asset_ids.update(item.asset_ids or [])
         self.client.lookup.assets.external_id(list(asset_ids))
 
+    def data_to_row(
+        self, data_chunk: Sequence[T_ResourceResponse], selector: AssetCentricSelector | None = None
+    ) -> list[dict[str, JsonVal]]:
+        rows: list[dict[str, JsonVal]] = []
+        for chunk in self.data_to_json_chunk(data_chunk, selector):
+            if "metadata" in chunk and isinstance(chunk["metadata"], dict):
+                metadata = chunk.pop("metadata")
+                # MyPy does understand that metadata is a dict here due to the check above.
+                for key, value in metadata.items():  # type: ignore[union-attr]
+                    chunk[f"metadata.{key}"] = value
+            rows.append(chunk)
+        return rows
+
+
+class UploadableAssetCentricIO(
+    Generic[T_ResourceResponse, T_ResourceRequest],
+    AssetCentricIO[T_ResourceResponse],
+    TableUploadableStorageIO[AssetCentricSelector, T_ResourceResponse, T_ResourceRequest],
+    ABC,
+):
     def _populate_data_set_external_id_cache(self, chunk: Sequence[dict[str, Any]]) -> None:
         data_set_external_ids: set[str] = set()
         for item in chunk:
@@ -222,22 +219,16 @@ class BaseAssetCentricIO(
                 security_category_names.add(security_category_external_id)
         self.client.lookup.security_categories.id(list(security_category_names))
 
-    def data_to_row(
-        self, data_chunk: Sequence[T_WritableCogniteResource], selector: AssetCentricSelector | None = None
-    ) -> list[dict[str, JsonVal]]:
-        rows: list[dict[str, JsonVal]] = []
-        for chunk in self.data_to_json_chunk(data_chunk, selector):
-            if "metadata" in chunk and isinstance(chunk["metadata"], dict):
-                metadata = chunk.pop("metadata")
-                # MyPy does understand that metadata is a dict here due to the check above.
-                for key, value in metadata.items():  # type: ignore[union-attr]
-                    chunk[f"metadata.{key}"] = value
-            rows.append(chunk)
-        return rows
+    def rows_to_data(
+        self, rows: list[tuple[str, dict[str, JsonVal]]], selector: AssetCentricSelector | None = None
+    ) -> Sequence[UploadItem[T_ResourceRequest]]:
+        # We need to populate caches for any external IDs used in the rows before converting to resources.
+        # Thus, we override this method instead of row_to_resource, and reuse the json_to_resource method that
+        # does the cache population.
+        return self.json_chunk_to_data([(source_id, self.row_to_json(row)) for source_id, row in rows])
 
-    def row_to_resource(
-        self, source_id: str, row: dict[str, JsonVal], selector: AssetCentricSelector | None = None
-    ) -> T_WriteClass:
+    @classmethod
+    def row_to_json(cls, row: dict[str, JsonVal]) -> dict[str, JsonVal]:
         metadata: dict[str, JsonVal] = {}
         cleaned_row: dict[str, JsonVal] = {}
         for key, value in row.items():
@@ -248,10 +239,17 @@ class BaseAssetCentricIO(
                 cleaned_row[key] = value
         if metadata:
             cleaned_row["metadata"] = metadata
-        return self.json_to_resource(cleaned_row)
+        return cleaned_row
+
+    def row_to_resource(
+        self, source_id: str, row: dict[str, JsonVal], selector: AssetCentricSelector | None = None
+    ) -> T_ResourceRequest:
+        raise NotImplementedError(
+            f"This method should not be called. {type(self).__name__} overrides rows_to_data instead."
+        )
 
 
-class AssetIO(BaseAssetCentricIO[str, AssetWrite, Asset, AssetWriteList, AssetList]):
+class AssetIO(UploadableAssetCentricIO[Asset, AssetWrite]):
     KIND = "Assets"
     RESOURCE_TYPE = "asset"
     SUPPORTED_DOWNLOAD_FORMATS = frozenset({".parquet", ".csv", ".ndjson"})
@@ -319,6 +317,9 @@ class AssetIO(BaseAssetCentricIO[str, AssetWrite, Asset, AssetWriteList, AssetLi
             asset_subtree_external_ids=asset_subtree_external_ids,
             data_set_external_ids=data_set_external_ids,
             aggregated_properties=["child_count", "path", "depth"],
+            # We cannot use partitions here as it is not thread safe. This spawn multiple threads
+            # that are not shut down until all data is downloaded. We need to be able to abort.
+            partitions=None,
         ):
             self._collect_dependencies(asset_list, selector)
             yield Page(worker_id="main", items=asset_list)
@@ -347,7 +348,9 @@ class AssetIO(BaseAssetCentricIO[str, AssetWrite, Asset, AssetWriteList, AssetLi
         return self.client.assets.retrieve_multiple(ids)
 
     @classmethod
-    def read_chunks(cls, reader: FileReader) -> Iterable[list[tuple[str, dict[str, JsonVal]]]]:
+    def read_chunks(
+        cls, reader: FileReader, selector: AssetCentricSelector
+    ) -> Iterable[list[tuple[str, dict[str, JsonVal]]]]:
         """Assets require special handling when reading data to ensure parent assets are created first."""
         current_depth = max_depth = 0
         data_name = "row" if isinstance(reader, TableReader) else "line"
@@ -375,7 +378,7 @@ class AssetIO(BaseAssetCentricIO[str, AssetWrite, Asset, AssetWriteList, AssetLi
             current_depth += 1
 
 
-class FileMetadataIO(BaseAssetCentricIO[str, FileMetadataWrite, FileMetadata, FileMetadataWriteList, FileMetadataList]):
+class FileMetadataIO(AssetCentricIO[FileMetadata]):
     KIND = "FileMetadata"
     RESOURCE_TYPE = "file"
     SUPPORTED_DOWNLOAD_FORMATS = frozenset({".parquet", ".csv", ".ndjson"})
@@ -430,54 +433,20 @@ class FileMetadataIO(BaseAssetCentricIO[str, FileMetadataWrite, FileMetadata, Fi
         ]
         return file_schema + metadata_schema
 
-    def stream_data(self, selector: AssetCentricSelector, limit: int | None = None) -> Iterable[Page]:
+    def stream_data(self, selector: AssetCentricSelector, limit: int | None = None) -> Iterable[Page[FileMetadata]]:
         asset_subtree_external_ids, data_set_external_ids = self._get_hierarchy_dataset_pair(selector)
         for file_list in self.client.files(
             chunk_size=self.CHUNK_SIZE,
             limit=limit,
             asset_subtree_external_ids=asset_subtree_external_ids,
             data_set_external_ids=data_set_external_ids,
+            # We cannot use partitions here as it is not thread safe. This spawn multiple threads
+            # that are not shut down until all data is downloaded. We need to be able to abort.
+            partitions=None,
         ):
             self._collect_dependencies(file_list, selector)
             yield Page(worker_id="main", items=file_list)
 
-    def upload_items(
-        self,
-        data_chunk: Sequence[UploadItem[FileMetadataWrite]],
-        http_client: HTTPClient,
-        selector: AssetCentricSelector | None = None,
-    ) -> Sequence[HTTPMessage]:
-        # The /files endpoint only supports creating one file at a time, so we override the default chunked
-        # upload behavior to upload one by one.
-        config = http_client.config
-        results: MutableSequence[HTTPMessage] = []
-        for item in data_chunk:
-            responses = http_client.request_with_retries(
-                message=SimpleBodyRequest(
-                    endpoint_url=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]
-                )
-            )
-            # Convert the responses to per-item responses
-            for message in responses:
-                if isinstance(message, SuccessResponse):
-                    results.append(
-                        SuccessResponseItems(status_code=message.status_code, ids=[item.as_id()], body=message.body)
-                    )
-                elif isinstance(message, FailedResponse):
-                    results.append(
-                        FailedResponseItems(
-                            status_code=message.status_code, ids=[item.as_id()], body=message.body, error=message.error
-                        )
-                    )
-                elif isinstance(message, FailedRequestMessage):
-                    results.append(FailedRequestItems(ids=[item.as_id()], error=message.error))
-                else:
-                    results.append(message)
-        return results
-
     def retrieve(self, ids: Sequence[int]) -> FileMetadataList:
         return self.client.files.retrieve_multiple(ids)
 
@@ -492,20 +461,8 @@ class FileMetadataIO(BaseAssetCentricIO[str, FileMetadataWrite, FileMetadata, Fi
 
         return [self._crud.dump_resource(item) for item in data_chunk]
 
-    def json_chunk_to_data(
-        self, data_chunk: list[tuple[str, dict[str, JsonVal]]]
-    ) -> Sequence[UploadItem[FileMetadataWrite]]:
-        chunks = [item_json for _, item_json in data_chunk]
-        self._populate_asset_external_ids_cache(chunks)
-        self._populate_data_set_external_id_cache(chunks)
-        self._populate_security_category_name_cache(chunks)
-        return super().json_chunk_to_data(data_chunk)
-
-    def json_to_resource(self, item_json: dict[str, JsonVal]) -> FileMetadataWrite:
-        return self._crud.load_resource(item_json)
-
 
-class TimeSeriesIO(BaseAssetCentricIO[str, TimeSeriesWrite, TimeSeries, TimeSeriesWriteList, TimeSeriesList]):
+class TimeSeriesIO(UploadableAssetCentricIO[TimeSeries, TimeSeriesWrite]):
     KIND = "TimeSeries"
     SUPPORTED_DOWNLOAD_FORMATS = frozenset({".parquet", ".csv", ".ndjson"})
     SUPPORTED_COMPRESSIONS = frozenset({".gz"})
@@ -533,6 +490,9 @@ class TimeSeriesIO(BaseAssetCentricIO[str, TimeSeriesWrite, TimeSeries, TimeSeri
             limit=limit,
             asset_subtree_external_ids=asset_subtree_external_ids,
             data_set_external_ids=data_set_external_ids,
+            # We cannot use partitions here as it is not thread safe. This spawn multiple threads
+            # that are not shut down until all data is downloaded. We need to be able to abort.
+            partitions=None,
         ):
             self._collect_dependencies(ts_list, selector)
             yield Page(worker_id="main", items=ts_list)
@@ -598,7 +558,7 @@ class TimeSeriesIO(BaseAssetCentricIO[str, TimeSeriesWrite, TimeSeries, TimeSeri
         return ts_schema + metadata_schema
 
 
-class EventIO(BaseAssetCentricIO[str, EventWrite, Event, EventWriteList, EventList]):
+class EventIO(UploadableAssetCentricIO[Event, EventWrite]):
     KIND = "Events"
     SUPPORTED_DOWNLOAD_FORMATS = frozenset({".parquet", ".csv", ".ndjson"})
     SUPPORTED_COMPRESSIONS = frozenset({".gz"})
@@ -660,6 +620,9 @@ class EventIO(BaseAssetCentricIO[str, EventWrite, Event, EventWriteList, EventLi
             limit=limit,
             asset_subtree_external_ids=asset_subtree_external_ids,
             data_set_external_ids=data_set_external_ids,
+            # We cannot use partitions here as it is not thread safe. This spawn multiple threads
+            # that are not shut down until all data is downloaded. We need to be able to abort.
+            partitions=None,
         ):
             self._collect_dependencies(event_list, selector)
             yield Page(worker_id="main", items=event_list)
@@ -698,7 +661,7 @@ class HierarchyIO(ConfigurableStorageIO[AssetCentricSelector, AssetCentricResour
         self._file_io = FileMetadataIO(client)
         self._timeseries_io = TimeSeriesIO(client)
         self._event_io = EventIO(client)
-        self._io_by_kind: dict[str, BaseAssetCentricIO] = {
+        self._io_by_kind: dict[str, AssetCentricIO] = {
             self._asset_io.KIND: self._asset_io,
             self._file_io.KIND: self._file_io,
             self._timeseries_io.KIND: self._timeseries_io,
@@ -706,7 +669,7 @@ class HierarchyIO(ConfigurableStorageIO[AssetCentricSelector, AssetCentricResour
         }
 
     def as_id(self, item: AssetCentricResource) -> str:
-        return item.external_id or BaseAssetCentricIO.create_internal_identifier(item.id, self.client.config.project)
+        return item.external_id or AssetCentricIO.create_internal_identifier(item.id, self.client.config.project)
 
     def stream_data(
         self, selector: AssetCentricSelector, limit: int | None = None
@@ -726,5 +689,5 @@ class HierarchyIO(ConfigurableStorageIO[AssetCentricSelector, AssetCentricResour
     def configurations(self, selector: AssetCentricSelector) -> Iterable[StorageIOConfig]:
         yield from self.get_resource_io(selector.kind).configurations(selector)
 
-    def get_resource_io(self, kind: str) -> BaseAssetCentricIO:
+    def get_resource_io(self, kind: str) -> AssetCentricIO:
         return self._io_by_kind[kind]

cognite_toolkit/_cdf_tk/storageio/_base.py

@@ -1,17 +1,15 @@
 from abc import ABC, abstractmethod
 from collections.abc import Iterable, Mapping, Sequence, Sized
 from dataclasses import dataclass
-from typing import ClassVar, Generic, TypeVar
-
-from cognite.client.data_classes._base import T_CogniteResource
+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.protocols import T_ResourceRequest, T_ResourceResponse
 from cognite_toolkit._cdf_tk.utils.collection import chunker
-from cognite_toolkit._cdf_tk.utils.fileio import FileReader, SchemaColumn
-from cognite_toolkit._cdf_tk.utils.fileio._readers import TableReader
+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 JsonVal, T_WriteCogniteResource
+from cognite_toolkit._cdf_tk.utils.useful_types import JsonVal
 
 from .selectors import DataSelector
 
@@ -28,9 +26,9 @@ T_Selector = TypeVar("T_Selector", bound=DataSelector)
 
 
 @dataclass
-class Page(Generic[T_CogniteResource], Sized):
+class Page(Generic[T_ResourceResponse], Sized):
     worker_id: str
-    items: Sequence[T_CogniteResource]
+    items: Sequence[T_ResourceResponse]
     next_cursor: str | None = None
 
     def __len__(self) -> int:
@@ -38,7 +36,7 @@ class Page(Generic[T_CogniteResource], Sized):
 
 
 @dataclass
-class UploadItem(Generic[T_WriteCogniteResource]):
+class UploadItem(Generic[T_ResourceRequest]):
     """An item to be uploaded to CDF, consisting of a source ID and the writable Cognite resource.
 
     Attributes:
@@ -47,7 +45,7 @@ class UploadItem(Generic[T_WriteCogniteResource]):
     """
 
     source_id: str
-    item: T_WriteCogniteResource
+    item: T_ResourceRequest
 
     def as_id(self) -> str:
         return self.source_id
@@ -56,7 +54,7 @@ class UploadItem(Generic[T_WriteCogniteResource]):
         return self.item.dump(camel_case=True)
 
 
-class StorageIO(ABC, Generic[T_Selector, T_CogniteResource]):
+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
@@ -80,7 +78,7 @@ class StorageIO(ABC, Generic[T_Selector, T_CogniteResource]):
         self.client = client
 
     @abstractmethod
-    def as_id(self, item: T_CogniteResource) -> str:
+    def as_id(self, item: T_ResourceResponse) -> str:
         """Convert an item to its corresponding ID.
         Args:
             item: The item to convert.
@@ -116,7 +114,7 @@ class StorageIO(ABC, Generic[T_Selector, T_CogniteResource]):
 
     @abstractmethod
     def data_to_json_chunk(
-        self, data_chunk: Sequence[T_CogniteResource], selector: T_Selector | None = None
+        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.
 
@@ -132,7 +130,7 @@ class StorageIO(ABC, Generic[T_Selector, T_CogniteResource]):
 
 
 class UploadableStorageIO(
-    Generic[T_Selector, T_CogniteResource, T_WriteCogniteResource], StorageIO[T_Selector, T_CogniteResource], ABC
+    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.
 
@@ -145,12 +143,14 @@ class UploadableStorageIO(
 
     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_WriteCogniteResource]],
+        data_chunk: Sequence[UploadItem[T_ResourceRequest]],
         http_client: HTTPClient,
         selector: T_Selector | None = None,
     ) -> Sequence[HTTPMessage]:
@@ -170,10 +170,17 @@ class UploadableStorageIO(
             raise ValueError(f"Data chunk size {len(data_chunk)} exceeds the maximum CHUNK_SIZE of {self.CHUNK_SIZE}.")
 
         config = http_client.config
+        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=config.create_api_url(self.UPLOAD_ENDPOINT),
-                method="POST",
+                endpoint_url=url,
+                method=self.UPLOAD_ENDPOINT_METHOD,
                 items=list(data_chunk),
                 extra_body_fields=dict(self.UPLOAD_EXTRA_ARGS or {}),
             )
@@ -181,7 +188,7 @@ class UploadableStorageIO(
 
     def json_chunk_to_data(
         self, data_chunk: list[tuple[str, dict[str, JsonVal]]]
-    ) -> Sequence[UploadItem[T_WriteCogniteResource]]:
+    ) -> Sequence[UploadItem[T_ResourceRequest]]:
         """Convert a JSON-compatible chunk of data back to a writable Cognite resource list.
 
         Args:
@@ -190,14 +197,14 @@ class UploadableStorageIO(
         Returns:
             A writable Cognite resource list representing the data.
         """
-        result: list[UploadItem[T_WriteCogniteResource]] = []
+        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 json_to_resource(self, item_json: dict[str, JsonVal]) -> T_WriteCogniteResource:
+    def json_to_resource(self, item_json: dict[str, JsonVal]) -> T_ResourceRequest:
         """Convert a JSON-compatible dictionary back to a writable Cognite resource.
 
         Args:
@@ -208,20 +215,45 @@ class UploadableStorageIO(
         raise NotImplementedError()
 
     @classmethod
-    def read_chunks(cls, reader: FileReader) -> Iterable[list[tuple[str, dict[str, JsonVal]]]]:
-        data_name = "row" if isinstance(reader, TableReader) else "line"
+    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:
+            The number of items in the reader.
+        """
+        return reader.count()
+
 
-class TableUploadableStorageIO(UploadableStorageIO[T_Selector, T_CogniteResource, T_WriteCogniteResource], ABC):
+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_WriteCogniteResource]]:
+    ) -> Sequence[UploadItem[T_ResourceRequest]]:
         """Convert a row-based JSON-compatible chunk of data back to a writable Cognite resource list.
 
         Args:
@@ -232,7 +264,7 @@ class TableUploadableStorageIO(UploadableStorageIO[T_Selector, T_CogniteResource
         Returns:
             A writable Cognite resource list representing the data.
         """
-        result: list[UploadItem[T_WriteCogniteResource]] = []
+        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))
@@ -241,7 +273,7 @@ class TableUploadableStorageIO(UploadableStorageIO[T_Selector, T_CogniteResource
     @abstractmethod
     def row_to_resource(
         self, source_id: str, row: dict[str, JsonVal], selector: T_Selector | None = None
-    ) -> T_WriteCogniteResource:
+    ) -> T_ResourceRequest:
         """Convert a row-based JSON-compatible dictionary back to a writable Cognite resource.
 
         Args:
@@ -254,7 +286,7 @@ class TableUploadableStorageIO(UploadableStorageIO[T_Selector, T_CogniteResource
         raise NotImplementedError()
 
 
-class ConfigurableStorageIO(StorageIO[T_Selector, T_CogniteResource], ABC):
+class ConfigurableStorageIO(StorageIO[T_Selector, T_ResourceResponse], ABC):
     """A base class for storage items that support configurations for different storage items."""
 
     @abstractmethod
@@ -263,7 +295,7 @@ class ConfigurableStorageIO(StorageIO[T_Selector, T_CogniteResource], ABC):
         raise NotImplementedError()
 
 
-class TableStorageIO(StorageIO[T_Selector, T_CogniteResource], ABC):
+class TableStorageIO(StorageIO[T_Selector, T_ResourceResponse], ABC):
     """A base class for storage items that support table schemas."""
 
     @abstractmethod
@@ -281,7 +313,7 @@ class TableStorageIO(StorageIO[T_Selector, T_CogniteResource], ABC):
 
     @abstractmethod
     def data_to_row(
-        self, data_chunk: Sequence[T_CogniteResource], selector: T_Selector | None = None
+        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.