cognite-toolkit 0.6.110__py3-none-any.whl → 0.6.112__py3-none-any.whl

cognite_toolkit/_cdf_tk/apps/_download_app.py

@@ -2,7 +2,9 @@ from enum import Enum
 from pathlib import Path
 from typing import Annotated, Any
 
+import questionary
 import typer
+from questionary import Choice
 from rich import print
 
 from cognite_toolkit._cdf_tk.client.data_classes.raw import RawTable
@@ -11,12 +13,14 @@ from cognite_toolkit._cdf_tk.constants import DATA_DEFAULT_DIR
 from cognite_toolkit._cdf_tk.storageio import (
     AssetIO,
     ChartIO,
+    EventIO,
+    FileMetadataIO,
     HierarchyIO,
     InstanceIO,
     RawIO,
+    TimeSeriesIO,
 )
 from cognite_toolkit._cdf_tk.storageio.selectors import (
-    AssetCentricSelector,
     AssetSubtreeSelector,
     ChartExternalIdSelector,
     ChartSelector,
@@ -28,11 +32,16 @@ from cognite_toolkit._cdf_tk.storageio.selectors import (
 )
 from cognite_toolkit._cdf_tk.utils.auth import EnvironmentVariables
 from cognite_toolkit._cdf_tk.utils.interactive_select import (
+    AssetCentricInteractiveSelect,
     AssetInteractiveSelect,
     DataModelingSelect,
+    EventInteractiveSelect,
+    FileMetadataInteractiveSelect,
     InteractiveChartSelect,
     RawTableInteractiveSelect,
+    TimeSeriesInteractiveSelect,
 )
+from cognite_toolkit._cdf_tk.utils.useful_types import AssetCentricKind
 
 
 class RawFormats(str, Enum):
@@ -77,6 +86,9 @@ class DownloadApp(typer.Typer):
         self.callback(invoke_without_command=True)(self.download_main)
         self.command("raw")(self.download_raw_cmd)
         self.command("assets")(self.download_assets_cmd)
+        self.command("timeseries")(self.download_timeseries_cmd)
+        self.command("events")(self.download_events_cmd)
+        self.command("files")(self.download_files_cmd)
         self.command("hierarchy")(self.download_hierarchy_cmd)
         self.command("instances")(self.download_instances_cmd)
         self.command("charts")(self.download_charts_cmd)
@@ -178,23 +190,15 @@ class DownloadApp(typer.Typer):
             )
         )
 
