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

cognite_toolkit/_cdf_tk/client/testing.py

@@ -9,6 +9,7 @@ from cognite.client._api.functions import FunctionCallsAPI, FunctionSchedulesAPI
 from cognite.client._api.raw import RawDatabasesAPI, RawRowsAPI, RawTablesAPI
 from cognite.client._api.synthetic_time_series import SyntheticDatapointsAPI
 from cognite.client.testing import CogniteClientMock
+from rich.console import Console
 
 from cognite_toolkit._cdf_tk.client._toolkit_client import ToolkitClient
 
@@ -20,18 +21,29 @@ from .api.extended_files import ExtendedFileMetadataAPI
 from .api.extended_functions import ExtendedFunctionsAPI
 from .api.extended_raw import ExtendedRawAPI
 from .api.extended_timeseries import ExtendedTimeSeriesAPI
+from .api.infield import InfieldAPI, InFieldCDMConfigAPI, InfieldConfigAPI
 from .api.location_filters import LocationFiltersAPI
 from .api.lookup import (
     AssetLookUpAPI,
     DataSetLookUpAPI,
+    EventLookUpAPI,
     ExtractionPipelineLookUpAPI,
+    FileMetadataLookUpAPI,
     FunctionLookUpAPI,
     LocationFiltersLookUpAPI,
     LookUpGroup,
     SecurityCategoriesLookUpAPI,
     TimeSeriesLookUpAPI,
 )
-from .api.migration import InstanceSourceAPI, MigrationAPI, ResourceViewMappingAPI
+from .api.migration import (
+    CreatedSourceSystemAPI,
+    InstanceSourceAPI,
+    LookupAPI,
+    MigrationAPI,
+    MigrationLookupAPI,
+    ResourceViewMappingAPI,
+)
+from .api.project import ProjectAPI
 from .api.robotics import RoboticsAPI
 from .api.robotics.capabilities import CapabilitiesAPI
 from .api.robotics.data_postprocessing import DataPostProcessingAPI
@@ -40,6 +52,7 @@ from .api.robotics.locations import LocationsAPI as RoboticsLocationsAPI
 from .api.robotics.maps import MapsAPI
 from .api.search import SearchAPI
 from .api.search_config import SearchConfigurationsAPI
+from .api.streams import StreamsAPI
 from .api.token import TokenAPI
 from .api.verify import VerifyAPI
 
@@ -55,6 +68,7 @@ class ToolkitClientMock(CogniteClientMock):
             super().__init__(*args, **kwargs)
             return None
         super().__init__(*args, **kwargs)
+        self.console = Console()
         # Developer note:
         # - Please add your mocked APIs in chronological order
         # - For nested APIs:
@@ -68,6 +82,11 @@ class ToolkitClientMock(CogniteClientMock):
         self.functions = MagicMock(spec=ExtendedFunctionsAPI)
         self.functions.calls = MagicMock(spec_set=FunctionCallsAPI)
         self.functions.schedules = MagicMock(spec_set=FunctionSchedulesAPI)
+        self.infield = MagicMock(spec=InfieldAPI)
+        self.infield.config = MagicMock(spec_set=InfieldConfigAPI)
+        self.infield.cdm_config = MagicMock(spec_set=InFieldCDMConfigAPI)
+
+        self.project = MagicMock(spec_set=ProjectAPI)
 
         self.search = MagicMock(spec=SearchAPI)
         self.search.locations = MagicMock(spec_set=LocationFiltersAPI)
@@ -77,13 +96,21 @@ class ToolkitClientMock(CogniteClientMock):
         self.lookup.data_sets = MagicMock(spec_set=DataSetLookUpAPI)
         self.lookup.assets = MagicMock(spec_set=AssetLookUpAPI)
         self.lookup.time_series = MagicMock(spec_set=TimeSeriesLookUpAPI)
+        self.lookup.files = MagicMock(spec_set=FileMetadataLookUpAPI)
+        self.lookup.events = MagicMock(spec_set=EventLookUpAPI)
         self.lookup.security_categories = MagicMock(spec_set=SecurityCategoriesLookUpAPI)
         self.lookup.location_filters = MagicMock(spec_set=LocationFiltersLookUpAPI)
         self.lookup.extraction_pipelines = MagicMock(spec_set=ExtractionPipelineLookUpAPI)
         self.lookup.functions = MagicMock(spec_set=FunctionLookUpAPI)
         self.migration = MagicMock(spec=MigrationAPI)
         self.migration.instance_source = MagicMock(spec_set=InstanceSourceAPI)
