cognite-toolkit 0.6.105__tar.gz → 0.6.106__tar.gz

{cognite_toolkit-0.6.105 → cognite_toolkit-0.6.106}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cognite_toolkit
-Version: 0.6.105
+Version: 0.6.106
 Summary: Official Cognite Data Fusion tool for project templates and configuration deployment
 Project-URL: Homepage, https://docs.cognite.com/cdf/deploy/cdf_toolkit/
 Project-URL: Changelog, https://github.com/cognitedata/toolkit/releases

{cognite_toolkit-0.6.105 → cognite_toolkit-0.6.106}/cdf.toml

@@ -33,7 +33,7 @@ dev = true
 [modules]
 # This is the version of the modules. It should not be changed manually.
 # It will be updated by the 'cdf modules upgrade' command.
-version = "0.6.105"
+version = "0.6.106"
 
 
 [library.toolkit-data]

{cognite_toolkit-0.6.105 → cognite_toolkit-0.6.106}/cognite_toolkit/_cdf_tk/commands/_migrate/conversion.py

@@ -15,7 +15,10 @@ from cognite.client.data_classes.data_modeling.views import ViewProperty
 from cognite.client.utils._identifier import InstanceId
 
 from cognite_toolkit._cdf_tk.client import ToolkitClient
-from cognite_toolkit._cdf_tk.client.data_classes.migration import AssetCentricId, ResourceViewMapping
+from cognite_toolkit._cdf_tk.client.data_classes.migration import (
+    AssetCentricId,
+    ResourceViewMappingApply,
+)
 from cognite_toolkit._cdf_tk.utils.collection import flatten_dict_json_path
 from cognite_toolkit._cdf_tk.utils.dtype_conversion import (
     asset_centric_convert_to_primary_property,
@@ -163,7 +166,7 @@ class DirectRelationCache:
 def asset_centric_to_dm(
     resource: AssetCentricResourceExtended,
     instance_id: InstanceId,
-    view_source: ResourceViewMapping,
+    view_source: ResourceViewMappingApply,
     view_properties: dict[str, ViewProperty],
     direct_relation_cache: DirectRelationCache,
 ) -> tuple[NodeApply | EdgeApply | None, ConversionIssue]:
@@ -172,7 +175,7 @@ def asset_centric_to_dm(
     Args:
         resource (CogniteResource): The asset-centric resource to convert.
         instance_id (NodeId | EdgeApply): The ID of the instance to create or update.
-        view_source (ResourceViewMapping): The view source defining how to map the resource to the data model.
+        view_source (ResourceViewMappingApply): The view source defining how to map the resource to the data model.
         view_properties (dict[str, ViewProperty]): The defined properties referenced in the view source mapping.
         direct_relation_cache (DirectRelationCache): Cache for direct relation references.
 

{cognite_toolkit-0.6.105 → cognite_toolkit-0.6.106}/cognite_toolkit/_cdf_tk/commands/_migrate/data_classes.py

@@ -14,7 +14,11 @@ from pydantic import BaseModel, BeforeValidator, field_validator, model_validato
 from cognite_toolkit._cdf_tk.client.data_classes.instances import InstanceApplyList
 from cognite_toolkit._cdf_tk.client.data_classes.migration import AssetCentricId
 from cognite_toolkit._cdf_tk.client.data_classes.pending_instances_ids import PendingInstanceId
-from cognite_toolkit._cdf_tk.commands._migrate.default_mappings import create_default_mappings
+from cognite_toolkit._cdf_tk.commands._migrate.default_mappings import (
+    ASSET_ANNOTATIONS_ID,
+    FILE_ANNOTATIONS_ID,
+    create_default_mappings,
+)
 from cognite_toolkit._cdf_tk.exceptions import ToolkitValueError
 from cognite_toolkit._cdf_tk.storageio._data_classes import ModelList
 from cognite_toolkit._cdf_tk.utils.useful_types import (
@@ -182,6 +186,18 @@ class FileMapping(MigrationMapping):
 class AnnotationMapping(MigrationMapping):
     resource_type: Literal["annotation"] = "annotation"
     instance_id: EdgeId
+    annotation_type: Literal["diagrams.AssetLink", "diagrams.FileLink"] | None = None
+
+    def get_ingestion_view(self) -> str:
+        """Get the ingestion view for the mapping. If not specified, return the default ingestion view."""
+        if self.ingestion_view:
+            return self.ingestion_view
+        elif self.annotation_type == "diagrams.AssetLink":
+            return ASSET_ANNOTATIONS_ID
+        elif self.annotation_type == "diagrams.FileLink":
+            return FILE_ANNOTATIONS_ID
+        else:
+            raise ToolkitValueError("Cannot determine default ingestion view for annotation without annotation_type")
 
     @field_validator("instance_id", mode="before")
     def _validate_instance_id(cls, v: Any) -> Any:

{cognite_toolkit-0.6.105 → cognite_toolkit-0.6.106}/cognite_toolkit/_cdf_tk/commands/_migrate/data_mapper.py

@@ -14,9 +14,10 @@ from cognite.client.data_classes.data_modeling import (
 )
 
 from cognite_toolkit._cdf_tk.client import ToolkitClient
-from cognite_toolkit._cdf_tk.client.data_classes.migration import ResourceViewMapping
+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
+from cognite_toolkit._cdf_tk.commands._migrate.default_mappings import create_default_mappings
 from cognite_toolkit._cdf_tk.commands._migrate.issues import ConversionIssue, MigrationIssue
 from cognite_toolkit._cdf_tk.commands._migrate.selectors import AssetCentricMigrationSelector
 from cognite_toolkit._cdf_tk.constants import MISSING_INSTANCE_SPACE
@@ -61,20 +62,22 @@ class AssetCentricMapper(
     def __init__(self, client: ToolkitClient) -> None:
         self.client = client
         self._ingestion_view_by_id: dict[ViewId, View] = {}
-        self._view_mapping_by_id: dict[str, ResourceViewMapping] = {}
+        self._view_mapping_by_id: dict[str, ResourceViewMappingApply] = {}
         self._direct_relation_cache = DirectRelationCache(client)
 
     def prepare(self, source_selector: AssetCentricMigrationSelector) -> None:
         ingestion_view_ids = source_selector.get_ingestion_mappings()
         ingestion_views = self.client.migration.resource_view_mapping.retrieve(ingestion_view_ids)
-        self._view_mapping_by_id = {view.external_id: view for view in ingestion_views}
+        defaults = {mapping.external_id: mapping for mapping in create_default_mappings()}
+        # Custom mappings from CDF override the default mappings
+        self._view_mapping_by_id = defaults | {view.external_id: view.as_write() for view in ingestion_views}
         missing_mappings = set(ingestion_view_ids) - set(self._view_mapping_by_id.keys())
         if missing_mappings:
             raise ToolkitValueError(
                 f"The following ingestion views were not found: {humanize_collection(missing_mappings)}"
             )
 
-        view_ids = list({view.view_id for view in ingestion_views})
+        view_ids = list({mapping.view_id for mapping in self._view_mapping_by_id.values()})
         views = self.client.data_modeling.views.retrieve(view_ids)
         self._ingestion_view_by_id = {view.as_id(): view for view in views}
         missing_views = set(view_ids) - set(self._ingestion_view_by_id.keys())

{cognite_toolkit-0.6.105 → cognite_toolkit-0.6.106}/cognite_toolkit/_cdf_tk/commands/_migrate/default_mappings.py

@@ -93,7 +93,7 @@ def create_default_mappings() -> list[ResourceViewMappingApply]:
             view_id=ViewId("cdf_cdm", "CogniteDiagramAnnotation", "v1"),
             property_mapping={
                 # We are ignoring the symbol region in the default mapping.
-                "annotatedResource.id": "edge.startNode",
+                "annotatedResourceId": "edge.startNode",
                 "annotationType": "edge.type.externalId",
                 "creatingUser": "sourceCreatedUser",
                 "creatingApp": "sourceId",
@@ -114,16 +114,16 @@ def create_default_mappings() -> list[ResourceViewMappingApply]:
         ResourceViewMappingApply(
             external_id=FILE_ANNOTATIONS_ID,
             resource_type="fileAnnotation",
-            view_id=ViewId("cdf_cdm", "CogniteFileAnnotation", "v1"),
+            view_id=ViewId("cdf_cdm", "CogniteDiagramAnnotation", "v1"),
             property_mapping={
-                "annotatedResource.id": "edge.startNode",
+                "annotatedResourceId": "edge.startNode",
                 "annotationType": "edge.type.externalId",
                 "creatingUser": "sourceCreatedUser",
                 "creatingApp": "sourceId",
                 "creatingAppVersion": "sourceContext",
                 "status": "status",
-                "data.fileRef.id": "edge.startNode",
-                "data.fileRef.externalId": "edge.startNode",
+                "data.fileRef.id": "edge.endNode",
+                "data.fileRef.externalId": "edge.endNode",
                 "data.description": "description",
                 "data.pageNumber": "startNodePageNumber",
                 "data.textRegion.confidence": "confidence",

{cognite_toolkit-0.6.105 → cognite_toolkit-0.6.106}/cognite_toolkit/_cdf_tk/commands/_migrate/migration_io.py

@@ -33,6 +33,7 @@ from .data_classes import (
     MigrationMappingList,
 )
 from .data_model import INSTANCE_SOURCE_VIEW_ID
+from .default_mappings import ASSET_ANNOTATIONS_ID, FILE_ANNOTATIONS_ID
 from .selectors import AssetCentricMigrationSelector, MigrateDataSetSelector, MigrationCSVFileSelector
 
 
@@ -213,6 +214,16 @@ class AssetCentricMigrationIO(
 class AnnotationMigrationIO(
     UploadableStorageIO[AssetCentricMigrationSelector, AssetCentricMapping[Annotation], InstanceApply]
 ):
+    """IO class for migrating Annotations.
+
+    Args:
+        client: The ToolkitClient to use for CDF interactions.
+        instance_space: The instance space to use for the migrated annotations.
+        default_asset_annotation_mapping: The default ingestion mapping to use for asset-linked annotations.
+        default_file_annotation_mapping: The default ingestion mappingto use for file-linked annotations.
+
+    """
+
     KIND = "AnnotationMigration"
     SUPPORTED_DOWNLOAD_FORMATS = frozenset({".parquet", ".csv", ".ndjson"})
     SUPPORTED_COMPRESSIONS = frozenset({".gz"})
@@ -220,17 +231,28 @@ class AnnotationMigrationIO(
     CHUNK_SIZE = 1000
     UPLOAD_ENDPOINT = InstanceIO.UPLOAD_ENDPOINT
 
-    def __init__(self, client: ToolkitClient, instance_space: str | None = None) -> None:
+    def __init__(
+        self,
+        client: ToolkitClient,
+        instance_space: str | None = None,
+        default_asset_annotation_mapping: str = ASSET_ANNOTATIONS_ID,
+        default_file_annotation_mapping: str = FILE_ANNOTATIONS_ID,
+    ) -> None:
         super().__init__(client)
         self.annotation_io = AnnotationIO(client)
         self.instance_space = instance_space
+        self.default_asset_annotation_mapping = default_asset_annotation_mapping
+        self.default_file_annotation_mapping = default_file_annotation_mapping
 
     def as_id(self, item: AssetCentricMapping[Annotation]) -> str:
         return f"Annotation_{item.mapping.id}"
 
     def count(self, selector: AssetCentricMigrationSelector) -> int | None:
-        # There is no efficient way to count annotations in CDF.
-        return None
+        if isinstance(selector, MigrationCSVFileSelector):
+            return len(selector.items)
+        else:
+            # There is no efficient way to count annotations in CDF.
+            return None
 
     def stream_data(self, selector: AssetCentricMigrationSelector, limit: int | None = None) -> Iterable[Page]:
         if isinstance(selector, MigrateDataSetSelector):
@@ -253,8 +275,10 @@ class AnnotationMigrationIO(
                 mapping = AnnotationMapping(
                     instance_id=EdgeId(space=self.instance_space, external_id=f"annotation_{resource.id!r}"),
                     id=resource.id,
-                    ingestion_view=selector.ingestion_mapping,
+                    ingestion_view=self._get_mapping(selector.ingestion_mapping, resource),
                     preferred_consumer_view=selector.preferred_consumer_view,
+                    # The PySDK is poorly typed.
+                    annotation_type=resource.annotation_type,  # type: ignore[arg-type]
                 )
                 mapping_list.append(AssetCentricMapping(mapping=mapping, resource=resource))
             yield mapping_list
@@ -275,6 +299,7 @@ class AnnotationMigrationIO(
                 if resource is None:
                     not_found += 1
                     continue
+                mapping.ingestion_view = self._get_mapping(mapping.ingestion_view, resource)
                 chunk.append(AssetCentricMapping(mapping=mapping, resource=resource))
             if chunk:
                 yield chunk
@@ -284,6 +309,21 @@ class AnnotationMigrationIO(
                     f"Could not find {not_found} annotations referenced in the CSV file. They will be skipped during migration."
                 ).print_warning(include_timestamp=True, console=self.client.console)
 
+    def _get_mapping(self, current_mapping: str | None, resource: Annotation) -> str:
+        try:
+            return (
+                current_mapping
+                or {
+                    "diagrams.AssetLink": self.default_asset_annotation_mapping,
+                    "diagrams.FileLink": self.default_file_annotation_mapping,
+                }[resource.annotation_type]
+            )
+        except KeyError as e:
+            raise ToolkitValueError(
+                f"Could not determine default ingestion view for annotation type '{resource.annotation_type}'. "
+                "Please specify the ingestion view explicitly in the CSV file."
+            ) from e
+
     def json_to_resource(self, item_json: dict[str, JsonVal]) -> InstanceApply:
         raise NotImplementedError("Deserializing Annotation Migrations from JSON is not supported.")
 

{cognite_toolkit-0.6.105 → cognite_toolkit-0.6.106}/cognite_toolkit/_cdf_tk/commands/_upload.py

@@ -27,7 +27,7 @@ from cognite_toolkit._cdf_tk.tk_warnings import HighSeverityWarning, MediumSever
 from cognite_toolkit._cdf_tk.tk_warnings.fileread import ResourceFormatWarning
 from cognite_toolkit._cdf_tk.utils.auth import EnvironmentVariables
 from cognite_toolkit._cdf_tk.utils.file import read_yaml_file
-from cognite_toolkit._cdf_tk.utils.fileio import TABLE_READ_CLS_BY_FORMAT, FileReader
+from cognite_toolkit._cdf_tk.utils.fileio import MultiFileReader
 from cognite_toolkit._cdf_tk.utils.http_client import HTTPClient, ItemMessage, SuccessResponseItems
 from cognite_toolkit._cdf_tk.utils.producer_worker import ProducerWorkerExecutor
 from cognite_toolkit._cdf_tk.utils.progress_tracker import ProgressTracker
@@ -87,7 +87,7 @@ class UploadCommand(ToolkitCommand):
         ├── datafile2.Manifest.yaml       # Manifest file for datafile2
         └── ...
         """
-        console = Console()
+        console = client.console
         data_files_by_selector = self._find_data_files(input_dir, kind)
 
         self._deploy_resource_folder(input_dir / DATA_RESOURCE_DIR, deploy_resources, client, console, dry_run, verbose)
@@ -225,49 +225,44 @@ class UploadCommand(ToolkitCommand):
                 io = self._create_selected_io(selector, datafiles[0], client)
                 if io is None:
                     continue
-                for data_file in datafiles:
-                    file_display = self._path_as_display_name(data_file)
-                    if verbose:
-                        console.print(f"{action} {selector.display_name} from {file_display.as_posix()!r}")
-                    reader = FileReader.from_filepath(data_file)
-                    is_table = reader.format in TABLE_READ_CLS_BY_FORMAT
-                    if is_table and not isinstance(io, TableUploadableStorageIO):
-                        raise ToolkitValueError(f"{selector.display_name} does not support {reader.format!r} files.")
-                    tracker = ProgressTracker[str]([self._UPLOAD])
-                    executor = ProducerWorkerExecutor[list[tuple[str, dict[str, JsonVal]]], Sequence[UploadItem]](
-                        download_iterable=io.read_chunks(reader),
-                        process=partial(io.rows_to_data, selector=selector)
-                        if is_table and isinstance(io, TableUploadableStorageIO)
-                        else io.json_chunk_to_data,
-                        write=partial(
-                            self._upload_items,
-                            upload_client=upload_client,
-                            io=io,
-                            dry_run=dry_run,
-                            selector=selector,
-                            tracker=tracker,
-                            console=console,
-                        ),
-                        iteration_count=None,
-                        max_queue_size=self._MAX_QUEUE_SIZE,
-                        download_description=f"Reading {file_count:,}/{total_file_count + 1:,}: {file_display.as_posix()!s}",
-                        process_description="Processing",
-                        write_description=f"{action} {selector.display_name!r}",
+                reader = MultiFileReader(datafiles)
+                if reader.is_table and not isinstance(io, TableUploadableStorageIO):
+                    raise ToolkitValueError(f"{selector.display_name} does not support {reader.format!r} files.")
+                tracker = ProgressTracker[str]([self._UPLOAD])
+                executor = ProducerWorkerExecutor[list[tuple[str, dict[str, JsonVal]]], Sequence[UploadItem]](
+                    download_iterable=io.read_chunks(reader),
+                    process=partial(io.rows_to_data, selector=selector)
+                    if reader.is_table and isinstance(io, TableUploadableStorageIO)
+                    else io.json_chunk_to_data,
+                    write=partial(
+                        self._upload_items,
+                        upload_client=upload_client,
+                        io=io,
+                        dry_run=dry_run,
+                        selector=selector,
+                        tracker=tracker,
                         console=console,
-                    )
-                    executor.run()
-                    file_count += 1
-                    executor.raise_on_error()
-                    final_action = "Uploaded" if not dry_run else "Would upload"
-                    suffix = " successfully" if not dry_run else ""
-                    results = tracker.aggregate()
-                    success = results.get((self._UPLOAD, "success"), 0)
-                    failed = results.get((self._UPLOAD, "failed"), 0)
-                    if failed > 0:
-                        suffix += f", {failed:,} failed"
-                    console.print(
-                        f"{final_action} {success:,} {selector.display_name} from {file_display.as_posix()!r}{suffix}."
-                    )
+                    ),
+                    iteration_count=None,
+                    max_queue_size=self._MAX_QUEUE_SIZE,
+                    download_description=f"Reading {selector.display_name!r} files",
+                    process_description="Processing",
+                    write_description=f"{action} {selector.display_name!r}",
+                    console=console,
+                )
+                executor.run()
+                file_count += len(datafiles)
+                executor.raise_on_error()
+                final_action = "Uploaded" if not dry_run else "Would upload"
+                suffix = " successfully" if not dry_run else ""
+                results = tracker.aggregate()
+                success = results.get((self._UPLOAD, "success"), 0)
+                failed = results.get((self._UPLOAD, "failed"), 0)
+                if failed > 0:
+                    suffix += f", {failed:,} failed"
+                console.print(
+                    f"{final_action} {success:,} {selector.display_name} from {len(datafiles)} files{suffix}."
+                )
 
     @staticmethod
     def _path_as_display_name(input_path: Path, cwd: Path = Path.cwd()) -> Path:

{cognite_toolkit-0.6.105 → cognite_toolkit-0.6.106}/cognite_toolkit/_cdf_tk/cruds/_resource_cruds/datamodel.py

@@ -328,6 +328,7 @@ class ContainerCRUD(ResourceContainerCRUD[ContainerId, ContainerApply, Container
 
     def dump_resource(self, resource: Container, local: dict[str, Any] | None = None) -> dict[str, Any]:
         dumped = resource.as_write().dump()
+        has_local = local is not None
         local = local or {}
         for key in ["constraints", "indexes"]:
             if not dumped.get(key) and key not in local:
@@ -339,15 +340,16 @@ class ContainerCRUD(ResourceContainerCRUD[ContainerId, ContainerApply, Container
                 continue
             local_prop = local_prop_by_id[prop_id]
             for key, default in [("immutable", False), ("autoIncrement", False), ("nullable", True)]:
-                if cdf_prop.get(key) is default and key not in local_prop:
+                if has_local and cdf_prop.get(key) is default and key not in local_prop:
                     cdf_prop.pop(key, None)
             cdf_type = cdf_prop.get("type", {})
             local_type = local_prop.get("type", {})
             for key, type_default in [("list", False), ("collation", "ucs_basic")]:
-                if cdf_type.get(key) == type_default and key not in local_type:
+                if has_local and cdf_type.get(key) == type_default and key not in local_type:
                     cdf_type.pop(key, None)
-        if "usedFor" not in local:
-            dumped.pop("usedFor", None)
+            if has_local and "usedFor" not in local and dumped.get("usedFor") == "node":
+                # Only drop if set to default by server.
+                dumped.pop("usedFor", None)
         return dumped
 
     def create(self, items: Sequence[ContainerApply]) -> ContainerList:

{cognite_toolkit-0.6.105 → cognite_toolkit-0.6.106}/cognite_toolkit/_cdf_tk/storageio/_annotations.py

@@ -26,21 +26,23 @@ class AnnotationIO(StorageIO[AssetCentricSelector, Annotation]):
     def stream_data(self, selector: AssetCentricSelector, limit: int | None = None) -> Iterable[Page[Annotation]]:
         total = 0
         for file_chunk in FileMetadataIO(self.client).stream_data(selector, None):
-            # Todo Support pagination. This is missing in the SDK.
-            results = self.client.annotations.list(
-                filter=AnnotationFilter(
-                    annotated_resource_type="file",
-                    annotated_resource_ids=[{"id": file_metadata.id} for file_metadata in file_chunk.items],
+            for annotation_type in ["diagrams.AssetLink", "diagrams.FileLink"]:
+                # Todo Support pagination. This is missing in the SDK.
+                results = self.client.annotations.list(
+                    filter=AnnotationFilter(
+                        annotated_resource_type="file",
+                        annotated_resource_ids=[{"id": file_metadata.id} for file_metadata in file_chunk.items],
+                        annotation_type=annotation_type,
+                    )
                 )
-            )
-            if limit is not None and total + len(results) > limit:
-                results = results[: limit - total]
-
-            for chunk in chunker_sequence(results, self.CHUNK_SIZE):
-                yield Page(worker_id="main", items=chunk)
-                total += len(chunk)
-            if limit is not None and total >= limit:
-                break
+                if limit is not None and total + len(results) > limit:
+                    results = results[: limit - total]
+
+                for chunk in chunker_sequence(results, self.CHUNK_SIZE):
+                    yield Page(worker_id="main", items=chunk)
+                    total += len(chunk)
+                if limit is not None and total >= limit:
+                    break
 
     def count(self, selector: AssetCentricSelector) -> int | None:
         """There is no efficient way to count annotations in CDF."""

{cognite_toolkit-0.6.105 → cognite_toolkit-0.6.106}/cognite_toolkit/_cdf_tk/storageio/_base.py

@@ -8,8 +8,7 @@ from cognite.client.data_classes._base import T_CogniteResource
 from cognite_toolkit._cdf_tk.client import ToolkitClient
 from cognite_toolkit._cdf_tk.exceptions import ToolkitNotImplementedError
 from cognite_toolkit._cdf_tk.utils.collection import chunker
-from cognite_toolkit._cdf_tk.utils.fileio import FileReader, SchemaColumn
-from cognite_toolkit._cdf_tk.utils.fileio._readers import TableReader
+from cognite_toolkit._cdf_tk.utils.fileio import MultiFileReader, SchemaColumn
 from cognite_toolkit._cdf_tk.utils.http_client import HTTPClient, HTTPMessage, ItemsRequest
 from cognite_toolkit._cdf_tk.utils.useful_types import JsonVal, T_WriteCogniteResource
 
@@ -217,8 +216,8 @@ class UploadableStorageIO(
         raise NotImplementedError()
 
     @classmethod
-    def read_chunks(cls, reader: FileReader) -> Iterable[list[tuple[str, dict[str, JsonVal]]]]:
-        data_name = "row" if isinstance(reader, TableReader) else "line"
+    def read_chunks(cls, reader: MultiFileReader) -> Iterable[list[tuple[str, dict[str, JsonVal]]]]:
+        data_name = "row" if reader.is_table else "line"
         # Include name of line for better error messages
         iterable = ((f"{data_name} {line_no}", item) for line_no, item in reader.read_chunks_with_line_numbers())
 

{cognite_toolkit-0.6.105 → cognite_toolkit-0.6.106}/cognite_toolkit/_cdf_tk/storageio/_datapoints.py

@@ -19,8 +19,7 @@ from cognite_toolkit._cdf_tk.utils.dtype_conversion import (
     _TextConverter,
     _ValueConverter,
 )
-from cognite_toolkit._cdf_tk.utils.fileio import FileReader
-from cognite_toolkit._cdf_tk.utils.fileio._readers import TableReader
+from cognite_toolkit._cdf_tk.utils.fileio._readers import MultiFileReader
 from cognite_toolkit._cdf_tk.utils.http_client import DataBodyRequest, HTTPClient, HTTPMessage
 from cognite_toolkit._cdf_tk.utils.useful_types import JsonVal
 
@@ -165,8 +164,8 @@ class DatapointsIO(TableUploadableStorageIO[DataPointsFileSelector, DataPointLis
         )
 
     @classmethod
-    def read_chunks(cls, reader: FileReader) -> Iterable[list[tuple[str, dict[str, JsonVal]]]]:
-        if not isinstance(reader, TableReader):
+    def read_chunks(cls, reader: MultiFileReader) -> Iterable[list[tuple[str, dict[str, JsonVal]]]]:
+        if not reader.is_table:
             raise RuntimeError("DatapointsIO can only read from TableReader instances.")
         iterator = iter(reader.read_chunks_with_line_numbers())
         try:

{cognite_toolkit-0.6.105 → cognite_toolkit-0.6.106}/cognite_toolkit/_cdf_tk/utils/fileio/__init__.py

@@ -12,6 +12,7 @@ from ._readers import (
     CSVReader,
     FailedParsing,
     FileReader,
+    MultiFileReader,
     NDJsonReader,
     ParquetReader,
     YAMLReader,
@@ -45,6 +46,7 @@ __all__ = [
     "FileReader",
     "FileWriter",
     "GzipCompression",
+    "MultiFileReader",
     "NDJsonReader",
     "NDJsonWriter",
     "ParquetReader",

{cognite_toolkit-0.6.105 → cognite_toolkit-0.6.106}/cognite_toolkit/_cdf_tk/utils/fileio/_base.py

@@ -14,7 +14,11 @@ T_IO = TypeVar("T_IO", bound=IO)
 
 
 class FileIO(ABC):
-    format: ClassVar[str]
+    FORMAT: ClassVar[str]
+
+    @property
+    def format(self) -> str:
+        return self.FORMAT
 
 
 @dataclass(frozen=True)

{cognite_toolkit-0.6.105 → cognite_toolkit-0.6.106}/cognite_toolkit/_cdf_tk/utils/fileio/_readers.py

@@ -1,5 +1,6 @@
 import csv
 import json
+import re
 from abc import ABC, abstractmethod
 from collections import Counter, defaultdict
 from collections.abc import Callable, Iterator, Mapping, Sequence
@@ -40,7 +41,7 @@ class FileReader(FileIO, ABC):
         raise NotImplementedError("This method should be implemented in subclasses.")
 
     @classmethod
-    def from_filepath(cls, filepath: Path) -> "FileReader":
+    def from_filepath(cls, filepath: Path) -> "type[FileReader]":
         if len(filepath.suffixes) == 0:
             raise ToolkitValueError(
                 f"File has no suffix. Available formats: {humanize_collection(FILE_READ_CLS_BY_FORMAT.keys())}."
@@ -55,15 +56,57 @@ class FileReader(FileIO, ABC):
                 )
 
         if suffix in FILE_READ_CLS_BY_FORMAT:
-            return FILE_READ_CLS_BY_FORMAT[suffix](input_file=filepath)
+            return FILE_READ_CLS_BY_FORMAT[suffix]
 
         raise ToolkitValueError(
             f"Unknown file format: {suffix}. Available formats: {humanize_collection(FILE_READ_CLS_BY_FORMAT.keys())}."
         )
 
 
+class MultiFileReader(FileReader):
+    """Reads multiple files and yields chunks from each file sequentially.
+
+    Args:
+        input_files (Sequence[Path]): The list of file paths to read.
+    """
+
+    PART_PATTERN = re.compile(r"part-(\d{4})$")
+
+    def __init__(self, input_files: Sequence[Path]) -> None:
+        super().__init__(input_file=input_files[0])
+        self.input_files = input_files
+        reader_classes = Counter([FileReader.from_filepath(input_file) for input_file in self.input_files])
+        if len(reader_classes) > 1:
+            raise ToolkitValueError(
+                "All input files must be of the same format. "
+                f"Found formats: {humanize_collection([cls.FORMAT for cls in reader_classes.keys()])}."
+            )
+        self.reader_class = reader_classes.most_common(1)[0][0]
+
+    @property
+    def is_table(self) -> bool:
+        return issubclass(self.reader_class, TableReader)
+
+    @property
+    def format(self) -> str:
+        return self.reader_class.FORMAT
+
+    def read_chunks(self) -> Iterator[dict[str, JsonVal]]:
+        for input_file in sorted(self.input_files, key=self._part_no):
+            yield from self.reader_class(input_file).read_chunks()
+
+    def _part_no(self, path: Path) -> int:
+        match = self.PART_PATTERN.search(path.stem)
+        if match:
+            return int(match.group(1))
+        return 99999
+
+    def _read_chunks_from_file(self, file: TextIOWrapper) -> Iterator[dict[str, JsonVal]]:
+        raise NotImplementedError("This method is not used in MultiFileReader.")
+
+
 class NDJsonReader(FileReader):
-    format = ".ndjson"
+    FORMAT = ".ndjson"
 
     def _read_chunks_from_file(self, file: TextIOWrapper) -> Iterator[dict[str, JsonVal]]:
         for line in file:
@@ -77,11 +120,11 @@ class YAMLBaseReader(FileReader, ABC):
 
 
 class YAMLReader(YAMLBaseReader):
-    format = ".yaml"
+    FORMAT = ".yaml"
 
 
 class YMLReader(YAMLBaseReader):
-    format = ".yml"
+    FORMAT = ".yml"
 
 
 @dataclass
@@ -171,8 +214,8 @@ class TableReader(FileReader, ABC):
 
         if not input_file.exists():
             raise ToolkitFileNotFoundError(f"File not found: {input_file.as_posix()!r}.")
-        if input_file.suffix != cls.format:
-            raise ToolkitValueError(f"Expected a {cls.format} file got a {input_file.suffix!r} file instead.")
+        if input_file.suffix != cls.FORMAT:
+            raise ToolkitValueError(f"Expected a {cls.FORMAT} file got a {input_file.suffix!r} file instead.")
 
         column_names, sample_rows = cls._read_sample_rows(input_file, sniff_rows)
         cls._check_column_names(column_names)
@@ -213,7 +256,7 @@ class TableReader(FileReader, ABC):
 class CSVReader(TableReader):
     """Reads CSV files and yields each row as a dictionary."""
 
-    format = ".csv"
+    FORMAT = ".csv"
 
     def _read_chunks_from_file(self, file: TextIOWrapper) -> Iterator[dict[str, JsonVal]]:
         if self.keep_failed_cells and self.failed_cell:
@@ -257,7 +300,7 @@ class CSVReader(TableReader):
 
 
 class ParquetReader(TableReader):
-    format = ".parquet"
+    FORMAT = ".parquet"
 
     def __init__(self, input_file: Path) -> None:
         # Parquet files have their own schema, so we don't need to sniff or provide one.
@@ -312,19 +355,19 @@ class ParquetReader(TableReader):
 FILE_READ_CLS_BY_FORMAT: Mapping[str, type[FileReader]] = {}
 TABLE_READ_CLS_BY_FORMAT: Mapping[str, type[TableReader]] = {}
 for subclass in get_concrete_subclasses(FileReader):  # type: ignore[type-abstract]
-    if not getattr(subclass, "format", None):
+    if not getattr(subclass, "FORMAT", None):
         continue
-    if subclass.format in FILE_READ_CLS_BY_FORMAT:
+    if subclass.FORMAT in FILE_READ_CLS_BY_FORMAT:
         raise TypeError(
-            f"Duplicate file format {subclass.format!r} found for classes "
-            f"{FILE_READ_CLS_BY_FORMAT[subclass.format].__name__!r} and {subclass.__name__!r}."
+            f"Duplicate file format {subclass.FORMAT!r} found for classes "
+            f"{FILE_READ_CLS_BY_FORMAT[subclass.FORMAT].__name__!r} and {subclass.__name__!r}."
         )
     # We know we have a dict, but we want to expose FILE_READ_CLS_BY_FORMAT as a Mapping
-    FILE_READ_CLS_BY_FORMAT[subclass.format] = subclass  # type: ignore[index]
+    FILE_READ_CLS_BY_FORMAT[subclass.FORMAT] = subclass  # type: ignore[index]
     if issubclass(subclass, TableReader):
-        if subclass.format in TABLE_READ_CLS_BY_FORMAT:
+        if subclass.FORMAT in TABLE_READ_CLS_BY_FORMAT:
             raise TypeError(
-                f"Duplicate table file format {subclass.format!r} found for classes "
-                f"{TABLE_READ_CLS_BY_FORMAT[subclass.format].__name__!r} and {subclass.__name__!r}."
+                f"Duplicate table file format {subclass.FORMAT!r} found for classes "
+                f"{TABLE_READ_CLS_BY_FORMAT[subclass.FORMAT].__name__!r} and {subclass.__name__!r}."
             )
-        TABLE_READ_CLS_BY_FORMAT[subclass.format] = subclass  # type: ignore[index]
+        TABLE_READ_CLS_BY_FORMAT[subclass.FORMAT] = subclass  # type: ignore[index]