cognite-toolkit 0.5.97__py3-none-any.whl → 0.7.34__py3-none-any.whl

cognite_toolkit/_cdf_tk/storageio/_raw.py

@@ -0,0 +1,124 @@
+from collections.abc import Iterable, Sequence
+from itertools import chain
+from uuid import uuid4
+
+from cognite.client.data_classes import Row, RowWrite
+
+from cognite_toolkit._cdf_tk.cruds import RawDatabaseCRUD, RawTableCRUD
+from cognite_toolkit._cdf_tk.exceptions import ToolkitValueError
+from cognite_toolkit._cdf_tk.utils import sanitize_filename
+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 (
+    ConfigurableStorageIO,
+    Page,
+    StorageIOConfig,
+    TableUploadableStorageIO,
+    UploadItem,
+)
+from .selectors import RawTableSelector
+
+
+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: Row) -> str:
+        return str(item.key)
+
+    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: 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: 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: Sequence[Row], selector: RawTableSelector | None = None
+    ) -> list[dict[str, JsonVal]]:
+        return [row.as_write().dump() for row in data_chunk]
+
+    def json_to_resource(self, item_json: dict[str, JsonVal]) -> RowWrite:
+        return RowWrite._load(item_json)
+
+    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 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)
+
+    @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,133 @@
+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
+from ._three_d import ThreeDModelFilteredSelector, ThreeDModelIdSelector, ThreeDSelector
+
+Selector = Annotated[
+    RawTableSelector
+    | ThreeDModelIdSelector
+    | ThreeDModelFilteredSelector
+    | InstanceViewSelector
+    | InstanceFileSelector
+    | InstanceSpaceSelector
+    | AllChartsSelector
+    | ChartOwnerSelector
+    | AssetSubtreeSelector
+    | AssetCentricFileSelector
+    | DataSetSelector
+    | DataPointsDataSetSelector
+    | DataPointsFileSelector
+    | ChartExternalIdSelector
+    | CanvasExternalIdSelector
+    | FileMetadataTemplateSelector
+    | FileDataModelingTemplateSelector
+    | FileIdentifierSelector,
+    Field(discriminator="type"),
+]
+
+ALPHA_SELECTORS = {FileIdentifierSelector}
+INTERNAL = {ThreeDModelIdSelector, ThreeDModelFilteredSelector}
+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."
+        )
+    elif type(selector) in INTERNAL:
+        return MediumSeverityWarning(
+            f"Selector type '{type(selector).__name__}' in file '{manifest_file}' is for internal use only and cannot be used."
+        )
+    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",
+    "ThreeDModelIdSelector",
+    "ThreeDSelector",
+    "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}"

cognite_toolkit/_cdf_tk/storageio/selectors/_base.py

@@ -0,0 +1,80 @@
+from abc import ABC, abstractmethod
+from pathlib import Path
+
+from pydantic import BaseModel, ConfigDict
+from pydantic.alias_generators import to_camel
+
+from cognite_toolkit._cdf_tk.constants import DATA_MANIFEST_SUFFIX
+from cognite_toolkit._cdf_tk.utils.file import safe_write, sanitize_filename, yaml_safe_dump
+from cognite_toolkit._cdf_tk.utils.text import to_sentence_case
+from cognite_toolkit._cdf_tk.utils.useful_types import JsonVal
+
+
+class SelectorObject(BaseModel):
+    """This is used as base class for all selector objects including nested ones."""
+
+    model_config = ConfigDict(frozen=True, alias_generator=to_camel, populate_by_name=True)
+
+
+class DataSelector(SelectorObject, ABC):
+    """A selector gives instructions on what data to select from CDF.
+
+    For example, for instances it can be a view or container, while for assets it can be a data set or asset subtree.
+
+    Args:
+        type: The type of selector.
+        kind: The kind of data the selector is for (e.g., 'RawRows', 'Assets').
+    """
+
+    type: str
+    kind: str
+
+    def dump(self) -> dict[str, JsonVal]:
+        return self.model_dump(by_alias=True)
+
+    def dump_to_file(self, directory: Path) -> Path:
+        """Dumps the selector to a YAML file in the specified directory.
+
+        The filename is derived from the string representation of the selector.
+
+        Args:
+            directory: The directory where the YAML file will be saved.
+        """
+
+        filepath = directory / f"{sanitize_filename(str(self))}{DATA_MANIFEST_SUFFIX}"
+        filepath.parent.mkdir(parents=True, exist_ok=True)
+        safe_write(file=filepath, content=yaml_safe_dump(self.model_dump(mode="json", by_alias=True)), encoding="utf-8")
+        return filepath
+
+    @property
+    @abstractmethod
+    def group(self) -> str:
+        """A string representing the group of the selector, used for organizing files.
+
+        It is used when downloading to determine the subdirectory within the output directory.
+
+        For example, for raw table, the group would be the database name, while the selector itself
+        would be the table name.
+        """
+        raise NotImplementedError()
+
+    @property
+    def display_name(self) -> str:
+        """A human-readable name for the selector."""
+        return to_sentence_case(self.kind)
+
+    def __str__(self) -> str:
+        # We want to force subclasses to implement __str__
+        raise NotImplementedError()
+
+    def find_data_files(self, input_dir: Path, manifest_file: Path) -> list[Path]:
+        """Find data files in the specified input directory that match this selector.
+
+        Args:
+            input_dir: The directory to search for data files.
+            manifest_file: The manifest file that describes the data files.
+        Returns:
+            A list of Paths to the data files that match this selector.
+        """
+        data_file_prefix = manifest_file.name.removesuffix(DATA_MANIFEST_SUFFIX)
+        return [file for file in input_dir.glob(f"{data_file_prefix}*") if not file.name.endswith(DATA_MANIFEST_SUFFIX)]

