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

cognite_toolkit/_cdf_tk/storageio/_asset_centric.py

@@ -1,77 +1,85 @@
 from abc import ABC, abstractmethod
 from collections import defaultdict
 from collections.abc import Iterable, Sequence
-from pathlib import Path
-from typing import Generic
+from typing import Any, ClassVar, Generic
 
 from cognite.client.data_classes import (
+    AggregateResultItem,
     Asset,
     AssetList,
     AssetWrite,
-    AssetWriteList,
-    DataSetWriteList,
+    Event,
+    EventList,
+    EventWrite,
     FileMetadata,
     FileMetadataList,
-    FileMetadataWrite,
-    FileMetadataWriteList,
     Label,
     LabelDefinition,
+    TimeSeries,
+    TimeSeriesList,
+    TimeSeriesWrite,
 )
-from cognite.client.data_classes._base import (
-    T_CogniteResourceList,
-    T_WritableCogniteResource,
-    T_WriteClass,
-)
-from cognite.client.data_classes.labels import LabelDefinitionWriteList
-from rich.console import Console
 
 from cognite_toolkit._cdf_tk.client import ToolkitClient
-from cognite_toolkit._cdf_tk.cruds import AssetCRUD, DataSetsCRUD, FileMetadataCRUD, LabelCRUD, ResourceCRUD
-from cognite_toolkit._cdf_tk.exceptions import ToolkitNotImplementedError
-from cognite_toolkit._cdf_tk.utils.aggregators import AssetAggregator, AssetCentricAggregator, FileAggregator
+from cognite_toolkit._cdf_tk.cruds import (
+    AssetCRUD,
+    DataSetsCRUD,
+    EventCRUD,
+    FileMetadataCRUD,
+    LabelCRUD,
+    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,
+    EventAggregator,
+    FileAggregator,
+    TimeSeriesAggregator,
+)
 from cognite_toolkit._cdf_tk.utils.cdf import metadata_key_counts
-from cognite_toolkit._cdf_tk.utils.file import find_files_with_suffix_and_prefix
-from cognite_toolkit._cdf_tk.utils.fileio import SchemaColumn
-from cognite_toolkit._cdf_tk.utils.useful_types import T_ID, JsonVal, T_WritableCogniteResourceList
+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.useful_types import (
+    AssetCentricResource,
+    AssetCentricType,
+    JsonVal,
+)
 
