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

cognite_toolkit/_cdf_tk/commands/dump_data.py

@@ -1,490 +0,0 @@
-import time
-from abc import ABC, abstractmethod
-from collections.abc import Callable, Iterable, Iterator
-from functools import lru_cache
-from itertools import groupby
-from pathlib import Path
-from typing import Any, ClassVar, Generic, Literal
-
-from cognite.client.data_classes import (
-    Asset,
-    AssetFilter,
-    DataSetList,
-    Event,
-    EventFilter,
-    FileMetadata,
-    FileMetadataFilter,
-    LabelDefinitionList,
-    TimeSeries,
-    TimeSeriesFilter,
-)
-from cognite.client.data_classes._base import T_CogniteResource
-from rich.console import Console
-from rich.progress import track
-
-from cognite_toolkit._cdf_tk.client import ToolkitClient
-from cognite_toolkit._cdf_tk.commands._base import ToolkitCommand
-from cognite_toolkit._cdf_tk.exceptions import (
-    ToolkitFileExistsError,
-    ToolkitIsADirectoryError,
-    ToolkitValueError,
-)
-from cognite_toolkit._cdf_tk.loaders import (
-    AssetLoader,
-    DataSetsLoader,
-    EventLoader,
-    FileMetadataLoader,
-    LabelLoader,
-    ResourceLoader,
-    TimeSeriesLoader,
-)
-from cognite_toolkit._cdf_tk.utils import humanize_collection
-from cognite_toolkit._cdf_tk.utils.cdf import metadata_key_counts
-from cognite_toolkit._cdf_tk.utils.file import safe_rmtree
-from cognite_toolkit._cdf_tk.utils.producer_worker import ProducerWorkerExecutor
-from cognite_toolkit._cdf_tk.utils.table_writers import (
-    FileFormat,
-    Schema,
-    SchemaColumn,
-    SchemaColumnList,
-    TableFileWriter,
-)
-
-
-class DataFinder:
-    supported_formats: ClassVar[frozenset[FileFormat]] = frozenset()
-    # This is the standard maximum items that can be returns by most CDF endpoints.
-    chunk_size: ClassVar[int] = 1000
-
-    def validate_format(self, format_: str) -> Literal[FileFormat]:
-        if format_ in self.supported_formats:
-            return format_  # type: ignore[return-value]
-        raise ToolkitValueError(
-            f"Unsupported format {format_}. Supported formats are {humanize_collection(self.supported_formats)}."
-        )
-
-    @abstractmethod
-    def create_iterators(
-        self, format_: FileFormat, limit: int | None
-    ) -> Iterator[tuple[Schema, int, Iterable, Callable]]:
-        """Create an iterator for the specified format."""
-        raise NotImplementedError("This method should be implemented in subclasses.")
-
-
-class AssetCentricFinder(DataFinder, ABC, Generic[T_CogniteResource]):
-    def __init__(self, client: ToolkitClient, hierarchies: list[str], data_sets: list[str]):
-        self.client = client
-        self.hierarchies = hierarchies
-        self.data_sets = data_sets
-        self.loader = self._create_loader(client)
-        self._hierarchy_set = set(self.hierarchies)
-        self._data_set_set = set(self.data_sets)
-        self._used_labels: set[str] = set()
-        self._used_data_sets: set[str] = set()
-
-    @abstractmethod
-    def _create_loader(self, client: ToolkitClient) -> ResourceLoader:
-        """Create the appropriate loader for the finder."""
-        raise NotImplementedError()
-
-    @lru_cache
-    def aggregate_count(self, hierarchies: tuple[str, ...], data_sets: tuple[str, ...]) -> int:
-        return self._aggregate_count(list(hierarchies), list(data_sets))
-
-    @abstractmethod
-    def _aggregate_count(self, hierarchies: list[str], data_sets: list[str]) -> int:
-        raise NotImplementedError()
-
-    @abstractmethod
-    def _get_resource_columns(self) -> SchemaColumnList:
-        """Get the columns for the schema."""
-        raise NotImplementedError()
-
-    @abstractmethod
-    def create_resource_iterator(self, limit: int | None) -> Iterable:
-        raise NotImplementedError()
-
-    @abstractmethod
-    def _resource_processor(self, items: Iterable[T_CogniteResource]) -> list[tuple[str, list[dict[str, Any]]]]:
-        """Process the resources and return them in a format suitable for writing."""
-        raise NotImplementedError()
-
-    def _to_write(self, items: Iterable[T_CogniteResource]) -> list[dict[str, Any]]:
-        write_items: list[dict[str, Any]] = []
-        for item in items:
-            dumped = self.loader.dump_resource(item)
-            if "metadata" in dumped:
-                metadata = dumped.pop("metadata")
-                for key, value in metadata.items():
-                    dumped[f"metadata.{key}"] = value
-            if isinstance(dumped.get("labels"), list):
-                dumped["labels"] = [label["externalId"] for label in dumped["labels"]]
-                self._used_labels.update(dumped["labels"])
-            if "dataSetExternalId" in dumped:
-                self._used_data_sets.add(dumped["dataSetExternalId"])
-            write_items.append(dumped)
-        return write_items
-
-    def create_iterators(
-        self, format_: FileFormat, limit: int | None
-    ) -> Iterator[tuple[Schema, int, Iterable, Callable]]:
-        total = self.aggregate_count(tuple(self.hierarchies), tuple(self.data_sets))
-        columns = self._get_resource_columns()
-
-        iteration_count = total // self.chunk_size + (1 if total % self.chunk_size > 0 else 0)
-        if iteration_count == 0:
-            return
-
-        yield (
-            Schema(
-                display_name=self.loader.display_name,
-                format_=format_,
-                columns=columns,
-                folder_name=self.loader.folder_name,
-                kind=self.loader.kind,
-            ),
-            iteration_count,
-            self.create_resource_iterator(limit),
-            self._resource_processor,
-        )
-        if self._used_data_sets:
-            yield self._data_sets()
-        if self._used_labels:
-            yield self._labels()
-
-    def _data_sets(self) -> tuple[Schema, int, Iterable, Callable]:
-        data_sets = self.client.data_sets.retrieve_multiple(
-            external_ids=list(self._used_data_sets), ignore_unknown_ids=True
-        )
-        loader = DataSetsLoader.create_loader(self.client)
-
-        def process_data_sets(items: DataSetList) -> list[tuple[str, list[dict[str, Any]]]]:
-            # All data sets are written to a single group, thus the empty string as the group key.
-            # (Group keys are for example used in CSV files to create separate files for each
-            # data set an asset belongs to.)
-            return [("", [loader.dump_resource(item) for item in items])]
-
-        return (
-            # YAML format does not need columns.
-            Schema(
-                display_name=loader.display_name,
-                format_="yaml",
-                columns=SchemaColumnList(),
-                folder_name=loader.folder_name,
-                kind=loader.kind,
-            ),
-            1,
-            [data_sets],
-            process_data_sets,
-        )
-
-    def _labels(self) -> tuple[Schema, int, Iterable, Callable]:
-        labels = self.client.labels.retrieve(external_id=list(self._used_labels))
-        loader = LabelLoader.create_loader(self.client)
-
-        def process_labels(items: LabelDefinitionList) -> list[tuple[str, list[dict[str, Any]]]]:
-            # All labels are written to a single group, thus the empty string as the group key.
-            # (Group keys are for example used in CSV files to create separate files for each
-            # label an asset belongs to.)
-            return [("", [loader.dump_resource(item) for item in items])]
-
-        return (
-            # YAML format does not need columns.
-            Schema(
-                display_name=loader.display_name,
-                format_="yaml",
-                columns=SchemaColumnList(),
-                folder_name=loader.folder_name,
-                kind=loader.kind,
-            ),
-            1,
-            [labels],
-            process_labels,
-        )
-
-
-class AssetFinder(AssetCentricFinder[Asset]):
-    supported_formats = frozenset({"csv", "parquet", "yaml"})
-
-    def _create_loader(self, client: ToolkitClient) -> ResourceLoader:
-        return AssetLoader.create_loader(client)
-
-    def _aggregate_count(self, hierarchies: list[str], data_sets: list[str]) -> int:
-        return self.client.assets.aggregate_count(
-            filter=AssetFilter(
-                data_set_ids=[{"externalId": item} for item in data_sets] or None,
-                asset_subtree_ids=[{"externalId": item} for item in hierarchies] or None,
-            )
-        )
-
-    def create_resource_iterator(self, limit: int | None) -> Iterator:
-        return self.client.assets(
-            chunk_size=self.chunk_size,
-            asset_subtree_external_ids=self.hierarchies or None,
-            data_set_external_ids=self.data_sets or None,
-            limit=limit,
-        )
-
-    def _resource_processor(self, assets: Iterable[Asset]) -> list[tuple[str, list[dict[str, Any]]]]:
-        grouped_assets: list[tuple[str, list[dict[str, object]]]] = []
-        for group, asset_group in groupby(
-            sorted([(self._group(asset), asset) for asset in assets], key=lambda x: x[0]), key=lambda x: x[0]
-        ):
-            grouped_assets.append((group, self._to_write([asset for _, asset in asset_group])))
-        return grouped_assets
-
-    def _group(self, item: Asset) -> str:
-        if self.hierarchies and self.data_sets:
-            asset_external_id = self.client.lookup.assets.external_id(item.root_id or 0)
-            data_set_external_id = self.client.lookup.data_sets.external_id(item.data_set_id or 0)
-            if asset_external_id and data_set_external_id:
-                return f"{asset_external_id}.{data_set_external_id}"
-            elif asset_external_id:
-                return asset_external_id
-            elif data_set_external_id:
-                return data_set_external_id
-            return ""
-        elif self.hierarchies:
-            return self.client.lookup.assets.external_id(item.root_id or 0) or ""
-        elif self.data_sets:
-            return self.client.lookup.data_sets.external_id(item.data_set_id or 0) or ""
-        return ""
-
-    def _get_resource_columns(self) -> SchemaColumnList:
-        columns = SchemaColumnList(
-            [
-                SchemaColumn(name="externalId", type="string"),
-                SchemaColumn(name="name", type="string"),
-                SchemaColumn(name="parentExternalId", type="string"),
-                SchemaColumn(name="description", type="string"),
-                SchemaColumn(name="dataSetExternalId", type="string"),
-                SchemaColumn(name="source", type="string"),
-                SchemaColumn(name="labels", type="string", is_array=True),
-                SchemaColumn(name="geoLocation", type="json"),
-            ]
-        )
-        data_set_ids = self.client.lookup.data_sets.id(self.data_sets) if self.data_sets else []
-        root_ids = self.client.lookup.assets.id(self.hierarchies) if self.hierarchies else []
-        metadata_keys = metadata_key_counts(self.client, "assets", data_set_ids or None, root_ids or None)
-        sorted_keys = sorted([key for key, count in metadata_keys if count > 0])
-        columns.extend([SchemaColumn(name=f"metadata.{key}", type="string") for key in sorted_keys])
-        return columns
-
-
-class FileMetadataFinder(AssetCentricFinder[FileMetadata]):
-    supported_formats = frozenset({"csv", "parquet"})
-
-    def _create_loader(self, client: ToolkitClient) -> ResourceLoader:
-        return FileMetadataLoader.create_loader(client)
-
-    def _aggregate_count(self, hierarchies: list[str], data_sets: list[str]) -> int:
-        result = self.client.files.aggregate(
-            filter=FileMetadataFilter(
-                data_set_ids=[{"externalId": item} for item in data_sets] or None,
-                asset_subtree_ids=[{"externalId": item} for item in hierarchies] or None,
-            )
-        )
-        return result[0].count if result else 0
-
-    def _get_resource_columns(self) -> SchemaColumnList:
-        columns = SchemaColumnList(
-            [
-                SchemaColumn(name="externalId", type="string"),
-                SchemaColumn(name="name", type="string"),
-                SchemaColumn(name="directory", type="string"),
-                SchemaColumn(name="source", type="string"),
-                SchemaColumn(name="mimeType", type="string"),
-                SchemaColumn(name="assetExternalIds", type="string", is_array=True),
-                SchemaColumn(name="dataSetExternalId", type="string"),
-                SchemaColumn(name="sourceCreatedTime", type="integer"),
-                SchemaColumn(name="sourceModifiedTime", type="integer"),
-                SchemaColumn(name="securityCategories", type="string", is_array=True),
-                SchemaColumn(name="labels", type="string", is_array=True),
-                SchemaColumn(name="geoLocation", type="json"),
-            ]
-        )
-        data_set_ids = self.client.lookup.data_sets.id(self.data_sets) if self.data_sets else []
-        root_ids = self.client.lookup.assets.id(self.hierarchies) if self.hierarchies else []
-        metadata_keys = metadata_key_counts(self.client, "files", data_set_ids or None, root_ids or None)
-        sorted_keys = sorted([key for key, count in metadata_keys if count > 0])
-        columns.extend([SchemaColumn(name=f"metadata.{key}", type="string") for key in sorted_keys])
-        return columns
-
-    def create_resource_iterator(self, limit: int | None) -> Iterable:
-        return self.client.files(
-            chunk_size=self.chunk_size,
-            asset_subtree_external_ids=self.hierarchies or None,
-            data_set_external_ids=self.data_sets or None,
-            limit=limit,
-        )
-
-    def _resource_processor(self, items: Iterable[FileMetadata]) -> list[tuple[str, list[dict[str, Any]]]]:
-        return [("", self._to_write(items))]
-
-
-class TimeSeriesFinder(AssetCentricFinder[TimeSeries]):
-    supported_formats = frozenset({"csv", "parquet", "yaml"})
-
-    def _create_loader(self, client: ToolkitClient) -> TimeSeriesLoader:
-        return TimeSeriesLoader.create_loader(client)
-
-    def _aggregate_count(self, hierarchies: list[str], data_sets: list[str]) -> int:
-        return self.client.time_series.aggregate_count(
-            filter=TimeSeriesFilter(
-                data_set_ids=[{"externalId": item} for item in data_sets] or None,
-                asset_subtree_ids=[{"externalId": item} for item in hierarchies] or None,
-            )
-        )
-
-    def create_resource_iterator(self, limit: int | None) -> Iterator:
-        return self.client.time_series(
-            chunk_size=self.chunk_size,
-            asset_subtree_external_ids=self.hierarchies or None,
-            data_set_external_ids=self.data_sets or None,
-            limit=limit,
-        )
-
-    def _resource_processor(self, time_series: Iterable[TimeSeries]) -> list[tuple[str, list[dict[str, Any]]]]:
-        return [("", self._to_write(time_series))]
-
-    def _get_resource_columns(self) -> SchemaColumnList:
-        columns = SchemaColumnList(
-            [
-                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="dataSetExternalId", type="string"),
-                SchemaColumn(name="securityCategories", type="string", is_array=True),
-            ]
-        )
-        data_set_ids = self.client.lookup.data_sets.id(self.data_sets) if self.data_sets else []
-        root_ids = self.client.lookup.assets.id(self.hierarchies) if self.hierarchies else []
-        metadata_keys = metadata_key_counts(self.client, "timeseries", data_set_ids or None, root_ids or None)
-        sorted_keys = sorted([key for key, count in metadata_keys if count > 0])
-        columns.extend([SchemaColumn(name=f"metadata.{key}", type="string") for key in sorted_keys])
-        return columns
-
-
-class EventFinder(AssetCentricFinder[Event]):
-    supported_formats = frozenset({"csv", "parquet"})
-
-    def _create_loader(self, client: ToolkitClient) -> ResourceLoader:
-        return EventLoader.create_loader(client)
-
-    def _aggregate_count(self, hierarchies: list[str], data_sets: list[str]) -> int:
-        return self.client.events.aggregate_count(
-            filter=EventFilter(
-                data_set_ids=[{"externalId": item} for item in data_sets] or None,
-                asset_subtree_ids=[{"externalId": item} for item in hierarchies] or None,
-            )
-        )
-
-    def _get_resource_columns(self) -> SchemaColumnList:
-        columns = SchemaColumnList(
-            [
-                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"),
-            ]
-        )
-        data_set_ids = self.client.lookup.data_sets.id(self.data_sets) if self.data_sets else []
-        root_ids = self.client.lookup.assets.id(self.hierarchies) if self.hierarchies else []
-        metadata_keys = metadata_key_counts(self.client, "events", data_set_ids or None, root_ids or None)
-        sorted_keys = sorted([key for key, count in metadata_keys if count > 0])
-        columns.extend([SchemaColumn(name=f"metadata.{key}", type="string") for key in sorted_keys])
-        return columns
-
-    def create_resource_iterator(self, limit: int | None) -> Iterable:
-        return self.client.events(
-            chunk_size=self.chunk_size,
-            asset_subtree_external_ids=self.hierarchies or None,
-            data_set_external_ids=self.data_sets or None,
-            limit=limit,
-        )
-
-    def _resource_processor(self, items: Iterable[Event]) -> list[tuple[str, list[dict[str, Any]]]]:
-        return [("", self._to_write(items))]
-
-
-class DumpDataCommand(ToolkitCommand):
-    def dump_table(
-        self,
-        finder: DataFinder,
-        output_dir: Path,
-        clean: bool,
-        limit: int | None = None,
-        format_: str = "csv",
-        verbose: bool = False,
-        parallel_threshold: int = 10,
-        max_queue_size: int = 10,
-    ) -> None:
-        """Dumps data from CDF to a file
-
-        Args:
-            finder (DataFinder): The finder object to use for fetching data.
-            output_dir (Path): The directory to write the output files to.
-            clean (bool): Whether to clean the output directory before writing files.
-            limit (int | None, optional): The maximum number of rows to write. Defaults to None.
-            format_ (Literal["yaml", "csv", "parquet"], optional): The format of the output file. Defaults to "csv".
-            verbose (bool, optional): Whether to print detailed progress information. Defaults to False.
-            parallel_threshold (int, optional): The iteration threshold for parallel processing. Defaults to 10.
-            max_queue_size (int, optional): If using parallel processing, the maximum size of the queue. Defaults to 10.
-
-        """
-        valid_format = finder.validate_format(format_)
-        self.validate_directory(output_dir, clean)
-
-        console = Console()
-        # The ignore is used as MyPy does not understand that is_supported_format
-        # above guarantees that the format is valid.
-        for schema, iteration_count, resource_iterator, resource_processor in finder.create_iterators(
-            valid_format, limit
-        ):
-            writer_cls = TableFileWriter.get_write_cls(schema.format_)
-            row_counts = 0
-            t0 = time.perf_counter()
-            with writer_cls(schema, output_dir) as writer:
-                if iteration_count > parallel_threshold:
-                    executor = ProducerWorkerExecutor(
-                        download_iterable=resource_iterator,
-                        process=resource_processor,
-                        write=writer.write_rows,
-                        iteration_count=iteration_count,
-                        max_queue_size=max_queue_size,
-                        download_description=f"Downloading {schema.display_name}",
-                        process_description=f"Processing {schema.display_name}",
-                        write_description=f"Writing {schema.display_name} to file",
-                    )
-                    executor.run()
-                    if executor.error_occurred:
-                        raise ToolkitValueError(executor.error_message)
-                    row_counts = executor.total_items
-                else:
-                    for resources in track(
-                        resource_iterator, total=iteration_count, description=f"Dumping {schema.display_name}"
-                    ):
-                        row_counts += len(resources)
-                        processed = resource_processor(resources)
-                        writer.write_rows(processed)
-            elapsed = time.perf_counter() - t0
-            console.print(f"Dumped {row_counts:,} rows to {output_dir} in {elapsed:,.2f} seconds.")
-
-    @staticmethod
-    def validate_directory(output_dir: Path, clean: bool) -> None:
-        if output_dir.exists() and clean:
-            safe_rmtree(output_dir)
-        elif output_dir.exists():
-            raise ToolkitFileExistsError(f"Output directory {output_dir!s} already exists. Use --clean to remove it.")
-        elif output_dir.suffix:
-            raise ToolkitIsADirectoryError(f"Output directory {output_dir!s} is not a directory.")

cognite_toolkit/_cdf_tk/commands/featureflag.py

@@ -1,27 +0,0 @@
-from rich import print
-from rich.table import Table
-
-from cognite_toolkit._cdf_tk.cdf_toml import CDFToml
-from cognite_toolkit._cdf_tk.commands._base import ToolkitCommand
-from cognite_toolkit._cdf_tk.feature_flags import Flags
-
-
-class FeatureFlagCommand(ToolkitCommand):
-    @staticmethod
-    def list() -> None:
-        user_settings = CDFToml.load().alpha_flags
-        table = Table(title="feature flags")
-        table.add_column("Name", justify="left")
-        table.add_column("Description", justify="left")
-        table.add_column("Status", justify="left")
-
-        for flag in Flags:
-            is_enabled = user_settings.get(flag.name, False)
-            if is_enabled or flag.value.get("visible", False):
-                table.add_row(
-                    flag.name,
-                    str(flag.value.get("description", "")),
-                    "enabled" if is_enabled else "disabled",
-                    style="yellow" if is_enabled else "",
-                )
-        print(table)