cognite_toolkit/_cdf_tk/storageio/selectors/_canvas.py

@@ -0,0 +1,22 @@
+import hashlib
+from abc import ABC
+from typing import Literal
+
+from ._base import DataSelector
+
+
+class CanvasSelector(DataSelector, ABC):
+    kind: Literal["IndustrialCanvas"] = "IndustrialCanvas"
+
+
+class CanvasExternalIdSelector(CanvasSelector):
+    type: Literal["canvasExternalId"] = "canvasExternalId"
+    external_ids: tuple[str, ...]
+
+    @property
+    def group(self) -> str:
+        return "Canvas"
+
+    def __str__(self) -> str:
+        hash_ = hashlib.md5(",".join(sorted(self.external_ids)).encode()).hexdigest()[:8]
+        return f"canvas_count_{len(self.external_ids)}_hash_{hash_}"

cognite_toolkit/_cdf_tk/storageio/selectors/_charts.py

@@ -0,0 +1,45 @@
+import hashlib
+from abc import ABC
+from typing import Literal
+
+from ._base import DataSelector
+
+
+class ChartSelector(DataSelector, ABC):
+    kind: Literal["Charts"] = "Charts"
+
+
+class ChartOwnerSelector(ChartSelector):
+    type: Literal["chartOwner"] = "chartOwner"
+    owner_id: str
+
+    @property
+    def group(self) -> str:
+        return "Charts"
+
+    def __str__(self) -> str:
+        return self.owner_id
+
+
+class AllChartsSelector(ChartSelector):
+    type: Literal["allCharts"] = "allCharts"
+
+    @property
+    def group(self) -> str:
+        return "Charts"
+
+    def __str__(self) -> str:
+        return "all"
+
+
+class ChartExternalIdSelector(ChartSelector):
+    type: Literal["chartExternalId"] = "chartExternalId"
+    external_ids: tuple[str, ...]
+
+    @property
+    def group(self) -> str:
+        return "Charts"
+
+    def __str__(self) -> str:
+        hash_ = hashlib.md5(",".join(sorted(self.external_ids)).encode()).hexdigest()[:8]
+        return f"chart_count_{len(self.external_ids)}_hash_{hash_}"

cognite_toolkit/_cdf_tk/storageio/selectors/_datapoints.py

@@ -0,0 +1,89 @@
+from abc import ABC, abstractmethod
+from functools import cached_property
+from typing import Annotated, Any, ClassVar, Literal
+
+from cognite.client._proto.data_points_pb2 import (
+    InstanceId,
+)
+from pydantic import Field
+
+from ._base import DataSelector, SelectorObject
+
+
+class Column(SelectorObject, ABC):
+    column_type: str
+    column: str
+    dtype: Literal["numeric", "string"]
+
+    @abstractmethod
+    def as_wrapped_id(self) -> dict[str, Any]: ...
+
+
+class InstanceColumn(Column):
+    column_type: Literal["instance"] = "instance"
+    space: str
+    external_id: str
+
+    def as_wrapped_id(self) -> dict[str, Any]:
+        return {"instanceId": InstanceId(space=self.space, externalId=self.external_id)}
+
+
+class ExternalIdColumn(Column):
+    column_type: Literal["externalId"] = "externalId"
+    external_id: str
+
+    def as_wrapped_id(self) -> dict[str, Any]:
+        return {"externalId": self.external_id}
+
+
+class InternalIdColumn(Column):
+    column_type: Literal["internalId"] = "internalId"
+    internal_id: int
+
+    def as_wrapped_id(self) -> dict[str, Any]:
+        return {"id": self.internal_id}
+
+
+TimeSeriesColumn = Annotated[
+    InstanceColumn | ExternalIdColumn | InternalIdColumn,
+    Field(discriminator="column_type"),
+]
+
+
+class DataPointsSelector(DataSelector, ABC):
+    kind: Literal["Datapoints"] = "Datapoints"
+
+
+class DataPointsFileSelector(DataPointsSelector):
+    type: Literal["datapointsFile"] = "datapointsFile"
+
+    timestamp_column: str
+    columns: tuple[TimeSeriesColumn, ...]
+
+    @property
+    def group(self) -> str:
+        return "Datapoints"
+
+    def __str__(self) -> str:
+        return "datapoints_file"
+
+    @cached_property
+    def id_by_column(self) -> dict[str, Column]:
+        return {col.column: col for col in self.columns}
+
+
+class DataPointsDataSetSelector(DataPointsSelector):
+    required_columns: ClassVar[frozenset[str]] = frozenset({"externalId", "timestamp", "value"})
+    type: Literal["datapointsDataSet"] = "datapointsDataSet"
+
+    data_set_external_id: str
+    start: int | str | None = None
+    end: int | str | None = None
+    data_type: Literal["numeric", "string"] = "numeric"
+
+    @property
+    def group(self) -> str:
+        return f"DataSet_{self.data_set_external_id}"
+
+    def __str__(self) -> str:
+        return f"datapoints_dataset_{self.data_set_external_id}"