-from ._base import StorageIOConfig, TableStorageIO
-from ._selectors import AssetCentricFileSelector, AssetCentricSelector, AssetSubtreeSelector, DataSetSelector
+from ._base import (
+    ConfigurableStorageIO,
+    Page,
+    StorageIOConfig,
+    TableStorageIO,
+    TableUploadableStorageIO,
+    UploadItem,
+)
+from .selectors import AssetCentricSelector, AssetSubtreeSelector, DataSetSelector
 
 
-class BaseAssetCentricIO(
-    Generic[T_ID, T_WriteClass, T_WritableCogniteResource, T_CogniteResourceList, T_WritableCogniteResourceList],
-    TableStorageIO[int, AssetCentricSelector, T_CogniteResourceList, T_WritableCogniteResourceList],
+class AssetCentricIO(
+    Generic[T_ResourceResponse],
+    TableStorageIO[AssetCentricSelector, T_ResourceResponse],
+    ConfigurableStorageIO[AssetCentricSelector, T_ResourceResponse],
     ABC,
 ):
+    RESOURCE_TYPE: ClassVar[AssetCentricType]
     CHUNK_SIZE = 1000
+    BASE_SELECTOR = AssetCentricSelector
 
     def __init__(self, client: ToolkitClient) -> None:
         super().__init__(client)
-        self._loader = self._get_loader()
         self._aggregator = self._get_aggregator()
         self._downloaded_data_sets_by_selector: dict[AssetCentricSelector, set[int]] = defaultdict(set)
         self._downloaded_labels_by_selector: dict[AssetCentricSelector, set[str]] = defaultdict(set)
 
-    def as_id(self, item: dict[str, JsonVal] | object) -> int:
-        if isinstance(item, dict) and isinstance(item.get("id"), int):
-            # MyPy checked above.
-            return item["id"]  # type: ignore[return-value]
-        raise TypeError(f"Cannot extract ID from item of type {type(item).__name__!r}")
-
-    @abstractmethod
-    def _get_loader(
-        self,
-    ) -> ResourceCRUD[
-        T_ID, T_WriteClass, T_WritableCogniteResource, T_CogniteResourceList, T_WritableCogniteResourceList
-    ]:
-        raise NotImplementedError()
-
     @abstractmethod
     def _get_aggregator(self) -> AssetCentricAggregator:
         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:
@@ -81,9 +89,6 @@ class BaseAssetCentricIO(
             return self._aggregator.count(hierarchy=selector.hierarchy)
         return None
 
-    def data_to_json_chunk(self, data_chunk: T_WritableCogniteResourceList) -> list[dict[str, JsonVal]]:
-        return [self._loader.dump_resource(item) for item in data_chunk]
-
     def configurations(self, selector: AssetCentricSelector) -> Iterable[StorageIOConfig]:
         data_set_ids = self._downloaded_data_sets_by_selector[selector]
         if data_set_ids:
@@ -94,17 +99,32 @@ class BaseAssetCentricIO(
             list(self._downloaded_labels_by_selector[selector]), LabelCRUD.create_loader(self.client)
         )
 
-    def _collect_dependencies(self, resources: AssetList | FileMetadataList, selector: AssetCentricSelector) -> None:
+    def _get_hierarchy_dataset_pair(self, selector: AssetCentricSelector) -> tuple[list[str] | None, list[str] | None]:
+        asset_subtree_external_ids: list[str] | None = None
+        data_set_external_ids: list[str] | None = None
+        if isinstance(selector, DataSetSelector):
+            data_set_external_ids = [selector.data_set_external_id]
+        elif isinstance(selector, AssetSubtreeSelector):
+            asset_subtree_external_ids = [selector.hierarchy]
+        else:
+            # This selector is for uploads, not for downloading from CDF.
+            raise ToolkitNotImplementedError(f"Selector type {type(selector)} not supported for {type(self).__name__}.")
+        return asset_subtree_external_ids, data_set_external_ids
+
+    def _collect_dependencies(
+        self, resources: AssetList | FileMetadataList | TimeSeriesList | EventList, selector: AssetCentricSelector
+    ) -> None:
         for resource in resources:
             if resource.data_set_id:
                 self._downloaded_data_sets_by_selector[selector].add(resource.data_set_id)
-            for label in resource.labels or []:
-                if isinstance(label, str):
-                    self._downloaded_labels_by_selector[selector].add(label)
-                elif isinstance(label, Label | LabelDefinition) and label.external_id:
-                    self._downloaded_labels_by_selector[selector].add(label.external_id)
-                elif isinstance(label, dict) and "externalId" in label:
-                    self._downloaded_labels_by_selector[selector].add(label["externalId"])
+            if isinstance(resource, Asset | FileMetadata):
+                for label in resource.labels or []:
+                    if isinstance(label, str):
+                        self._downloaded_labels_by_selector[selector].add(label)
+                    elif isinstance(label, Label | LabelDefinition) and label.external_id:
+                        self._downloaded_labels_by_selector[selector].add(label.external_id)
+                    elif isinstance(label, dict) and "externalId" in label:
+                        self._downloaded_labels_by_selector[selector].add(label["externalId"])
 
     @classmethod
     def _configurations(
@@ -122,64 +142,127 @@ class BaseAssetCentricIO(
             value=[loader.dump_resource(item) for item in items],  # type: ignore[arg-type]
         )
 
-    def load_selector(self, datafile: Path) -> AssetCentricSelector:
-        return AssetCentricFileSelector(datafile=datafile)
+    def _create_identifier(self, internal_id: int) -> str:
+        return self.create_internal_identifier(internal_id, self.client.config.project)
 
-    def ensure_configurations(self, selector: AssetCentricSelector, console: Console | None = None) -> None:
-        """Ensures that all data sets and labels referenced by the asset selection exist in CDF."""
-        if not isinstance(selector, AssetCentricFileSelector):
-            return None
-        datafile = selector.datafile
-        filepaths = find_files_with_suffix_and_prefix(
-            datafile.parent.parent / DataSetsCRUD.folder_name, datafile.name, suffix=f".{DataSetsCRUD.kind}.yaml"
-        )
-        self._create_if_not_exists(filepaths, DataSetsCRUD.create_loader(self.client), console)
-
-        filepaths = find_files_with_suffix_and_prefix(
-            datafile.parent.parent / LabelCRUD.folder_name, datafile.name, suffix=f".{LabelCRUD.kind}.yaml"
-        )
-        self._create_if_not_exists(filepaths, LabelCRUD.create_loader(self.client), console)
-        return None
+    @classmethod
+    def create_internal_identifier(cls, internal_id: int, project: str) -> str:
+        return f"INTERNAL_ID_project_{project}_{internal_id!s}"
+
+    def _populate_data_set_id_cache(self, chunk: Sequence[Asset | FileMetadata | TimeSeries | Event]) -> None:
+        data_set_ids = {item.data_set_id for item in chunk if item.data_set_id is not None}
+        self.client.lookup.data_sets.external_id(list(data_set_ids))
+
+    def _populate_security_category_cache(self, chunk: Sequence[FileMetadata | TimeSeries]) -> None:
+        security_category_ids: set[int] = set()
+        for item in chunk:
+            security_category_ids.update(item.security_categories or [])
+        self.client.lookup.security_categories.external_id(list(security_category_ids))
+
+    def _populate_asset_id_cache(self, chunk: Sequence[FileMetadata | Event]) -> None:
+        asset_ids: set[int] = set()
+        for item in chunk:
+            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:
+            data_set_external_id = item.get("dataSetExternalId")
+            if isinstance(data_set_external_id, str):
+                data_set_external_ids.add(data_set_external_id)
+        self.client.lookup.data_sets.id(list(data_set_external_ids))
+
+    def _populate_asset_external_ids_cache(self, chunk: Sequence[dict[str, Any]]) -> None:
+        asset_external_id_set: set[str] = set()
+        for item in chunk:
+            asset_external_ids = item.get("assetExternalIds")
+            if isinstance(asset_external_ids, list):
+                for asset_external_id_item in asset_external_ids:
+                    if isinstance(asset_external_id_item, str):
+                        asset_external_id_set.add(asset_external_id_item)
+            asset_external_id = item.get("assetExternalId")
+            if isinstance(asset_external_id, str):
+                asset_external_id_set.add(asset_external_id)
+        self.client.lookup.assets.id(list(asset_external_id_set))
+
+    def _populate_security_category_name_cache(self, chunk: Sequence[dict[str, Any]]) -> None:
+        security_category_names: set[str] = set()
+        for item in chunk:
+            security_category_external_ids_list = item.get("securityCategoryNames")
+            if isinstance(security_category_external_ids_list, list):
+                for security_category_external_id_item in security_category_external_ids_list:
+                    if isinstance(security_category_external_id_item, str):
+                        security_category_names.add(security_category_external_id_item)
+            security_category_external_id = item.get("securityCategoryNames")
+            if isinstance(security_category_external_id, str):
+                security_category_names.add(security_category_external_id)
+        self.client.lookup.security_categories.id(list(security_category_names))
+
+    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])
 
     @classmethod
-    def _create_if_not_exists(
-        cls,
-        filepaths: list[Path],
-        loader: DataSetsCRUD | LabelCRUD,
-        console: Console | None = None,
-    ) -> None:
-        items: LabelDefinitionWriteList | DataSetWriteList = loader.list_write_cls([])
-        for filepath in filepaths:
-            if not filepath.exists():
-                continue
-            for loaded in loader.load_resource_file(filepath):
-                items.append(loader.load_resource(loaded))
-        # MyPy fails to understand that existing, existing_ids and missing will be consistent for given loader.
-        existing = loader.retrieve(loader.get_ids(items))
-        existing_ids = set(loader.get_ids(existing))
-        if missing := [item for item in items if loader.get_id(item) not in existing_ids]:  # type: ignore[arg-type]
-            loader.create(loader.list_write_cls(missing))  # type: ignore[arg-type]
-            if console:
-                console.print(
-                    f"Created {loader.kind} for {len(missing)} items: {', '.join(str(item) for item in loader.get_ids(missing))}"  # type: ignore[arg-type]
-                )
+    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():
+            if key.startswith("metadata."):
+                metadata_key = key[len("metadata.") :]
+                metadata[metadata_key] = value
+            else:
+                cleaned_row[key] = value
+        if metadata:
+            cleaned_row["metadata"] = metadata
+        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]):
-    FOLDER_NAME = "classic"
+class AssetIO(UploadableAssetCentricIO[Asset, AssetWrite]):
     KIND = "Assets"
