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

cognite_toolkit/_cdf_tk/storageio/_instances.py

@@ -1,25 +1,49 @@
-from collections.abc import Iterable, Mapping
-from pathlib import Path
+from collections.abc import Iterable, Mapping, Sequence
 from types import MappingProxyType
 from typing import ClassVar
 
 from cognite.client.data_classes.aggregations import Count
-from cognite.client.data_classes.data_modeling import Edge, EdgeApply, Node, NodeApply
+from cognite.client.data_classes.data_modeling import (
+    ContainerId,
+    ContainerList,
+    EdgeApply,
+    NodeApply,
+    SpaceList,
+    ViewId,
+    ViewList,
+)
+from cognite.client.data_classes.data_modeling.instances import Instance, InstanceApply
 from cognite.client.utils._identifier import InstanceId
-from rich.console import Console
 
-from cognite_toolkit._cdf_tk.client.data_classes.instances import InstanceApplyList, InstanceList
+from cognite_toolkit._cdf_tk import constants
+from cognite_toolkit._cdf_tk.client import ToolkitClient
+from cognite_toolkit._cdf_tk.client.data_classes.instances import InstanceList
+from cognite_toolkit._cdf_tk.cruds import ContainerCRUD, SpaceCRUD, ViewCRUD
+from cognite_toolkit._cdf_tk.utils import sanitize_filename
 from cognite_toolkit._cdf_tk.utils.cdf import iterate_instances
 from cognite_toolkit._cdf_tk.utils.collection import chunker_sequence
-from cognite_toolkit._cdf_tk.utils.fileio import SchemaColumn
 from cognite_toolkit._cdf_tk.utils.useful_types import JsonVal
 
-from ._base import StorageIOConfig, TableStorageIO
-from ._selectors import InstanceFileSelector, InstanceSelector, InstanceViewSelector
+from . import StorageIOConfig
+from ._base import ConfigurableStorageIO, Page, UploadableStorageIO
+from .selectors import InstanceFileSelector, InstanceSelector, InstanceSpaceSelector, InstanceViewSelector
 
 
-class InstanceIO(TableStorageIO[InstanceId, InstanceSelector, InstanceApplyList, InstanceList]):
-    FOLDER_NAME = "instances"
+class InstanceIO(
+    ConfigurableStorageIO[InstanceSelector, Instance],
+    UploadableStorageIO[InstanceSelector, Instance, InstanceApply],
+):
+    """This class provides functionality to interact with instances in Cognite Data Fusion (CDF).
+
+    It is used to download, upload, and purge instances, as well as spaces,views, and containers related to instances.
+
+    Args:
+        client (ToolkitClient): An instance of ToolkitClient to interact with the CDF API.
+        remove_existing_version (bool): Whether to remove existing versions of instances during upload.
+            Default is True. Existing versions are used to safeguard against accidental overwrites.
+
+    """
+
     KIND = "Instances"
     DISPLAY_NAME = "Instances"
     SUPPORTED_DOWNLOAD_FORMATS = frozenset({".parquet", ".csv", ".ndjson"})