-    @staticmethod
     def download_assets_cmd(
+        self,
         ctx: typer.Context,
         data_sets: Annotated[
             list[str] | None,
             typer.Option(
                 "--data-set",
                 "-d",
-                help="List of data sets to download assets from. If this and hierarchy are not provided, an interactive selection will be made.",
-            ),
-        ] = None,
-        hierarchy: Annotated[
-            list[str] | None,
-            typer.Option(
-                "--hierarchy",
-                "-r",
-                help="List of asset hierarchies to download assets from. If this and data sets are not provided, an interactive selection will be made.",
+                help="List of data sets to download assets from. If this is not provided, an interactive selection will be made.",
             ),
         ] = None,
         file_format: Annotated[
@@ -241,20 +245,17 @@ class DownloadApp(typer.Typer):
     ) -> None:
         """This command will download assets from CDF into a temporary directory."""
         client = EnvironmentVariables.create_from_environment().get_client()
-        is_interactive = not data_sets and not hierarchy
-        if is_interactive:
-            interactive = AssetInteractiveSelect(client, "download assets")
-            selector_type = interactive.select_hierarchies_or_data_sets()
-            if selector_type == "Data Set":
-                data_sets = interactive.select_data_sets()
-            else:
-                hierarchy = interactive.select_hierarchies()
-
-        selectors: list[AssetCentricSelector] = []
-        if data_sets:
-            selectors.extend([DataSetSelector(data_set_external_id=ds, kind="Assets") for ds in data_sets])
-        if hierarchy:
-            selectors.extend([AssetSubtreeSelector(hierarchy=h, kind="Assets") for h in hierarchy])
+        if data_sets is None:
+            data_sets, file_format, compression, output_dir, limit = self._asset_centric_interactive(
+                AssetInteractiveSelect(client, "download"),
+                file_format,
+                compression,
+                output_dir,
+                limit,
+                "Assets",
+            )
+
+        selectors = [DataSetSelector(kind="Assets", data_set_external_id=data_set) for data_set in data_sets]
         cmd = DownloadCommand()
         cmd.run(
             lambda: cmd.download(
@@ -268,6 +269,287 @@ class DownloadApp(typer.Typer):
             )
         )
 
+    @classmethod
+    def _asset_centric_interactive(
+        cls,
+        selector: AssetCentricInteractiveSelect,
+        file_format: AssetCentricFormats,
+        compression: CompressionFormat,
+        output_dir: Path,
+        limit: int,
+        kind: AssetCentricKind,
+    ) -> tuple[list[str], AssetCentricFormats, CompressionFormat, Path, int]:
+        data_sets = selector.select_data_sets()
+        display_name = kind.casefold() + "s"
+        file_format = questionary.select(
+            f"Select format to download the {display_name} in:",
+            choices=[Choice(title=format_.value, value=format_) for format_ in AssetCentricFormats],
+            default=file_format,
+        ).ask()
+        compression = questionary.select(
+            f"Select compression format to use when downloading the {display_name}:",
+            choices=[Choice(title=comp.value, value=comp) for comp in CompressionFormat],
+            default=compression,
+        ).ask()
+        output_dir = Path(
+            questionary.path(
+                "Where to download the assets:",
+                default=str(output_dir),
+                only_directories=True,
+            ).ask()
+        )
+        while True:
+            limit_str = questionary.text(
+                f"The maximum number of {display_name} to download from each dataset. Use -1 to download all {display_name}.",
+                default=str(limit),
+            ).ask()
+            if limit_str is None:
+                raise typer.Abort()
+            try:
+                limit = int(limit_str)
+                break
+            except ValueError:
+                print("[red]Please enter a valid integer for the limit.[/]")
+        return data_sets, file_format, compression, output_dir, limit
+
+    def download_timeseries_cmd(
+        self,
+        ctx: typer.Context,
+        data_sets: Annotated[
+            list[str] | None,
+            typer.Option(
+                "--data-set",
+                "-d",
+                help="List of data sets to download time series from. If this is not provided, an interactive selection will be made.",
+            ),
+        ] = None,
+        file_format: Annotated[
+            AssetCentricFormats,
+            typer.Option(
+                "--format",
+                "-f",
+                help="Format to download the time series in.",
+            ),
+        ] = AssetCentricFormats.csv,
+        compression: Annotated[
+            CompressionFormat,
+            typer.Option(
+                "--compression",
+                "-z",
+                help="Compression format to use when downloading the time series.",
+            ),
+        ] = CompressionFormat.none,
+        output_dir: Annotated[
+            Path,
+            typer.Option(
+                "--output-dir",
+                "-o",
+                help="Where to download the time series.",
+                allow_dash=True,
+            ),
+        ] = DEFAULT_DOWNLOAD_DIR,
+        limit: Annotated[
+            int,
+            typer.Option(
+                "--limit",
+                "-l",
+                help="The maximum number of time series to download from each dataset. Use -1 to download all time series.",
+            ),
+        ] = 100_000,
+        verbose: Annotated[
+            bool,
+            typer.Option(
+                "--verbose",
+                "-v",
+                help="Turn on to get more verbose output when running the command",
+            ),
+        ] = False,
+    ) -> None:
+        """This command will download time series from CDF into a temporary directory."""
+        client = EnvironmentVariables.create_from_environment().get_client()
+        if data_sets is None:
+            data_sets, file_format, compression, output_dir, limit = self._asset_centric_interactive(
+                TimeSeriesInteractiveSelect(client, "download"),
+                file_format,
+                compression,
+                output_dir,
+                limit,
+                "TimeSeries",
+            )
+
+        selectors = [DataSetSelector(kind="TimeSeries", data_set_external_id=data_set) for data_set in data_sets]
+        cmd = DownloadCommand()
+        cmd.run(
+            lambda: cmd.download(
+                selectors=selectors,
+                io=TimeSeriesIO(client),
+                output_dir=output_dir,
+                file_format=f".{file_format.value}",
+                compression=compression.value,
+                limit=limit if limit != -1 else None,
+                verbose=verbose,
+            )
+        )
+
+    def download_events_cmd(
+        self,
+        ctx: typer.Context,
+        data_sets: Annotated[
+            list[str] | None,
+            typer.Option(
+                "--data-set",
+                "-d",
+                help="List of data sets to download events from. If this is not provided, an interactive selection will be made.",
+            ),
+        ] = None,
+        file_format: Annotated[
+            AssetCentricFormats,
+            typer.Option(
+                "--format",
+                "-f",
+                help="Format to download the events in.",
+            ),
+        ] = AssetCentricFormats.csv,
+        compression: Annotated[
+            CompressionFormat,
+            typer.Option(
+                "--compression",
+                "-z",
+                help="Compression format to use when downloading the events.",
+            ),
+        ] = CompressionFormat.none,
+        output_dir: Annotated[
+            Path,
+            typer.Option(
+                "--output-dir",
+                "-o",
+                help="Where to download the events.",
+                allow_dash=True,
+            ),
+        ] = DEFAULT_DOWNLOAD_DIR,
+        limit: Annotated[
+            int,
+            typer.Option(
+                "--limit",
+                "-l",
+                help="The maximum number of events to download from each dataset. Use -1 to download all events.",
+            ),
+        ] = 100_000,
+        verbose: Annotated[
+            bool,
+            typer.Option(
+                "--verbose",
+                "-v",
+                help="Turn on to get more verbose output when running the command",
+            ),
+        ] = False,
+    ) -> None:
+        """This command will download events from CDF into a temporary directory."""
+        client = EnvironmentVariables.create_from_environment().get_client()
+        if data_sets is None:
+            data_sets, file_format, compression, output_dir, limit = self._asset_centric_interactive(
+                EventInteractiveSelect(client, "download"),
+                file_format,
+                compression,
+                output_dir,
+                limit,
+                "Events",
+            )
+
+        selectors = [DataSetSelector(kind="Events", data_set_external_id=data_set) for data_set in data_sets]
+        cmd = DownloadCommand()
+
+        cmd.run(
+            lambda: cmd.download(
+                selectors=selectors,
+                io=EventIO(client),
+                output_dir=output_dir,
+                file_format=f".{file_format.value}",
+                compression=compression.value,
+                limit=limit if limit != -1 else None,
+                verbose=verbose,
+            )
+        )
+
+    def download_files_cmd(
+        self,
+        ctx: typer.Context,
+        data_sets: Annotated[
+            list[str] | None,
+            typer.Option(
+                "--data-set",
+                "-d",
+                help="List of data sets to download file metadata from. If this is not provided, an interactive selection will be made.",
+            ),
+        ] = None,
+        file_format: Annotated[
+            AssetCentricFormats,
+            typer.Option(
+                "--format",
+                "-f",
+                help="Format to download the file metadata in.",
+            ),
+        ] = AssetCentricFormats.csv,
+        compression: Annotated[
+            CompressionFormat,
+            typer.Option(
+                "--compression",
+                "-z",
+                help="Compression format to use when downloading the file metadata.",
+            ),
+        ] = CompressionFormat.none,
+        output_dir: Annotated[
+            Path,
+            typer.Option(
+                "--output-dir",
+                "-o",
+                help="Where to download the file metadata.",
+                allow_dash=True,
+            ),
+        ] = DEFAULT_DOWNLOAD_DIR,
+        limit: Annotated[
+            int,
+            typer.Option(
+                "--limit",
+                "-l",
+                help="The maximum number of file metadata to download from each dataset. Use -1 to download all file metadata.",
+            ),
+        ] = 100_000,
+        verbose: Annotated[
+            bool,
+            typer.Option(
+                "--verbose",
+                "-v",
+                help="Turn on to get more verbose output when running the command",
+            ),
+        ] = False,
+    ) -> None:
+        """This command will download file metadata from CDF into a temporary directory."""
+        client = EnvironmentVariables.create_from_environment().get_client()
+        if data_sets is None:
+            data_sets, file_format, compression, output_dir, limit = self._asset_centric_interactive(
+                FileMetadataInteractiveSelect(client, "download"),
+                file_format,
+                compression,
+                output_dir,
+                limit,
+                "FileMetadata",
+            )
+
+        selectors = [DataSetSelector(kind="FileMetadata", data_set_external_id=data_set) for data_set in data_sets]
+        cmd = DownloadCommand()
+        cmd.run(
+            lambda: cmd.download(
+                selectors=selectors,
+                io=FileMetadataIO(client),
+                output_dir=output_dir,
+                file_format=f".{file_format.value}",
+                compression=compression.value,
+                limit=limit if limit != -1 else None,
+                verbose=verbose,
+            )
+        )
+
     @staticmethod
     def download_hierarchy_cmd(
         ctx: typer.Context,

cognite_toolkit/_cdf_tk/apps/_migrate_app.py

@@ -18,7 +18,7 @@ from cognite_toolkit._cdf_tk.commands._migrate.creators import (
     InstanceSpaceCreator,
     SourceSystemCreator,
 )
-from cognite_toolkit._cdf_tk.commands._migrate.data_mapper import AssetCentricMapper
+from cognite_toolkit._cdf_tk.commands._migrate.data_mapper import AssetCentricMapper, ChartMapper
 from cognite_toolkit._cdf_tk.commands._migrate.migration_io import (
     AnnotationMigrationIO,
     AssetCentricMigrationIO,
@@ -28,12 +28,15 @@ from cognite_toolkit._cdf_tk.commands._migrate.selectors import (
     MigrateDataSetSelector,
     MigrationCSVFileSelector,
 )
+from cognite_toolkit._cdf_tk.storageio import ChartIO
+from cognite_toolkit._cdf_tk.storageio.selectors import ChartExternalIdSelector
 from cognite_toolkit._cdf_tk.utils.auth import EnvironmentVariables
 from cognite_toolkit._cdf_tk.utils.cli_args import parse_view_str
 from cognite_toolkit._cdf_tk.utils.interactive_select import (
     AssetInteractiveSelect,
     DataModelingSelect,
     FileMetadataInteractiveSelect,
+    InteractiveChartSelect,
     ResourceViewMappingInteractiveSelect,
 )
 from cognite_toolkit._cdf_tk.utils.useful_types import AssetCentricKind
@@ -54,6 +57,7 @@ class MigrateApp(typer.Typer):
         self.command("files")(self.files)
         self.command("annotations")(self.annotations)
         self.command("canvas")(self.canvas)
+        self.command("charts")(self.charts)
         # Uncomment when infield v2 config migration is ready
         # self.command("infield-configs")(self.infield_configs)
 
@@ -887,6 +891,68 @@ class MigrateApp(typer.Typer):
             )
         )
 
+    @staticmethod
+    def charts(
+        ctx: typer.Context,
+        external_id: Annotated[
+            list[str] | None,
+            typer.Argument(
+                help="The external ID of the Chart to migrate. If not provided, an interactive selection will be "
+                "performed to select the Charts to migrate."
+            ),
+        ] = None,
+        log_dir: Annotated[
+            Path,
+            typer.Option(
+                "--log-dir",
+                "-l",
+                help="Path to the directory where migration logs will be stored.",
+            ),
+        ] = Path(f"migration_logs_{TODAY}"),
+        dry_run: Annotated[
+            bool,
+            typer.Option(
+                "--dry-run",
+                "-d",
+                help="If set, the migration will not be executed, but only a report of "
+                "what would be done is printed. This is useful for checking that all time series referenced by the Charts "
+                "have been migrated to the new data modeling resources in CDF.",
+            ),
+        ] = False,
+        verbose: Annotated[
+            bool,
+            typer.Option(
+                "--verbose",
+                "-v",
+                help="Turn on to get more verbose output when running the command",
+            ),
+        ] = False,
+    ) -> None:
+        """Migrate Charts from time series references to data modeling in CDF.
+
+        This command expects that the CogniteMigration data model is already deployed, and that the Mapping view
+        is populated with the mapping from time series to the new data modeling resources.
+        """
+        client = EnvironmentVariables.create_from_environment().get_client()
+
+        selected_external_ids: list[str]
+        if external_id:
+            selected_external_ids = external_id
+        else:
+            selected_external_ids = InteractiveChartSelect(client).select_external_ids()
+
+        cmd = MigrationCommand()
+        cmd.run(
+            lambda: cmd.migrate(
+                selected=ChartExternalIdSelector(external_ids=tuple(selected_external_ids)),
+                data=ChartIO(client),
+                mapper=ChartMapper(client),
+                log_dir=log_dir,
+                dry_run=dry_run,
+                verbose=verbose,
+            )
+        )
+
     @staticmethod
     def infield_configs(
         ctx: typer.Context,

cognite_toolkit/_cdf_tk/client/data_classes/base.py

@@ -1,10 +1,14 @@
 import sys
 from abc import ABC, abstractmethod
-from typing import Any, Generic, TypeVar
+from collections import UserList
+from typing import TYPE_CHECKING, Any, Generic, TypeVar
 
 from pydantic import BaseModel, ConfigDict
 from pydantic.alias_generators import to_camel
 
+if TYPE_CHECKING:
+    from cognite.client import CogniteClient
+
 if sys.version_info >= (3, 11):
     from typing import Self
 else:
@@ -61,3 +65,23 @@ class Identifier(BaseModel):
 
     def as_id(self) -> Self:
         return self
+
+
+T_Resource = TypeVar("T_Resource", bound=RequestResource | ResponseResource)
+
+
+class BaseResourceList(UserList[T_Resource]):
+    """Base class for resource lists."""
+
+    _RESOURCE: type[T_Resource]
+
+    def __init__(self, initlist: list[T_Resource] | None = None, **_: Any) -> None:
+        super().__init__(initlist or [])
+
+    def dump(self, camel_case: bool = True) -> list[dict[str, Any]]:
+        return [item.dump(camel_case) for item in self.data]
+
+    @classmethod
+    def load(cls, data: list[dict[str, Any]], cognite_client: "CogniteClient | None" = None) -> Self:
+        items = [cls._RESOURCE.model_validate(item) for item in data]
+        return cls(items)  # type: ignore[arg-type]

cognite_toolkit/_cdf_tk/client/data_classes/infield.py

@@ -1,15 +1,16 @@
 import sys
-from collections import UserList
 from typing import Any, ClassVar, Literal
 
-from cognite.client import CogniteClient
 from pydantic import JsonValue, field_validator
 from pydantic_core.core_schema import ValidationInfo
 
-from cognite_toolkit._cdf_tk.protocols import ResourceRequestListProtocol, ResourceResponseListProtocol
+from cognite_toolkit._cdf_tk.protocols import (
+    ResourceRequestListProtocol,
+    ResourceResponseListProtocol,
+)
 from cognite_toolkit._cdf_tk.utils.text import sanitize_instance_external_id
 
-from .base import ResponseResource
+from .base import BaseResourceList, ResponseResource
 from .instance_api import InstanceRequestResource, ViewReference
 
 if sys.version_info >= (3, 11):
@@ -74,29 +75,13 @@ class InfieldLocationConfig(
 
 
 class InfieldLocationConfigList(
-    UserList[InfieldLocationConfig],
+    BaseResourceList[InfieldLocationConfig],
     ResourceResponseListProtocol,
     ResourceRequestListProtocol,
 ):
     """A list of InfieldLocationConfig objects."""
 
     _RESOURCE = InfieldLocationConfig
-    data: list[InfieldLocationConfig]
-
-    def __init__(self, initlist: list[InfieldLocationConfig] | None = None, **_: Any) -> None:
-        super().__init__(initlist or [])
-
-    def dump(self, camel_case: bool = True) -> list[dict[str, Any]]:
-        """Serialize the list of InfieldLocationConfig objects to a list of dictionaries."""
-        return [item.dump(camel_case) for item in self.data]
-
-    @classmethod
-    def load(
-        cls, data: list[dict[str, Any]], cognite_client: CogniteClient | None = None
-    ) -> "InfieldLocationConfigList":
-        """Deserialize a list of dictionaries to an InfieldLocationConfigList."""
-        items = [InfieldLocationConfig.model_validate(item) for item in data]
-        return cls(items)
 
     def as_write(self) -> Self:
         return self

cognite_toolkit/_cdf_tk/client/data_classes/streams.py

@@ -1,18 +1,12 @@
-import sys
-from collections import UserList
-from typing import Any, Literal
-
-from cognite.client import CogniteClient
+from typing import Literal
 
 from cognite_toolkit._cdf_tk.constants import StreamTemplateName
-from cognite_toolkit._cdf_tk.protocols import ResourceRequestListProtocol, ResourceResponseListProtocol
-
-from .base import BaseModelObject, RequestResource, ResponseResource
+from cognite_toolkit._cdf_tk.protocols import (
+    ResourceRequestListProtocol,
+    ResourceResponseListProtocol,
+)
 
-if sys.version_info >= (3, 11):
-    from typing import Self
-else:
-    from typing_extensions import Self
+from .base import BaseModelObject, BaseResourceList, RequestResource, ResponseResource
 
 
 class StreamRequest(RequestResource):
@@ -25,22 +19,11 @@ class StreamRequest(RequestResource):
         return self.external_id
 
 
-class StreamRequestList(UserList[StreamRequest], ResourceRequestListProtocol):
+class StreamRequestList(BaseResourceList[StreamRequest], ResourceRequestListProtocol):
     """List of Stream request resources."""
 
     _RESOURCE = StreamRequest
 
-    def __init__(self, initlist: list[StreamRequest] | None = None, **_: Any) -> None:
-        super().__init__(initlist or [])
-
-    def dump(self, camel_case: bool = True) -> list[dict[str, Any]]:
-        return [item.dump(camel_case) for item in self.data]
-
-    @classmethod
-    def load(cls, data: list[dict[str, Any]], cognite_client: CogniteClient | None = None) -> "StreamRequestList":
-        items = [StreamRequest.model_validate(item) for item in data]
-        return cls(items)
-
 
 class LifecycleObject(BaseModelObject):
     """Lifecycle object."""
@@ -98,22 +81,10 @@ class StreamResponse(ResponseResource["StreamRequest"]):
         )
 
 
-class StreamResponseList(UserList[StreamResponse], ResourceResponseListProtocol):
+class StreamResponseList(BaseResourceList[StreamResponse], ResourceResponseListProtocol):
     """List of Stream response resources."""
 
     _RESOURCE = StreamResponse
-    data: list[StreamResponse]
-
-    def __init__(self, initlist: list[StreamResponse] | None = None, **_: Any) -> None:
-        super().__init__(initlist or [])
-
-    def dump(self, camel_case: bool = True) -> list[dict[str, Any]]:
-        return [item.dump(camel_case) for item in self.data]
-
-    @classmethod
-    def load(cls, data: list[dict[str, Any]], cognite_client: CogniteClient | None = None) -> "StreamResponseList":
-        items = [StreamResponse.model_validate(item) for item in data]
-        return cls(items)
 
-    def as_write(self) -> Self:
-        return self
+    def as_write(self) -> StreamRequestList:
+        return StreamRequestList([item.as_write() for item in self.data])

cognite_toolkit/_cdf_tk/commands/_migrate/data_mapper.py

@@ -17,7 +17,11 @@ from cognite.client.data_classes.data_modeling import (
 
 from cognite_toolkit._cdf_tk.client import ToolkitClient
 from cognite_toolkit._cdf_tk.client.data_classes.charts import Chart, ChartWrite
-from cognite_toolkit._cdf_tk.client.data_classes.charts_data import ChartCoreTimeseries, ChartSource, ChartTimeseries
+from cognite_toolkit._cdf_tk.client.data_classes.charts_data import (
+    ChartCoreTimeseries,
+    ChartSource,
+    ChartTimeseries,
+)
 from cognite_toolkit._cdf_tk.client.data_classes.migration import ResourceViewMappingApply
 from cognite_toolkit._cdf_tk.commands._migrate.conversion import DirectRelationCache, asset_centric_to_dm
 from cognite_toolkit._cdf_tk.commands._migrate.data_classes import AssetCentricMapping