+        self.migration.lookup = MagicMock(spec=MigrationLookupAPI)
+        self.migration.lookup.assets = MagicMock(spec_set=LookupAPI)
+        self.migration.lookup.events = MagicMock(spec_set=LookupAPI)
+        self.migration.lookup.files = MagicMock(spec_set=LookupAPI)
+        self.migration.lookup.time_series = MagicMock(spec_set=LookupAPI)
         self.migration.resource_view_mapping = MagicMock(spec_set=ResourceViewMappingAPI)
+        self.migration.created_source_system = MagicMock(spec_set=CreatedSourceSystemAPI)
         self.raw = MagicMock(spec=ExtendedRawAPI)
         self.raw.databases = MagicMock(spec_set=RawDatabasesAPI)
         self.raw.rows = MagicMock(spec_set=RawRowsAPI)
@@ -104,6 +131,8 @@ class ToolkitClientMock(CogniteClientMock):
         self.time_series.data.synthetic = MagicMock(spec_set=SyntheticDatapointsAPI)
         self.time_series.subscriptions = MagicMock(spec_set=DatapointsSubscriptionAPI)
 
+        self.streams = MagicMock(spec=StreamsAPI)
+
         # This is a helper API, not a real API.
         self.token = TokenAPI(self)
         self.verify = MagicMock(spec_set=VerifyAPI)
@@ -112,6 +141,8 @@ class ToolkitClientMock(CogniteClientMock):
 @contextmanager
 def monkeypatch_toolkit_client() -> Iterator[ToolkitClientMock]:
     toolkit_client_mock = ToolkitClientMock()
-    ToolkitClient.__new__ = lambda *args, **kwargs: toolkit_client_mock  # type: ignore[method-assign]
-    yield toolkit_client_mock
-    ToolkitClient.__new__ = lambda cls, *args, **kwargs: object.__new__(cls)  # type: ignore[method-assign]
+    try:
+        ToolkitClient.__new__ = lambda *args, **kwargs: toolkit_client_mock  # type: ignore[method-assign]
+        yield toolkit_client_mock
+    finally:
+        ToolkitClient.__new__ = lambda cls, *args, **kwargs: object.__new__(cls)  # type: ignore[method-assign]

cognite_toolkit/_cdf_tk/commands/__init__.py

@@ -1,45 +1,38 @@
 from ._download import DownloadCommand
 from ._migrate import (
-    MigrateFilesCommand,
-    MigrateTimeseriesCommand,
-    MigrationCanvasCommand,
+    MigrationCommand,
     MigrationPrepareCommand,
 )
-from ._populate import PopulateCommand
 from ._profile import ProfileAssetCentricCommand, ProfileAssetCommand, ProfileRawCommand, ProfileTransformationCommand
 from ._purge import PurgeCommand
 from ._upload import UploadCommand
+from .about import AboutCommand
 from .auth import AuthCommand
 from .build_cmd import BuildCommand
 from .clean import CleanCommand
 from .collect import CollectCommand
 from .deploy import DeployCommand
-from .dump_data import DumpDataCommand
 from .dump_resource import DumpResourceCommand
-from .featureflag import FeatureFlagCommand
 from .init import InitCommand
 from .modules import ModulesCommand
 from .pull import PullCommand
 from .repo import RepoCommand