@@ -28,83 +52,174 @@ class InstanceIO(TableStorageIO[InstanceId, InstanceSelector, InstanceApplyList,
     CHUNK_SIZE = 1000
     UPLOAD_ENDPOINT = "/models/instances"
     UPLOAD_EXTRA_ARGS: ClassVar[Mapping[str, JsonVal] | None] = MappingProxyType({"autoCreateDirectRelations": True})
-
-    def as_id(self, item: dict[str, JsonVal] | object) -> InstanceId:
-        if isinstance(item, dict) and isinstance(item.get("space"), str) and isinstance(item.get("externalId"), str):
-            # MyPy checked above.
-            return InstanceId(space=item["space"], external_id=item["externalId"])  # type: ignore[arg-type]
-        if isinstance(item, InstanceId):
-            return item
-        elif isinstance(item, Node | Edge | NodeApply | EdgeApply):
-            return item.as_id()
-        raise TypeError(f"Cannot extract ID from item of type {type(item).__name__!r}")
-
-    def stream_data(self, selector: InstanceSelector, limit: int | None = None) -> Iterable[InstanceList]:
-        if isinstance(selector, InstanceViewSelector):
+    BASE_SELECTOR = InstanceSelector
+
+    def __init__(self, client: ToolkitClient, remove_existing_version: bool = True) -> None:
+        super().__init__(client)
+        self._remove_existing_version = remove_existing_version
+        # Cache for view to read-only properties mapping
+        self._view_readonly_properties_cache: dict[ViewId, set[str]] = {}
+        self._view_crud = ViewCRUD.create_loader(self.client)
+
+    def as_id(self, item: Instance) -> str:
+        return f"{item.space}:{item.external_id}"
+
+    def _filter_readonly_properties(self, instance: InstanceApply) -> None:
+        """
+        Filter out read-only properties from the instance.
+
+        Args:
+            instance: The instance to filter readonly properties from
+        """
+
+        for source in instance.sources:
+            readonly_properties = set()
+            if isinstance(source.source, ViewId):
+                if source.source not in self._view_readonly_properties_cache:
+                    self._view_readonly_properties_cache[source.source] = self._view_crud.get_readonly_properties(
+                        source.source
+                    )
+                readonly_properties = self._view_readonly_properties_cache[source.source]
+            elif isinstance(source.source, ContainerId):
+                if source.source in constants.READONLY_CONTAINER_PROPERTIES:
+                    readonly_properties = constants.READONLY_CONTAINER_PROPERTIES[source.source]
+
+            source.properties = {k: v for k, v in source.properties.items() if k not in readonly_properties}
+
+    def stream_data(self, selector: InstanceSelector, limit: int | None = None) -> Iterable[Page]:
+        if isinstance(selector, InstanceViewSelector | InstanceSpaceSelector):
             chunk = InstanceList([])
             total = 0
-            for instance in iterate_instances(
-                client=self.client,
-                source=selector.view,
-                instance_type=selector.instance_type,
-                space=list(selector.instance_spaces) if selector.instance_spaces else None,
-            ):
+            for instance in iterate_instances(client=self.client, **selector.as_filter_args()):
                 if limit is not None and total >= limit:
                     break
                 total += 1
                 chunk.append(instance)
                 if len(chunk) >= self.CHUNK_SIZE:
-                    yield chunk
+                    yield Page(worker_id="main", items=chunk)
                     chunk = InstanceList([])
             if chunk:
-                yield chunk
+                yield Page(worker_id="main", items=chunk)
         elif isinstance(selector, InstanceFileSelector):
             for node_chunk in chunker_sequence(selector.node_ids, self.CHUNK_SIZE):
-                yield InstanceList(self.client.data_modeling.instances.retrieve(nodes=node_chunk).nodes)
+                yield Page(
+                    worker_id="main",
+                    items=InstanceList(self.client.data_modeling.instances.retrieve(nodes=node_chunk).nodes),
+                )
             for edge_chunk in chunker_sequence(selector.edge_ids, self.CHUNK_SIZE):
-                yield InstanceList(self.client.data_modeling.instances.retrieve(edges=edge_chunk).edges)
+                yield Page(
+                    worker_id="main",
+                    items=InstanceList(self.client.data_modeling.instances.retrieve(edges=edge_chunk).edges),
+                )
         else:
             raise NotImplementedError()
 
     def download_ids(self, selector: InstanceSelector, limit: int | None = None) -> Iterable[list[InstanceId]]:
-        if isinstance(selector, InstanceFileSelector) and selector.validate is False:
+        if isinstance(selector, InstanceFileSelector) and selector.validate_instance is False:
             instances_to_yield = selector.instance_ids
             if limit is not None:
                 instances_to_yield = instances_to_yield[:limit]
             yield from chunker_sequence(instances_to_yield, self.CHUNK_SIZE)
         else:
-            yield from ([instance.as_id() for instance in chunk] for chunk in self.stream_data(selector, limit))  # type: ignore[attr-defined]
+            yield from ([instance.as_id() for instance in chunk.items] for chunk in self.stream_data(selector, limit))
 
     def count(self, selector: InstanceSelector) -> int | None:
-        if isinstance(selector, InstanceViewSelector):
+        if isinstance(selector, InstanceViewSelector) or (
+            isinstance(selector, InstanceSpaceSelector) and selector.view
+        ):
             result = self.client.data_modeling.instances.aggregate(
-                view=selector.view,
+                # MyPy do not understand that selector.view is always defined here.
+                view=selector.view.as_id(),  # type: ignore[union-attr]
                 aggregates=Count("externalId"),
                 instance_type=selector.instance_type,
-                space=list(selector.instance_spaces) if selector.instance_spaces else None,
+                space=selector.get_instance_spaces(),
             )
             return int(result.value or 0)
+        elif isinstance(selector, InstanceSpaceSelector):
+            statistics = self.client.data_modeling.statistics.spaces.retrieve(space=selector.instance_space)
+            if statistics is None:
+                return None
+            if selector.instance_type == "node":
+                return statistics.nodes
+            elif selector.instance_type == "edge":
+                return statistics.edges
+            # This should never happen due to validation in the selector.
+            raise ValueError(f"Unknown instance type {selector.instance_type!r}")
         elif isinstance(selector, InstanceFileSelector):
             return len(selector.instance_ids)
         raise NotImplementedError()
 
-    def upload_items(self, data_chunk: InstanceApplyList, selector: InstanceSelector) -> None:
-        raise NotImplementedError()
-
-    def data_to_json_chunk(self, data_chunk: InstanceList) -> list[dict[str, JsonVal]]:
-        raise NotImplementedError()
-
-    def json_chunk_to_data(self, data_chunk: list[dict[str, JsonVal]]) -> InstanceApplyList:
-        raise NotImplementedError()
+    def data_to_json_chunk(
+        self, data_chunk: Sequence[Instance], selector: InstanceSelector | None = None
+    ) -> list[dict[str, JsonVal]]:
+        return [instance.as_write().dump(camel_case=True) for instance in data_chunk]
+
+    def json_to_resource(self, item_json: dict[str, JsonVal]) -> InstanceApply:
+        # There is a bug in the SDK where InstanceApply._load turns all keys to snake_case.
+        # So we cannot use InstanceApplyList._load here.
+        instance_type = item_json.get("instanceType")
+        item_to_load = dict(item_json)
+        if self._remove_existing_version and "existingVersion" in item_to_load:
+            del item_to_load["existingVersion"]
+        instance: InstanceApply
+        if instance_type == "node":
+            instance = NodeApply._load(item_to_load, cognite_client=self.client)
+        elif instance_type == "edge":
+            instance = EdgeApply._load(item_to_load, cognite_client=self.client)
+        else:
+            raise ValueError(f"Unknown instance type {instance_type!r}")
+        # Filter out read-only properties if applicable
+        self._filter_readonly_properties(instance)
+        return instance
 
     def configurations(self, selector: InstanceSelector) -> Iterable[StorageIOConfig]:
-        raise NotImplementedError()
-
-    def load_selector(self, datafile: Path) -> InstanceSelector:
-        raise NotImplementedError()
-
-    def ensure_configurations(self, selector: InstanceSelector, console: Console | None = None) -> None:
-        raise NotImplementedError()
-
-    def get_schema(self, selector: InstanceSelector) -> list[SchemaColumn]:
-        raise NotImplementedError()
+        if not isinstance(selector, InstanceViewSelector | InstanceSpaceSelector):
+            return
+        spaces = list(set((selector.get_instance_spaces() or []) + (selector.get_schema_spaces() or [])))
+        if not spaces:
+            return
+        space_crud = SpaceCRUD.create_loader(self.client)
+        retrieved_spaces = space_crud.retrieve(spaces)
+        retrieved_spaces = SpaceList([space for space in retrieved_spaces if not space.is_global])
+        if not retrieved_spaces:
+            return
+        for space in retrieved_spaces:
+            yield StorageIOConfig(
+                kind=SpaceCRUD.kind,
+                folder_name=SpaceCRUD.folder_name,
+                value=space_crud.dump_resource(space),
+                filename=sanitize_filename(space.space),
+            )
+        if not selector.view:
+            return
+        view_id = selector.view.as_id()
+        view_crud = ViewCRUD(self.client, None, None, topological_sort_implements=True)
+        views = view_crud.retrieve([view_id])
+        views = ViewList([view for view in views if not view.is_global])
+        if not views:
+            return
+        for view in views:
+            filename = f"{view.space}_{view.external_id}"
+            if view.version is not None:
+                filename += f"_{view.version}"
+            yield StorageIOConfig(
+                kind=ViewCRUD.kind,
+                folder_name=ViewCRUD.folder_name,
+                value=view_crud.dump_resource(view),
+                filename=sanitize_filename(filename),
+            )
+        container_ids = list({container for view in views for container in view.referenced_containers() or []})
+        if not container_ids:
+            return
+        container_crud = ContainerCRUD.create_loader(self.client)
+        containers = container_crud.retrieve(container_ids)
+        containers = ContainerList([container for container in containers if not container.is_global])
+        if not containers:
+            return
+        for container in containers:
+            yield StorageIOConfig(
+                kind=ContainerCRUD.kind,
+                folder_name=ContainerCRUD.folder_name,
+                value=container_crud.dump_resource(container),
+                filename=sanitize_filename(f"{container.space}_{container.external_id}"),
+            )

cognite_toolkit/_cdf_tk/storageio/_raw.py

@@ -1,79 +1,124 @@
-from collections.abc import Iterable
-from pathlib import Path
+from collections.abc import Iterable, Sequence
+from itertools import chain
+from uuid import uuid4
 
-from cognite.client.data_classes import RowList, RowWrite, RowWriteList
-from rich.console import Console
+from cognite.client.data_classes import Row, RowWrite
 
-from cognite_toolkit._cdf_tk.client.data_classes.raw import RawDatabase, RawDatabaseList, RawTable, RawTableList
 from cognite_toolkit._cdf_tk.cruds import RawDatabaseCRUD, RawTableCRUD
 from cognite_toolkit._cdf_tk.exceptions import ToolkitValueError
-from cognite_toolkit._cdf_tk.utils.file import find_adjacent_files, read_yaml_file
+from cognite_toolkit._cdf_tk.utils import sanitize_filename
+from cognite_toolkit._cdf_tk.utils.collection import chunker
+from cognite_toolkit._cdf_tk.utils.fileio import MultiFileReader
+from cognite_toolkit._cdf_tk.utils.http_client import HTTPClient, HTTPMessage, ItemsRequest
 from cognite_toolkit._cdf_tk.utils.useful_types import JsonVal
 
-from ._base import StorageIO, StorageIOConfig
+from ._base import (
+    ConfigurableStorageIO,
+    Page,
+    StorageIOConfig,
+    TableUploadableStorageIO,
+    UploadItem,
+)
+from .selectors import RawTableSelector
 
 
-class RawIO(StorageIO[RawTable, RawTable, RowWriteList, RowList]):
-    FOLDER_NAME = "raw"
+class RawIO(
+    ConfigurableStorageIO[RawTableSelector, Row],
+    TableUploadableStorageIO[RawTableSelector, Row, RowWrite],
+):
     KIND = "RawRows"
     DISPLAY_NAME = "Raw Rows"
     SUPPORTED_DOWNLOAD_FORMATS = frozenset({".yaml", ".ndjson"})
     SUPPORTED_COMPRESSIONS = frozenset({".gz"})
     SUPPORTED_READ_FORMATS = frozenset({".parquet", ".csv", ".ndjson", ".yaml"})
     CHUNK_SIZE = 10_000
+    UPLOAD_ENDPOINT = "/raw/dbs/{dbName}/tables/{tableName}/rows"
+    BASE_SELECTOR = RawTableSelector
 
-    def as_id(self, item: dict[str, JsonVal] | object) -> RawTable:
-        raise ValueError("You cannot extract an ID from a Raw Table row. Use a RawTable selector instead.")
+    def as_id(self, item: Row) -> str:
+        return str(item.key)
 
-    def count(self, selector: RawTable) -> int | None:
+    def count(self, selector: RawTableSelector) -> int | None:
         # Raw tables do not support aggregation queries, so we do not know the count
         # up front.
         return None
 
-    def stream_data(self, selector: RawTable, limit: int | None = None) -> Iterable[RowList]:
-        yield from self.client.raw.rows(
-            db_name=selector.db_name,
-            table_name=selector.table_name,
+    def stream_data(self, selector: RawTableSelector, limit: int | None = None) -> Iterable[Page]:
+        for chunk in self.client.raw.rows(
+            db_name=selector.table.db_name,
+            table_name=selector.table.table_name,
             limit=limit,
             # 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,
             chunk_size=self.CHUNK_SIZE,
-        )
+        ):
+            yield Page(worker_id="main", items=chunk)
 
-    def upload_items(self, data_chunk: RowWriteList, selector: RawTable) -> None:
-        self.client.raw.rows.insert(db_name=selector.db_name, table_name=selector.table_name, row=data_chunk)
+    def upload_items(
+        self,
+        data_chunk: Sequence[UploadItem[RowWrite]],
+        http_client: HTTPClient,
+        selector: RawTableSelector | None = None,
+    ) -> Sequence[HTTPMessage]:
+        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(
+                endpoint_url=config.create_api_url(url),
+                method="POST",
+                items=list(data_chunk),
+            )
+        )
 
-    def data_to_json_chunk(self, data_chunk: RowList) -> list[dict[str, JsonVal]]:
+    def data_to_json_chunk(
+        self, data_chunk: Sequence[Row], selector: RawTableSelector | None = None
+    ) -> list[dict[str, JsonVal]]:
         return [row.as_write().dump() for row in data_chunk]
 
-    def json_chunk_to_data(self, data_chunk: list[dict[str, JsonVal]]) -> RowWriteList:
-        return RowWriteList([RowWrite._load(row) for row in data_chunk])
-
-    def configurations(self, selector: RawTable) -> Iterable[StorageIOConfig]:
-        yield StorageIOConfig(kind=RawTableCRUD.kind, folder_name=RawTableCRUD.folder_name, value=selector.dump())
+    def json_to_resource(self, item_json: dict[str, JsonVal]) -> RowWrite:
+        return RowWrite._load(item_json)
 
-    def load_selector(self, datafile: Path) -> RawTable:
-        config_files = find_adjacent_files(datafile, suffix=f".{RawTableCRUD.kind}.yaml")
-        if not config_files:
-            raise ToolkitValueError(f"No configuration file found for {datafile.as_posix()!r}")
-        if len(config_files) > 1:
-            raise ToolkitValueError(f"Multiple configuration files found for {datafile.as_posix()!r}: {config_files}")
-        config_file = config_files[0]
-        loader = RawTableCRUD.create_loader(self.client)
-        return loader.load_resource(read_yaml_file(config_file, expected_output="dict"))
+    def configurations(self, selector: RawTableSelector) -> Iterable[StorageIOConfig]:
+        yield StorageIOConfig(
+            kind=RawDatabaseCRUD.kind,
+            folder_name=RawDatabaseCRUD.folder_name,
+            value={"dbName": selector.table.db_name},
+            filename=sanitize_filename(selector.table.db_name),
+        )
+        yield StorageIOConfig(
+            kind=RawTableCRUD.kind, folder_name=RawTableCRUD.folder_name, value=selector.table.model_dump(by_alias=True)
+        )
 
-    def ensure_configurations(self, selector: RawTable, console: Console | None = None) -> None:
-        """Ensure that the Raw table exists in CDF."""
-        db_loader = RawDatabaseCRUD.create_loader(self.client, console=console)
-        db = RawDatabase(db_name=selector.db_name)
-        if not db_loader.retrieve([db]):
-            db_loader.create(RawDatabaseList([db]))
-            if console:
-                console.print(f"Created raw database: [bold]{db.db_name}[/bold]")
+    def row_to_resource(
+        self, source_id: str, row: dict[str, JsonVal], selector: RawTableSelector | None = None
+    ) -> RowWrite:
+        key = str(uuid4())
+        if selector is not None and selector.key is not None and selector.key in row:
+            key = str(row.pop(selector.key))
+        return RowWrite(key=key, columns=row)
 
-        table_loader = RawTableCRUD.create_loader(self.client, console=console)
-        if not table_loader.retrieve([selector]):
-            table_loader.create(RawTableList([selector]))
-            if console:
-                console.print(f"Created raw table: [bold]{selector.db_name}.{selector.table_name}[/bold]")
+    @classmethod
+    def read_chunks(
+        cls, reader: MultiFileReader, selector: RawTableSelector
+    ) -> Iterable[list[tuple[str, dict[str, JsonVal]]]]:
+        if not reader.is_table or selector.key is None:
+            yield from super().read_chunks(reader, selector)
+            return
+        data_name = "row" if reader.is_table else "line"
+        # Validate that the key exists in all files
+        for input_file in sorted(reader.input_files, key=reader._part_no):
+            iterable = reader.reader_class(input_file).read_chunks()
+            try:
+                first = next(iterable)
+            except StopIteration:
+                continue
+            if selector.key not in first:
+                raise ToolkitValueError(
+                    f"Column '{selector.key}' not found in file {input_file.as_posix()!r}. Please ensure the specified column exists."
+                )
+            full_iterator = chain([first], iterable)
+            line_numbered_iterator = ((f"{data_name} {i}", row) for i, row in enumerate(full_iterator, start=1))
+            yield from chunker(line_numbered_iterator, cls.CHUNK_SIZE)

cognite_toolkit/_cdf_tk/storageio/selectors/__init__.py

@@ -0,0 +1,124 @@
+from pathlib import Path
+from typing import Annotated
+
+from pydantic import Field, TypeAdapter, ValidationError
+
+from cognite_toolkit._cdf_tk.feature_flags import Flags
+from cognite_toolkit._cdf_tk.tk_warnings import MediumSeverityWarning, ToolkitWarning
+from cognite_toolkit._cdf_tk.tk_warnings.fileread import ResourceFormatWarning
+from cognite_toolkit._cdf_tk.utils.file import read_yaml_file
+from cognite_toolkit._cdf_tk.validation import humanize_validation_error
+
+from ._asset_centric import AssetCentricFileSelector, AssetCentricSelector, AssetSubtreeSelector, DataSetSelector
+from ._base import DataSelector
+from ._canvas import CanvasExternalIdSelector, CanvasSelector
+from ._charts import AllChartsSelector, ChartExternalIdSelector, ChartOwnerSelector, ChartSelector
+from ._datapoints import (
+    DataPointsDataSetSelector,
+    DataPointsFileSelector,
+    DataPointsSelector,
+    ExternalIdColumn,
+    InstanceColumn,
+    InternalIdColumn,
+    TimeSeriesColumn,
+)
+from ._file_content import (
+    FileContentSelector,
+    FileDataModelingTemplate,
+    FileDataModelingTemplateSelector,
+    FileIdentifierSelector,
+    FileMetadataTemplate,
+    FileMetadataTemplateSelector,
+)
+from ._instances import (
+    InstanceFileSelector,
+    InstanceSelector,
+    InstanceSpaceSelector,
+    InstanceViewSelector,
+    SelectedView,
+)
+from ._raw import RawTableSelector, SelectedTable
+
+Selector = Annotated[
+    RawTableSelector
+    | InstanceViewSelector
+    | InstanceFileSelector
+    | InstanceSpaceSelector
+    | AllChartsSelector
+    | ChartOwnerSelector
+    | AssetSubtreeSelector
+    | AssetCentricFileSelector
+    | DataSetSelector
+    | DataPointsDataSetSelector
+    | DataPointsFileSelector
+    | ChartExternalIdSelector
+    | CanvasExternalIdSelector
+    | FileMetadataTemplateSelector
+    | FileDataModelingTemplateSelector
+    | FileIdentifierSelector,
+    Field(discriminator="type"),
+]
+
+ALPHA_SELECTORS = {FileIdentifierSelector}
+
+SelectorAdapter: TypeAdapter[Selector] = TypeAdapter(Selector)
+
+
+def load_selector(manifest_file: Path) -> Selector | ToolkitWarning:
+    """Loads a selector from a manifest file.
+
+    Args:
+        manifest_file: Path to the manifest file.
+
+    Returns:
+        A selector object or a toolkit warning if loading fails or the selector is an alpha feature that is not enabled.
+    """
+    selector_dict = read_yaml_file(manifest_file, expected_output="dict")
+    try:
+        selector = SelectorAdapter.validate_python(selector_dict)
+    except ValidationError as e:
+        errors = humanize_validation_error(e)
+        return ResourceFormatWarning(manifest_file, tuple(errors), text="Invalid selector in metadata file, skipping.")
+    if not Flags.EXTEND_UPLOAD.is_enabled() and type(selector) in ALPHA_SELECTORS:
+        return MediumSeverityWarning(
+            f"Selector type '{type(selector).__name__}' in file '{manifest_file}' is in alpha. To enable it set the alpha flag 'extend-upload = true' in your CDF.toml file."
+        )
+    return selector
+
+
+__all__ = [
+    "AllChartsSelector",
+    "AssetCentricFileSelector",
+    "AssetCentricSelector",
+    "AssetSubtreeSelector",
+    "CanvasExternalIdSelector",
+    "CanvasSelector",
+    "ChartExternalIdSelector",
+    "ChartOwnerSelector",
+    "ChartSelector",
+    "DataPointsDataSetSelector",
+    "DataPointsFileSelector",
+    "DataPointsSelector",
+    "DataSelector",
+    "DataSetSelector",
+    "ExternalIdColumn",
+    "FileContentSelector",
+    "FileDataModelingTemplate",
+    "FileDataModelingTemplateSelector",
+    "FileIdentifierSelector",
+    "FileMetadataTemplate",
+    "FileMetadataTemplateSelector",
+    "InstanceColumn",
+    "InstanceFileSelector",
+    "InstanceSelector",
+    "InstanceSpaceSelector",
+    "InstanceViewSelector",
+    "InternalIdColumn",
+    "RawTableSelector",
+    "SelectedTable",
+    "SelectedView",
+    "Selector",
+    "SelectorAdapter",
+    "TimeSeriesColumn",
+    "load_selector",
+]

cognite_toolkit/_cdf_tk/storageio/selectors/_asset_centric.py

@@ -0,0 +1,60 @@
+from abc import ABC
+from pathlib import Path
+from typing import Any, Literal
+
+from cognite_toolkit._cdf_tk.utils.useful_types import AssetCentricKind
+
+from ._base import DataSelector
+
+
+class AssetCentricSelector(DataSelector, ABC):
+    kind: AssetCentricKind
+
+
+class DataSetSelector(AssetCentricSelector):
+    """Select data associated with a specific data set."""
+
+    type: Literal["dataSet"] = "dataSet"
+
+    data_set_external_id: str
+
+    @property
+    def group(self) -> str:
+        return f"DataSet_{self.data_set_external_id}"
+
+    def __str__(self) -> str:
+        return self.kind
+
+    def as_filter(self) -> dict[str, Any]:
+        return {"dataSetIds": [{"externalId": self.data_set_external_id}]}
+
+
+class AssetSubtreeSelector(AssetCentricSelector):
+    """Select data associated with an asset and its subtree."""
+
+    type: Literal["assetSubtree"] = "assetSubtree"
+    hierarchy: str
+
+    @property
+    def group(self) -> str:
+        return f"Hierarchy_{self.hierarchy}"
+
+    def __str__(self) -> str:
+        return self.kind
+
+    def as_filter(self) -> dict[str, Any]:
+        return {"assetSubtreeIds": [{"externalId": self.hierarchy}]}
+
+
+class AssetCentricFileSelector(AssetCentricSelector):
+    """Select data from a specific file."""
+
+    type: Literal["assetFile"] = "assetFile"
+    datafile: Path
+
+    @property
+    def group(self) -> str:
+        return "AssetCentric"
+
+    def __str__(self) -> str:
+        return f"file_{self.datafile.name}"