-    DISPLAY_NAME = "Assets"
+    RESOURCE_TYPE = "asset"
     SUPPORTED_DOWNLOAD_FORMATS = frozenset({".parquet", ".csv", ".ndjson"})
     SUPPORTED_COMPRESSIONS = frozenset({".gz"})
     SUPPORTED_READ_FORMATS = frozenset({".parquet", ".csv", ".ndjson", ".yaml", ".yml"})
+    UPLOAD_ENDPOINT = "/assets"
 
-    def as_id(self, item: dict[str, JsonVal] | object) -> int:
-        if isinstance(item, Asset | AssetWrite) and item.id is not None:  # type: ignore[union-attr]
-            return item.id  # type: ignore[union-attr]
-        return super().as_id(item)
+    def __init__(self, client: ToolkitClient) -> None:
+        super().__init__(client)
+        self._crud = AssetCRUD.create_loader(self.client)
 
-    def _get_loader(self) -> AssetCRUD:
-        return AssetCRUD.create_loader(self.client)
+    def as_id(self, item: Asset) -> str:
+        return item.external_id if item.external_id is not None else self._create_identifier(item.id)
 
     def _get_aggregator(self) -> AssetCentricAggregator:
         return AssetAggregator(self.client)
@@ -187,10 +270,18 @@ class AssetIO(BaseAssetCentricIO[str, AssetWrite, Asset, AssetWriteList, AssetLi
     def get_schema(self, selector: AssetCentricSelector) -> list[SchemaColumn]:
         data_set_ids: list[int] = []
         if isinstance(selector, DataSetSelector):
-            data_set_ids.append(self.client.lookup.data_sets.id(selector.data_set_external_id))
+            data_set_id = self.client.lookup.data_sets.id(selector.data_set_external_id)
+            if data_set_id is None:
+                raise ToolkitMissingResourceError(
+                    f"Data set with external ID {selector.data_set_external_id} not found."
+                )
+            data_set_ids.append(data_set_id)
         hierarchy: list[int] = []
         if isinstance(selector, AssetSubtreeSelector):
-            hierarchy.append(self.client.lookup.assets.id(selector.hierarchy))
+            asset_id = self.client.lookup.assets.id(selector.hierarchy)
+            if asset_id is None:
+                raise ToolkitMissingResourceError(f"Asset with external ID {selector.hierarchy} not found.")
+            hierarchy.append(asset_id)
 
         if hierarchy or data_set_ids:
             metadata_keys = metadata_key_counts(
@@ -212,55 +303,95 @@ class AssetIO(BaseAssetCentricIO[str, AssetWrite, Asset, AssetWriteList, AssetLi
             SchemaColumn(name="source", type="string"),
             SchemaColumn(name="labels", type="string", is_array=True),
             SchemaColumn(name="geoLocation", type="json"),
+            SchemaColumn(name="childCount", type="integer"),
+            SchemaColumn(name="depth", type="integer"),
+            SchemaColumn(name="path", type="string", is_array=True),
         ]
         return asset_schema + metadata_schema
 
-    def stream_data(self, selector: AssetCentricSelector, limit: int | None = None) -> Iterable[AssetList]:
-        asset_subtree_external_ids: list[str] | None = None
-        data_set_external_ids: list[str] | None = None
-        if isinstance(selector, DataSetSelector):
-            data_set_external_ids = [selector.data_set_external_id]
-        elif isinstance(selector, AssetSubtreeSelector):
-            asset_subtree_external_ids = [selector.hierarchy]
-        else:
-            # This selector is for uploads, not for downloading from CDF.
-            raise ToolkitNotImplementedError(f"Selector type {type(selector)} not supported for AssetIO.")
+    def stream_data(self, selector: AssetCentricSelector, limit: int | None = None) -> Iterable[Page]:
+        asset_subtree_external_ids, data_set_external_ids = self._get_hierarchy_dataset_pair(selector)
         for asset_list in self.client.assets(
             chunk_size=self.CHUNK_SIZE,
             limit=limit,
             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 asset_list
-
-    def upload_items(self, data_chunk: AssetWriteList, selector: AssetCentricSelector) -> None:
-        if not data_chunk:
-            return
-        self.client.assets.upsert(data_chunk, mode="patch")
-
-    def json_chunk_to_data(self, data_chunk: list[dict[str, JsonVal]]) -> AssetWriteList:
-        return AssetWriteList([self._loader.load_resource(item) for item in data_chunk])
+            yield Page(worker_id="main", items=asset_list)
+
+    def data_to_json_chunk(
+        self, data_chunk: Sequence[Asset], selector: AssetCentricSelector | None = None
+    ) -> list[dict[str, JsonVal]]:
+        # Ensure data sets are looked up to populate cache.
+        # This is to avoid looking up each data set id individually in the .dump_resource call.
+        self._populate_data_set_id_cache(data_chunk)
+        asset_ids = {
+            segment["id"]
+            for item in data_chunk
+            if isinstance(item.aggregates, AggregateResultItem)
+            for segment in item.aggregates.path or []
+            if "id" in segment
+        }
+        self.client.lookup.assets.external_id(list(asset_ids))
+
+        return [self._crud.dump_resource(item) for item in data_chunk]
+
+    def json_to_resource(self, item_json: dict[str, JsonVal]) -> AssetWrite:
+        return self._crud.load_resource(item_json)
 
     def retrieve(self, ids: Sequence[int]) -> AssetList:
         return self.client.assets.retrieve_multiple(ids)
 
-
-class FileMetadataIO(BaseAssetCentricIO[str, FileMetadataWrite, FileMetadata, FileMetadataWriteList, FileMetadataList]):
-    FOLDER_NAME = FileMetadataCRUD.folder_name
+    @classmethod
+    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"
+        batch: list[tuple[str, dict[str, JsonVal]]] = []
+        # We read the file multiple times, once for each depth level, to ensure parents are created before children.
+        while current_depth <= max_depth:
+            for line_number, item in reader.read_chunks_with_line_numbers():
+                try:
+                    depth = int(item["depth"])  # type: ignore[arg-type]
+                except (TypeError, ValueError, KeyError):
+                    if current_depth == 0:
+                        # If depth is not set, we yield it at depth 0
+                        batch.append((f"{data_name} {line_number}", item))
+                else:
+                    if depth == current_depth:
+                        batch.append((f"{data_name} {line_number}", item))
+                    elif current_depth == 0:
+                        max_depth = max(max_depth, depth)
+                if len(batch) >= cls.CHUNK_SIZE:
+                    yield batch
+                    batch = []
+            if batch:
+                yield batch
+                batch = []
+            current_depth += 1
+
+
+class FileMetadataIO(AssetCentricIO[FileMetadata]):
     KIND = "FileMetadata"
-    DISPLAY_NAME = "file metadata"
+    RESOURCE_TYPE = "file"
     SUPPORTED_DOWNLOAD_FORMATS = frozenset({".parquet", ".csv", ".ndjson"})
     SUPPORTED_COMPRESSIONS = frozenset({".gz"})
     SUPPORTED_READ_FORMATS = frozenset({".parquet", ".csv", ".ndjson"})
+    UPLOAD_ENDPOINT = "/files"
 
-    def as_id(self, item: dict[str, JsonVal] | object) -> int:
-        if isinstance(item, FileMetadata | FileMetadataWrite) and item.id is not None:  # type: ignore[union-attr]
-            return item.id  # type: ignore[union-attr]
-        return super().as_id(item)
+    def __init__(self, client: ToolkitClient) -> None:
+        super().__init__(client)
+        self._crud = FileMetadataCRUD.create_loader(self.client)
 
-    def _get_loader(self) -> FileMetadataCRUD:
-        return FileMetadataCRUD.create_loader(self.client)
+    def as_id(self, item: FileMetadata) -> str:
+        return item.external_id if item.external_id is not None else self._create_identifier(item.id)
 
     def _get_aggregator(self) -> AssetCentricAggregator:
         return FileAggregator(self.client)
@@ -268,7 +399,12 @@ class FileMetadataIO(BaseAssetCentricIO[str, FileMetadataWrite, FileMetadata, Fi
     def get_schema(self, selector: AssetCentricSelector) -> list[SchemaColumn]:
         data_set_ids: list[int] = []
         if isinstance(selector, DataSetSelector):
-            data_set_ids.append(self.client.lookup.data_sets.id(selector.data_set_external_id))
+            data_set_id = self.client.lookup.data_sets.id(selector.data_set_external_id)
+            if data_set_id is None:
+                raise ToolkitMissingResourceError(
+                    f"Data set with external ID {selector.data_set_external_id} not found."
+                )
+            data_set_ids.append(data_set_id)
         if isinstance(selector, AssetSubtreeSelector):
             raise ToolkitNotImplementedError(f"Selector type {type(selector)} not supported for FileIO.")
 
@@ -297,32 +433,261 @@ class FileMetadataIO(BaseAssetCentricIO[str, FileMetadataWrite, FileMetadata, Fi
         ]
         return file_schema + metadata_schema
 
-    def stream_data(self, selector: AssetCentricSelector, limit: int | None = None) -> Iterable[FileMetadataList]:
-        asset_subtree_external_ids: list[str] | None = None
-        data_set_external_ids: list[str] | None = None
-        if isinstance(selector, DataSetSelector):
-            data_set_external_ids = [selector.data_set_external_id]
-        elif isinstance(selector, AssetSubtreeSelector):
-            asset_subtree_external_ids = [selector.hierarchy]
-        else:
-            # This selector is for uploads, not for downloading from CDF.
-            raise ToolkitNotImplementedError(f"Selector type {type(selector)} not supported for FileMetadataIO.")
+    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 file_list
+            yield Page(worker_id="main", items=file_list)
 
     def retrieve(self, ids: Sequence[int]) -> FileMetadataList:
         return self.client.files.retrieve_multiple(ids)
 
-    def upload_items(self, data_chunk: FileMetadataWriteList, selector: AssetCentricSelector) -> None:
-        if not data_chunk:
-            return
-        self._loader.create(data_chunk)
+    def data_to_json_chunk(
+        self, data_chunk: Sequence[FileMetadata], selector: AssetCentricSelector | None = None
+    ) -> list[dict[str, JsonVal]]:
+        # Ensure data sets/assets/security-categories are looked up to populate cache.
+        # This is to avoid looking up each data set id individually in the .dump_resource call
+        self._populate_data_set_id_cache(data_chunk)
+        self._populate_asset_id_cache(data_chunk)
+        self._populate_security_category_cache(data_chunk)
+
+        return [self._crud.dump_resource(item) for item in data_chunk]
+
+
+class TimeSeriesIO(UploadableAssetCentricIO[TimeSeries, TimeSeriesWrite]):
+    KIND = "TimeSeries"
+    SUPPORTED_DOWNLOAD_FORMATS = frozenset({".parquet", ".csv", ".ndjson"})
+    SUPPORTED_COMPRESSIONS = frozenset({".gz"})
+    SUPPORTED_READ_FORMATS = frozenset({".parquet", ".csv", ".ndjson"})
+    UPLOAD_ENDPOINT = "/timeseries"
+    RESOURCE_TYPE = "timeseries"
+
+    def __init__(self, client: ToolkitClient) -> None:
+        super().__init__(client)
+        self._crud = TimeSeriesCRUD.create_loader(self.client)
+
+    def as_id(self, item: TimeSeries) -> str:
+        return item.external_id if item.external_id is not None else self._create_identifier(item.id)
+
+    def _get_aggregator(self) -> AssetCentricAggregator:
+        return TimeSeriesAggregator(self.client)
+
+    def retrieve(self, ids: Sequence[int]) -> TimeSeriesList:
+        return self.client.time_series.retrieve_multiple(ids=ids)
+
+    def stream_data(self, selector: AssetCentricSelector, limit: int | None = None) -> Iterable[Page]:
+        asset_subtree_external_ids, data_set_external_ids = self._get_hierarchy_dataset_pair(selector)
+        for ts_list in self.client.time_series(
+            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(ts_list, selector)
+            yield Page(worker_id="main", items=ts_list)
+
+    def data_to_json_chunk(
+        self, data_chunk: Sequence[TimeSeries], selector: AssetCentricSelector | None = None
+    ) -> list[dict[str, JsonVal]]:
+        # Ensure data sets/assets/security categories are looked up to populate cache.
+        self._populate_data_set_id_cache(data_chunk)
+        self._populate_security_category_cache(data_chunk)
+        asset_ids = {item.asset_id for item in data_chunk if item.asset_id is not None}
+        self.client.lookup.assets.external_id(list(asset_ids))
+
+        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[TimeSeriesWrite]]:
+        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]) -> TimeSeriesWrite:
+        return self._crud.load_resource(item_json)
+
+    def get_schema(self, selector: AssetCentricSelector) -> list[SchemaColumn]:
+        data_set_ids: list[int] = []
+        if isinstance(selector, DataSetSelector):
+            data_set_id = self.client.lookup.data_sets.id(selector.data_set_external_id)
+            if data_set_id is None:
+                raise ToolkitMissingResourceError(
+                    f"Data set with external ID {selector.data_set_external_id} not found."
+                )
+            data_set_ids.append(data_set_id)
+        elif isinstance(selector, AssetSubtreeSelector):
+            raise ToolkitNotImplementedError(f"Selector type {type(selector)} not supported for {type(self).__name__}.")
+
+        if data_set_ids:
+            metadata_keys = metadata_key_counts(
+                self.client, "timeseries", data_sets=data_set_ids or None, hierarchies=None
+            )
+        else:
+            metadata_keys = []
+        metadata_schema: list[SchemaColumn] = []
+        if metadata_keys:
+            metadata_schema.extend(
+                [SchemaColumn(name=f"metadata.{key}", type="string", is_array=False) for key, _ in metadata_keys]
+            )
+        ts_schema = [
+            SchemaColumn(name="externalId", type="string"),
+            SchemaColumn(name="name", type="string"),
+            SchemaColumn(name="isString", type="boolean"),
+            SchemaColumn(name="unit", type="string"),
+            SchemaColumn(name="unitExternalId", type="string"),
+            SchemaColumn(name="assetExternalId", type="string"),
+            SchemaColumn(name="isStep", type="boolean"),
+            SchemaColumn(name="description", type="string"),
+            SchemaColumn(name="securityCategories", type="string", is_array=True),
+            SchemaColumn(name="dataSetExternalId", type="string"),
+        ]
+        return ts_schema + metadata_schema
+
+
+class EventIO(UploadableAssetCentricIO[Event, EventWrite]):
+    KIND = "Events"
+    SUPPORTED_DOWNLOAD_FORMATS = frozenset({".parquet", ".csv", ".ndjson"})
+    SUPPORTED_COMPRESSIONS = frozenset({".gz"})
+    SUPPORTED_READ_FORMATS = frozenset({".parquet", ".csv", ".ndjson"})
+    UPLOAD_ENDPOINT = "/events"
+    RESOURCE_TYPE = "event"
+
+    def __init__(self, client: ToolkitClient) -> None:
+        super().__init__(client)
+        self._crud = EventCRUD.create_loader(self.client)
+
+    def as_id(self, item: Event) -> str:
+        return item.external_id if item.external_id is not None else self._create_identifier(item.id)
+
+    def _get_aggregator(self) -> AssetCentricAggregator:
+        return EventAggregator(self.client)
+
+    def get_schema(self, selector: AssetCentricSelector) -> list[SchemaColumn]:
+        data_set_ids: list[int] = []
+        if isinstance(selector, DataSetSelector):
+            data_set_id = self.client.lookup.data_sets.id(selector.data_set_external_id)
+            if data_set_id is None:
+                raise ToolkitMissingResourceError(
+                    f"Data set with external ID {selector.data_set_external_id} not found."
+                )
+            data_set_ids.append(data_set_id)
+        hierarchy: list[int] = []
+        if isinstance(selector, AssetSubtreeSelector):
+            raise ToolkitNotImplementedError(f"Selector type {type(selector)} not supported for {type(self).__name__}.")
+
+        if hierarchy or data_set_ids:
+            metadata_keys = metadata_key_counts(
+                self.client, "events", data_sets=data_set_ids or None, hierarchies=hierarchy or None
+            )
+        else:
+            metadata_keys = []
+        metadata_schema: list[SchemaColumn] = []
+        if metadata_keys:
+            metadata_schema.extend(
+                [SchemaColumn(name=f"metadata.{key}", type="string", is_array=False) for key, _ in metadata_keys]
+            )
+        event_schema = [
+            SchemaColumn(name="externalId", type="string"),
+            SchemaColumn(name="dataSetExternalId", type="string"),
+            SchemaColumn(name="startTime", type="integer"),
+            SchemaColumn(name="endTime", type="integer"),
+            SchemaColumn(name="type", type="string"),
+            SchemaColumn(name="subtype", type="string"),
+            SchemaColumn(name="description", type="string"),
+            SchemaColumn(name="assetExternalIds", type="string", is_array=True),
+            SchemaColumn(name="source", type="string"),
+        ]
+        return event_schema + metadata_schema
+
+    def stream_data(self, selector: AssetCentricSelector, limit: int | None = None) -> Iterable[Page]:
+        asset_subtree_external_ids, data_set_external_ids = self._get_hierarchy_dataset_pair(selector)
+        for event_list in self.client.events(
+            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(event_list, selector)
+            yield Page(worker_id="main", items=event_list)
+
+    def data_to_json_chunk(
+        self, data_chunk: Sequence[Event], selector: AssetCentricSelector | None = None
+    ) -> list[dict[str, JsonVal]]:
+        # Ensure data sets/assets are looked up to populate cache.
+        self._populate_data_set_id_cache(data_chunk)
+        self._populate_asset_id_cache(data_chunk)
+
+        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[EventWrite]]:
+        chunks = [item_json for _, item_json in data_chunk]
+        self._populate_asset_external_ids_cache(chunks)
+        self._populate_data_set_external_id_cache(chunks)
+        return super().json_chunk_to_data(data_chunk)
+
+    def json_to_resource(self, item_json: dict[str, JsonVal]) -> EventWrite:
+        return self._crud.load_resource(item_json)
+
+    def retrieve(self, ids: Sequence[int]) -> EventList:
+        return self.client.events.retrieve_multiple(ids)
+
+
+class HierarchyIO(ConfigurableStorageIO[AssetCentricSelector, AssetCentricResource]):
+    CHUNK_SIZE = 1000
+    BASE_SELECTOR = AssetCentricSelector
+    SUPPORTED_DOWNLOAD_FORMATS = frozenset({".parquet", ".csv", ".ndjson"})
+    SUPPORTED_COMPRESSIONS = frozenset({".gz"})
+
+    def __init__(self, client: ToolkitClient) -> None:
+        super().__init__(client)
+        self._asset_io = AssetIO(client)
+        self._file_io = FileMetadataIO(client)
+        self._timeseries_io = TimeSeriesIO(client)
+        self._event_io = EventIO(client)
+        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,
+            self._event_io.KIND: self._event_io,
+        }
+
+    def as_id(self, item: AssetCentricResource) -> str:
+        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
+    ) -> Iterable[Page[AssetCentricResource]]:
+        yield from self.get_resource_io(selector.kind).stream_data(selector, limit)
+
+    def count(self, selector: AssetCentricSelector) -> int | None:
+        return self.get_resource_io(selector.kind).count(selector)
+
+    def data_to_json_chunk(
+        self, data_chunk: Sequence[AssetCentricResource], selector: AssetCentricSelector | None = None
+    ) -> list[dict[str, JsonVal]]:
+        if selector is None:
+            raise ValueError(f"Selector must be provided to convert data to JSON chunk for {type(self).__name__}.)")
+        return self.get_resource_io(selector.kind).data_to_json_chunk(data_chunk, selector)
+
+    def configurations(self, selector: AssetCentricSelector) -> Iterable[StorageIOConfig]:
+        yield from self.get_resource_io(selector.kind).configurations(selector)
 
-    def json_chunk_to_data(self, data_chunk: list[dict[str, JsonVal]]) -> FileMetadataWriteList:
-        return FileMetadataWriteList([self._loader.load_resource(item) for item in data_chunk])
+    def get_resource_io(self, kind: str) -> AssetCentricIO:
+        return self._io_by_kind[kind]