+from .resources import ResourcesCommand
 from .run import RunFunctionCommand, RunTransformationCommand, RunWorkflowCommand
 
 __all__ = [
+    "AboutCommand",
     "AuthCommand",
     "BuildCommand",
     "CleanCommand",
     "CollectCommand",
     "DeployCommand",
     "DownloadCommand",
-    "DumpDataCommand",
     "DumpResourceCommand",
-    "FeatureFlagCommand",
     "InitCommand",
-    "MigrateFilesCommand",
-    "MigrateTimeseriesCommand",
-    "MigrationCanvasCommand",
+    "MigrationCommand",
     "MigrationPrepareCommand",
     "ModulesCommand",
-    "PopulateCommand",
     "ProfileAssetCentricCommand",
     "ProfileAssetCommand",
     "ProfileRawCommand",
@@ -47,6 +40,7 @@ __all__ = [
     "PullCommand",
     "PurgeCommand",
     "RepoCommand",
+    "ResourcesCommand",
     "RunFunctionCommand",
     "RunTransformationCommand",
     "RunWorkflowCommand",

cognite_toolkit/_cdf_tk/commands/_base.py

@@ -5,6 +5,7 @@ from typing import Any
 from rich import print
 from rich.console import Console
 
+from cognite_toolkit._cdf_tk.data_classes import CommandTrackingInfo
 from cognite_toolkit._cdf_tk.tk_warnings import (
     ToolkitWarning,
     WarningList,
@@ -20,13 +21,19 @@ class ToolkitCommand:
         self.silent = silent
         self.warning_list = WarningList[ToolkitWarning]()
         self.tracker = Tracker(skip_tracking)
+        self._additional_tracking_info = CommandTrackingInfo()
 
     @property
     def print_warning(self) -> bool:
         return self._print_warning and not self.silent
 
     def _track_command(self, result: str | Exception) -> None:
-        self.tracker.track_cli_command(self.warning_list, result, type(self).__name__.removesuffix("Command"))
+        self.tracker.track_cli_command(
+            self.warning_list,
+            result,
+            type(self).__name__.removesuffix("Command"),
+            self._additional_tracking_info,
+        )
 
     def run(self, execute: Callable[..., Any], *args: Any, **kwargs: Any) -> Any:
         global _HAS_PRINTED_COLLECT_MESSAGE

cognite_toolkit/_cdf_tk/commands/_changes.py

@@ -11,9 +11,8 @@ from packaging.version import Version
 from packaging.version import parse as parse_version
 from rich import print
 
-from cognite_toolkit._cdf_tk.builders import get_loader
+from cognite_toolkit._cdf_tk.cdf_toml import CDFToml
 from cognite_toolkit._cdf_tk.constants import DOCKER_IMAGE_NAME
-from cognite_toolkit._cdf_tk.data_classes import ModuleDirectories
 from cognite_toolkit._cdf_tk.utils import iterate_modules, read_yaml_file, safe_read, safe_write
 from cognite_toolkit._version import __version__
 
@@ -51,40 +50,6 @@ class ManualChange(Change):
         return ""
 
 
-class SetKindOnFile(AutomaticChange):
-    """Adds the kind to the filename of all resource files.
-
-Before `your_file.yaml`:
-After `your_file.FileMetadata.yaml`:
-    """
-
-    deprecated_from = Version("0.4.0")
-    has_file_changes = True
-
-    def do(self) -> set[Path]:
-        module_directories = ModuleDirectories.load(self._organization_dir)
-        changed: set[Path] = set()
-        for module in module_directories:
-            for resource_folder, source_files in module.source_paths_by_resource_folder.items():
-                for source_file in source_files:
-                    if source_file.suffix not in {".yaml", ".yml"}:
-                        continue
-                    loader, _ = get_loader(source_file, resource_folder, force_pattern=True)
-                    if loader is None:
-                        print(f"Could not find loader for {source_file}")
-                        continue
-                    if source_file.stem.casefold().endswith(loader.kind.casefold()):
-                        continue
-                    new_name = source_file.with_name(f"{source_file.stem}.{loader.kind}{source_file.suffix}")
-                    source_file.rename(new_name)
-                    changed.add(source_file)
-                    for suffix in [".sql", ".csv", ".parquet"]:
-                        if (adjacent_file := source_file.with_suffix(suffix)).exists():
-                            adjacent_file.rename(new_name.with_suffix(suffix))
-                            changed.add(adjacent_file)
-
-        return changed
-
 
 class FixViewBasedLocationFilter(AutomaticChange):
     """The created view-based location filter has been fixed to be compatible with the CDF API.
@@ -362,7 +327,6 @@ After:
 
     def do(self) -> set[Path]:
         # Avoid circular import
-        from .modules import ModulesCommand
 
         system_yaml = self._organization_dir / "_system.yaml"
         if not system_yaml.exists():
@@ -372,11 +336,8 @@ After:
         content = read_yaml_file(system_yaml)
         current_version = content.get("cdf_toolkit_version", __version__)
 
-        cdf_toml_content = ModulesCommand(skip_tracking=True).create_cdf_toml(self._organization_dir)
-        cdf_toml_content = cdf_toml_content.replace(f'version = "{__version__}"', f'version = "{current_version}"')
-
-        cdf_toml_path = Path.cwd() / "cdf.toml"
-        cdf_toml_path.write_text(cdf_toml_content, encoding="utf-8")
+        CDFToml.write(self._organization_dir, version=current_version)
+        cdf_toml_path = Path.cwd() / CDFToml.file_name
         system_yaml.unlink()
         return {cdf_toml_path, system_yaml}
 

cognite_toolkit/_cdf_tk/commands/_download.py

@@ -1,17 +1,27 @@
-from collections import Counter
-from collections.abc import Iterable
+from collections.abc import Callable, Iterable
 from functools import partial
 from pathlib import Path
 
-from cognite.client.data_classes._base import T_CogniteResourceList
-from rich.console import Console
-
+from cognite_toolkit._cdf_tk.constants import DATA_MANIFEST_STEM, DATA_RESOURCE_DIR
 from cognite_toolkit._cdf_tk.exceptions import ToolkitValueError
-from cognite_toolkit._cdf_tk.storageio import StorageIO, TableStorageIO
-from cognite_toolkit._cdf_tk.utils.file import safe_write, to_directory_compatible, yaml_safe_dump
-from cognite_toolkit._cdf_tk.utils.fileio import TABLE_WRITE_CLS_BY_FORMAT, Compression, FileWriter, SchemaColumn
+from cognite_toolkit._cdf_tk.protocols import T_ResourceResponse
+from cognite_toolkit._cdf_tk.storageio import (
+    ConfigurableStorageIO,
+    Page,
+    StorageIO,
+    T_Selector,
+    TableStorageIO,
+)
+from cognite_toolkit._cdf_tk.tk_warnings import LowSeverityWarning
+from cognite_toolkit._cdf_tk.utils.file import safe_write, sanitize_filename, yaml_safe_dump
+from cognite_toolkit._cdf_tk.utils.fileio import (
+    TABLE_WRITE_CLS_BY_FORMAT,
+    Compression,
+    FileWriter,
+    SchemaColumn,
+)
 from cognite_toolkit._cdf_tk.utils.producer_worker import ProducerWorkerExecutor
-from cognite_toolkit._cdf_tk.utils.useful_types import T_ID, JsonVal, T_Selector, T_WritableCogniteResourceList
+from cognite_toolkit._cdf_tk.utils.useful_types import JsonVal
 
 from ._base import ToolkitCommand
 
@@ -20,7 +30,7 @@ class DownloadCommand(ToolkitCommand):
     def download(
         self,
         selectors: Iterable[T_Selector],
-        io: StorageIO[T_ID, T_Selector, T_CogniteResourceList, T_WritableCogniteResourceList],
+        io: StorageIO[T_Selector, T_ResourceResponse],
         output_dir: Path,
         verbose: bool,
         file_format: str,
@@ -38,59 +48,64 @@ class DownloadCommand(ToolkitCommand):
             compression: The compression method to use for the downloaded files (e.g., "none", "gzip").
             limit: The maximum number of items to download for each selected set. If None, all items will be downloaded.
         """
-        target_directory = output_dir / io.FOLDER_NAME
-        target_directory.mkdir(parents=True, exist_ok=True)
         compression_cls = Compression.from_name(compression)
 
-        console = Console()
-        filestem_counter: dict[str, int] = Counter()
+        console = io.client.console
         for selector in selectors:
+            target_dir = output_dir / sanitize_filename(selector.group)
             if verbose:
-                console.print(f"Downloading {io.DISPLAY_NAME} '{selector!s}' to {target_directory.as_posix()!r}")
+                console.print(f"Downloading {selector.display_name} '{selector!s}' to {target_dir.as_posix()!r}")
 
-            filestem = to_directory_compatible(str(selector))
-            if filestem_counter[filestem] > 0:
-                filestem = f"{filestem}_{filestem_counter[filestem]}"
-            filestem_counter[filestem] += 1
             iteration_count = self._get_iteration_count(io, selector, limit)
+            filestem = sanitize_filename(str(selector))
+            if self._already_downloaded(target_dir, filestem):
+                warning = LowSeverityWarning(
+                    f"Data for {selector!s} already exists in {target_dir.as_posix()!r}. Skipping download."
+                )
+                self.warn(warning, console=console)
+                continue
 
+            selector.dump_to_file(target_dir)
             columns: list[SchemaColumn] | None = None
-            if file_format in TABLE_WRITE_CLS_BY_FORMAT and isinstance(io, TableStorageIO):
+            is_table = file_format in TABLE_WRITE_CLS_BY_FORMAT
+            if is_table and isinstance(io, TableStorageIO):
                 columns = io.get_schema(selector)
-            elif file_format in TABLE_WRITE_CLS_BY_FORMAT:
+            elif is_table:
                 raise ToolkitValueError(
-                    f"Cannot download {io.KIND} in {file_format!r} format. The {io.KIND!r} storage type does not support table schemas."
+                    f"Cannot download {selector.kind} in {file_format!r} format. The {selector.kind!r} storage type does not support table schemas."
                 )
 
             with FileWriter.create_from_format(
-                file_format, target_directory, io.KIND, compression_cls, columns=columns
+                file_format, target_dir, selector.kind, compression_cls, columns=columns
             ) as writer:
-                executor = ProducerWorkerExecutor[T_WritableCogniteResourceList, list[dict[str, JsonVal]]](
+                executor = ProducerWorkerExecutor[Page[T_ResourceResponse], list[dict[str, JsonVal]]](
                     download_iterable=io.stream_data(selector, limit),
-                    process=io.data_to_json_chunk,
+                    process=self.create_data_process(io=io, selector=selector, is_table=is_table),
                     write=partial(writer.write_chunks, filestem=filestem),
                     iteration_count=iteration_count,
                     # Limit queue size to avoid filling up memory before the workers can write to disk.
                     max_queue_size=8 * 10,  # 8 workers, 10 items per worker
                     download_description=f"Downloading {selector!s}",
                     process_description="Processing",
-                    write_description=f"Writing to {target_directory.as_posix()!r} in files with stem {filestem!r}",
+                    write_description=f"Writing to {target_dir.as_posix()!r} in files with stem {filestem!r}",
                     console=console,
                 )
                 executor.run()
                 executor.raise_on_error()
                 file_count = writer.file_count
 
-            for config in io.configurations(selector):
-                config_file = output_dir / config.folder_name / f"{filestem}.{config.kind}.yaml"
-                config_file.parent.mkdir(parents=True, exist_ok=True)
-                safe_write(config_file, yaml_safe_dump(config.value))
+            if isinstance(io, ConfigurableStorageIO):
+                for config in io.configurations(selector):
+                    filename = config.filename or filestem
+                    config_file = target_dir / DATA_RESOURCE_DIR / config.folder_name / f"{filename}.{config.kind}.yaml"
+                    config_file.parent.mkdir(parents=True, exist_ok=True)
+                    safe_write(config_file, yaml_safe_dump(config.value))
 
-            console.print(f"Downloaded {selector!s} to {file_count} file(s) in {target_directory.as_posix()!r}.")
+            console.print(f"Downloaded {selector!s} to {file_count} file(s) in {target_dir.as_posix()!r}.")
 
     @staticmethod
     def _get_iteration_count(
-        io: StorageIO[T_ID, T_Selector, T_CogniteResourceList, T_WritableCogniteResourceList],
+        io: StorageIO[T_Selector, T_ResourceResponse],
         selector: T_Selector,
         limit: int | None,
     ) -> int | None:
@@ -101,3 +116,39 @@ class DownloadCommand(ToolkitCommand):
         if total is not None:
             iteration_count = total // io.CHUNK_SIZE + (1 if total % io.CHUNK_SIZE > 0 else 0)
         return iteration_count
+
+    @staticmethod
+    def _already_downloaded(output_dir: Path, filestem: str) -> bool:
+        if not output_dir.exists():
+            return False
+
+        # Check for multi-part files (e.g. ndjson, csv, parquet)
+        if any(output_dir.glob(f"{filestem}-part-*")):
+            return True
+
+        # Check for single files (e.g. yaml) and exclude the metadata file.
+        manifest_file_name = f"{filestem}.{DATA_MANIFEST_STEM}.yaml"
+        for f in output_dir.glob(f"{filestem}.*"):
+            if f.name != manifest_file_name:
+                return True
+
+        return False
+
+    @staticmethod
+    def create_data_process(
+        io: StorageIO[T_Selector, T_ResourceResponse],
+        selector: T_Selector,
+        is_table: bool,
+    ) -> Callable[[Page[T_ResourceResponse]], list[dict[str, JsonVal]]]:
+        """Creates a data processing function based on the IO type and whether the output is a table."""
+        if is_table and isinstance(io, TableStorageIO):
+
+            def row_data_process(chunk: Page[T_ResourceResponse]) -> list[dict[str, JsonVal]]:
+                return io.data_to_row(chunk.items, selector)
+
+            return row_data_process
+
+        def chunk_data_process(data_page: Page[T_ResourceResponse]) -> list[dict[str, JsonVal]]:
+            return io.data_to_json_chunk(data_page.items, selector)
+
+        return chunk_data_process

cognite_toolkit/_cdf_tk/commands/_migrate/__init__.py

@@ -1,13 +1,7 @@
-from .canvas import MigrationCanvasCommand
 from .command import MigrationCommand
-from .files import MigrateFilesCommand
 from .prepare import MigrationPrepareCommand
-from .timeseries import MigrateTimeseriesCommand
 
 __all__ = [
-    "MigrateFilesCommand",
-    "MigrateTimeseriesCommand",
-    "MigrationCanvasCommand",
     "MigrationCommand",
     "MigrationPrepareCommand",
 ]