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

cognite_toolkit/_cdf_tk/tracker.py

@@ -14,7 +14,7 @@ from mixpanel import Consumer, Mixpanel, MixpanelException
 
 from cognite_toolkit._cdf_tk.cdf_toml import CDFToml
 from cognite_toolkit._cdf_tk.constants import IN_BROWSER
-from cognite_toolkit._cdf_tk.data_classes._built_modules import BuiltModule
+from cognite_toolkit._cdf_tk.data_classes import CommandTrackingInfo
 from cognite_toolkit._cdf_tk.tk_warnings import ToolkitWarning, WarningList
 from cognite_toolkit._cdf_tk.utils import get_cicd_environment
 from cognite_toolkit._version import __version__
@@ -38,13 +38,19 @@ class Tracker:
 
     @property
     def opted_out(self) -> bool:
-        return self._opt_status == "opted-out"
+        return False
 
     @property
     def opted_in(self) -> bool:
-        return self._opt_status == "opted-in"
+        return True
 
-    def track_cli_command(self, warning_list: WarningList[ToolkitWarning], result: str | Exception, cmd: str) -> bool:
+    def track_cli_command(
+        self,
+        warning_list: WarningList[ToolkitWarning],
+        result: str | Exception,
+        cmd: str,
+        additional_tracking_info: CommandTrackingInfo | None = None,
+    ) -> bool:
         warning_count = Counter([type(w).__name__ for w in warning_list])
 
         warning_details: dict[str, str | int] = {}
@@ -52,7 +58,7 @@ class Tracker:
             warning_details[f"warningMostCommon{no}Count"] = count
             warning_details[f"warningMostCommon{no}Name"] = warning
 
-        positional_args, optional_args = self._parse_sys_args()
+        subcommands, optional_args = self._parse_sys_args()
         event_information = {
             "userInput": self.user_command,
             "toolkitVersion": __version__,
@@ -63,23 +69,16 @@ class Tracker:
             **warning_details,
             "result": type(result).__name__ if isinstance(result, Exception) else result,
             "error": str(result) if isinstance(result, Exception) else "",
-            **positional_args,
+            "subcommands": subcommands,
             **optional_args,
-            **{f"alphaFlag-{name}": value for name, value in self._cdf_toml.alpha_flags.items()},
-            **{f"plugin-{name}": value for name, value in self._cdf_toml.plugins.items()},
+            "alphaFlags": [name for name, value in self._cdf_toml.alpha_flags.items() if value],
+            "plugins": [name for name, value in self._cdf_toml.plugins.items() if value],
         }
 
-        return self._track(f"command{cmd.capitalize()}", event_information)
+        if additional_tracking_info:
+            event_information.update(additional_tracking_info.to_dict())
 
-    def track_module_build(self, module: BuiltModule) -> bool:
-        event_information = {
-            "module": module.name,
-            "location_path": module.location.path.as_posix(),
-            "warning_count": module.warning_count,
-            "status": module.status,
-            **{resource_type: len(resource_build) for resource_type, resource_build in module.resources.items()},
-        }
-        return self._track("moduleBuild", event_information)
+        return self._track(f"command{cmd.capitalize()}", event_information)
 
     def _track(self, event_name: str, event_information: dict[str, Any]) -> bool:
         if self.skip_tracking or not self.opted_in or "PYTEST_CURRENT_TEST" in os.environ:
@@ -129,9 +128,9 @@ class Tracker:
         return distinct_id
 
     @staticmethod
-    def _parse_sys_args() -> tuple[dict[str, str], dict[str, str | bool]]:
+    def _parse_sys_args() -> tuple[list[str], dict[str, str | bool]]:
         optional_args: dict[str, str | bool] = {}
-        positional_args: dict[str, str] = {}
+        subcommands: list[str] = []
         last_key: str | None = None
         if sys.argv and len(sys.argv) > 1:
             for arg in sys.argv[1:]:
@@ -148,11 +147,11 @@ class Tracker:
                     optional_args[last_key] = arg
                     last_key = None
                 else:
-                    positional_args[f"positionalArg{len(positional_args)}"] = arg
+                    subcommands.append(arg)
 
             if last_key:
                 optional_args[last_key] = True
-        return positional_args, optional_args
+        return subcommands, optional_args
 
     @property
     def _cicd(self) -> str:

cognite_toolkit/_cdf_tk/utils/__init__.py

@@ -9,9 +9,9 @@ from .file import (
     read_yaml_file,
     safe_read,
     safe_write,
+    sanitize_filename,
     stringify_value_by_key_in_yaml,
     tmp_build_directory,
-    to_directory_compatible,
 )
 from .graphql_parser import GraphQLParser
 from .hashing import (
@@ -45,9 +45,9 @@ __all__ = [
     "safe_read",
     "safe_write",
     "safe_write",
+    "sanitize_filename",
     "sentry_exception_filter",
     "stringify_value_by_key_in_yaml",
     "tmp_build_directory",
     "to_diff",
-    "to_directory_compatible",
 ]

cognite_toolkit/_cdf_tk/utils/aggregators.py

@@ -18,6 +18,7 @@ from cognite.client.data_classes.sequences import SequenceProperty
 from cognite.client.data_classes.time_series import TimeSeriesProperty
 
 from cognite_toolkit._cdf_tk.client import ToolkitClient
+from cognite_toolkit._cdf_tk.exceptions import ToolkitMissingResourceError
 from cognite_toolkit._cdf_tk.utils.cdf import (
     label_aggregate_count,
     label_count,
@@ -105,6 +106,15 @@ class AssetCentricAggregator(ABC):
                 seen.add(int_id)
         return ids
 
+    def _to_dataset_id(self, data_set_external_id: str | list[str] | None) -> list[int] | None:
+        """Converts data set external IDs to data set IDs."""
+        dataset_id: list[int] | None = None
+        if data_set_external_id is not None:
+            if isinstance(data_set_external_id, str):
+                data_set_external_id = [data_set_external_id]
+            dataset_id = self.client.lookup.data_sets.id(data_set_external_id, allow_empty=False)
+        return dataset_id
+
 
 class MetadataAggregator(AssetCentricAggregator, ABC, Generic[T_CogniteFilter]):
     filter_cls: type[T_CogniteFilter]
@@ -125,7 +135,9 @@ class MetadataAggregator(AssetCentricAggregator, ABC, Generic[T_CogniteFilter]):
         self, hierarchy: str | list[str] | None = None, data_sets: str | list[str] | None = None
     ) -> list[tuple[str, int]]:
         """Returns a list of metadata keys and their counts."""
-        hierarchy_ids, data_set_ids = self._lookup_hierarchy_data_set_pair(hierarchy, data_sets)
+        hierarchy_ids, data_set_ids = self._lookup_hierarchy_data_set_pair(
+            hierarchy, data_sets, operation="find metadata keys"
+        )
         return self._used_metadata_keys(hierarchy=hierarchy_ids, data_sets=data_set_ids)
 
     @lru_cache(maxsize=1)
@@ -140,22 +152,44 @@ class MetadataAggregator(AssetCentricAggregator, ABC, Generic[T_CogniteFilter]):
         )
 
     def _lookup_hierarchy_data_set_pair(
-        self, hierarchy: str | list[str] | None = None, data_sets: str | list[str] | None = None
+        self, hierarchy: str | list[str] | None, data_sets: str | list[str] | None, operation: str
     ) -> tuple[tuple[int, ...] | None, tuple[int, ...] | None]:
         """Returns a tuple of hierarchy and data sets."""
         hierarchy_ids: tuple[int, ...] | None = None
         if isinstance(hierarchy, str):
             asset_id = self.client.lookup.assets.id(external_id=hierarchy, allow_empty=False)
+            if asset_id is None:
+                raise ToolkitMissingResourceError(f"Cannot {operation}. Asset with external ID {hierarchy!r} not found")
             hierarchy_ids = (asset_id,)
         elif isinstance(hierarchy, list) and all(isinstance(item, str) for item in hierarchy):
-            hierarchy_ids = tuple(sorted(self.client.lookup.assets.id(external_id=hierarchy, allow_empty=False)))
+            asset_ids = self.client.lookup.assets.id(external_id=hierarchy, allow_empty=False)
+            if len(asset_ids) != len(hierarchy):
+                missing = set(hierarchy) - set(
+                    self.client.lookup.assets.external_id([id_ for id_ in asset_ids if id_ is not None])
+                )
+                raise ToolkitMissingResourceError(
+                    f"Cannot {operation}. Assets with external IDs {sorted(missing)!r} not found"
+                )
+            hierarchy_ids = tuple(sorted(asset_ids))
 
         data_set_ids: tuple[int, ...] | None = None
         if isinstance(data_sets, str):
             data_set_id = self.client.lookup.data_sets.id(external_id=data_sets, allow_empty=False)
+            if data_set_id is None:
+                raise ToolkitMissingResourceError(
+                    f"Cannot {operation}. Data set with external ID {data_sets!r} not found"
+                )
             data_set_ids = (data_set_id,)
         elif isinstance(data_sets, list) and all(isinstance(item, str) for item in data_sets):
-            data_set_ids = tuple(sorted(self.client.lookup.data_sets.id(external_id=data_sets, allow_empty=False)))
+            data_set_ids_list = self.client.lookup.data_sets.id(external_id=data_sets, allow_empty=False)
+            if len(data_set_ids_list) != len(data_sets):
+                missing = set(data_sets) - set(
+                    self.client.lookup.data_sets.external_id([id_ for id_ in data_set_ids_list if id_ is not None])
+                )
+                raise ToolkitMissingResourceError(
+                    f"Cannot {operation}. Data sets with external IDs {sorted(missing)!r} not found"
+                )
+            data_set_ids = tuple(sorted(data_set_ids_list))
 
         return hierarchy_ids, data_set_ids
 
@@ -201,7 +235,9 @@ class LabelAggregator(MetadataAggregator, ABC, Generic[T_CogniteFilter]):
         self, hierarchy: str | list[str] | None = None, data_sets: str | list[str] | None = None
     ) -> list[tuple[str, int]]:
         """Returns a list of labels and their counts."""
-        hierarchy_ids, data_set_ids = self._lookup_hierarchy_data_set_pair(hierarchy, data_sets)
+        hierarchy_ids, data_set_ids = self._lookup_hierarchy_data_set_pair(
+            hierarchy, data_sets, operation="find labels"
+        )
         return self._used_labels(hierarchy=hierarchy_ids, data_sets=data_set_ids)
 
     @lru_cache(maxsize=1)
@@ -359,9 +395,10 @@ class RelationshipAggregator(AssetCentricAggregator):
     def count(
         self, hierarchy: str | list[str] | None = None, data_set_external_id: str | list[str] | None = None
     ) -> int:
-        if hierarchy is not None or data_set_external_id is not None:
+        if hierarchy is not None:
             raise NotImplementedError()
-        results = relationship_aggregate_count(self.client)
+        dataset_id = self._to_dataset_id(data_set_external_id)
+        results = relationship_aggregate_count(self.client, dataset_id)
         return sum(result.count for result in results)
 
     def used_data_sets(self, hierarchy: str | None = None) -> list[str]:
@@ -378,9 +415,10 @@ class LabelCountAggregator(AssetCentricAggregator):
     def count(
         self, hierarchy: str | list[str] | None = None, data_set_external_id: str | list[str] | None = None
     ) -> int:
-        if hierarchy is not None or data_set_external_id is not None:
+        if hierarchy is not None:
             raise NotImplementedError()
-        return label_aggregate_count(self.client)
+        data_set_id = self._to_dataset_id(data_set_external_id)
+        return label_aggregate_count(self.client, data_set_id)
 
     def used_data_sets(self, hierarchy: str | None = None) -> list[str]:
         raise NotImplementedError()

cognite_toolkit/_cdf_tk/utils/cdf.py

@@ -36,7 +36,7 @@ from cognite_toolkit._cdf_tk.tk_warnings import (
 )
 from cognite_toolkit._cdf_tk.utils import humanize_collection
 
-from .file import to_directory_compatible
+from .file import sanitize_filename
 from .sql_parser import SQLParser
 
 if sys.version_info < (3, 11):
@@ -103,10 +103,14 @@ def iterate_instances(
     # By making the sort be on external_id, Postgres should pick the index that's on (project_id, space, external_id)
     # WHERE deleted_at IS NULL. In other words, avoiding soft deleted instances.
     body["sort"] = [
+        {
+            "property": [instance_type, "space"],
+            "direction": "ascending",
+        },
         {
             "property": [instance_type, "externalId"],
             "direction": "ascending",
-        }
+        },
     ]
     url = f"/api/{client._API_VERSION}/projects/{client.config.project}/models/instances/list"
     if space:
@@ -137,7 +141,7 @@ def iterate_instances(
 
 @overload
 def read_auth(
-    authentication: object,
+    authentication: Any,
     client_config: ToolkitClientConfig,
     identifier: Hashable,
     resource_name: str,
@@ -148,7 +152,7 @@ def read_auth(
 
 @overload
 def read_auth(
-    authentication: object,
+    authentication: Any,
     client_config: ToolkitClientConfig,
     identifier: Hashable,
     resource_name: str,
@@ -158,7 +162,7 @@ def read_auth(
 
 
 def read_auth(
-    authentication: object,
+    authentication: Any,
     client_config: ToolkitClientConfig,
     identifier: Hashable,
     resource_name: str,
@@ -406,7 +410,7 @@ class ThrottlerState:
     @classmethod
     def _filepath(cls, project: str) -> Path:
         # Enure the project name is a valid string for hashing.
-        project = to_directory_compatible(project)
+        project = sanitize_filename(project)
         filename = cls._FILENAME.format(project=project)
         return Path(tempfile.gettempdir()) / filename
 

cognite_toolkit/_cdf_tk/utils/cli_args.py

@@ -0,0 +1,34 @@
+from cognite.client.data_classes.data_modeling import ViewId
+
+from cognite_toolkit._cdf_tk.exceptions import ToolkitValueError
+
+
+def parse_view_str(view_str: str) -> ViewId:
+    """Parse a view string into a ViewId.
+
+    Args:
+        view_str (str): The view string to parse.
+
+    Returns:
+        ViewId: The parsed ViewId.
+
+    Raises:
+        ToolkitValueError: If the view string is not in a valid format.
+
+    >>> parse_view_str("my_space:my_view/v1")
+    ViewId(space='my_space', external_id='my_view', version='v1')
+    """
+    space_and_rest = view_str.split(":", 1)
+    if len(space_and_rest) != 2 or not space_and_rest[0] or not space_and_rest[1]:
+        raise ToolkitValueError(
+            f"Invalid view string format: '{view_str}'. Expected format 'space:externalId/version'."
+        )
+    space = space_and_rest[0]
+    external_id_and_version = space_and_rest[1].rsplit("/", 1)
+    if len(external_id_and_version) != 2 or not external_id_and_version[0] or not external_id_and_version[1]:
+        raise ToolkitValueError(
+            f"Invalid view string format: '{view_str}'. Expected format 'space:externalId/version'."
+        )
+    external_id = external_id_and_version[0]
+    version = external_id_and_version[1]
+    return ViewId(space=space, external_id=external_id, version=version)

cognite_toolkit/_cdf_tk/utils/dtype_conversion.py

@@ -7,7 +7,12 @@ from typing import ClassVar, Literal, overload
 
 from cognite.client.data_classes import Label, LabelDefinition
 from cognite.client.data_classes.data_modeling import ContainerId
-from cognite.client.data_classes.data_modeling.data_types import Enum, ListablePropertyType, PropertyType
+from cognite.client.data_classes.data_modeling.data_types import (
+    DirectRelationReference,
+    Enum,
+    ListablePropertyType,
+    PropertyType,
+)
 from cognite.client.data_classes.data_modeling.instances import PropertyValueWrite
 from cognite.client.utils import ms_to_datetime
 from dateutil import parser
@@ -15,7 +20,13 @@ from dateutil import parser
 from cognite_toolkit._cdf_tk.constants import CDF_UNIT_SPACE
 from cognite_toolkit._cdf_tk.exceptions import ToolkitNotSupported
 from cognite_toolkit._cdf_tk.utils._auxiliary import get_concrete_subclasses
-from cognite_toolkit._cdf_tk.utils.useful_types import AssetCentric, DataType, JsonVal, PythonTypes
+from cognite_toolkit._cdf_tk.utils.useful_types import (
+    AssetCentricType,
+    AssetCentricTypeExtended,
+    DataType,
+    JsonVal,
+    PythonTypes,
+)
 
 from .collection import humanize_collection
 
@@ -30,18 +41,27 @@ def asset_centric_convert_to_primary_property(
     type_: PropertyType,
     nullable: bool,
     destination_container_property: tuple[ContainerId, str],
-    source_property: tuple[AssetCentric, str],
+    source_property: tuple[AssetCentricTypeExtended, str],
+    direct_relation_lookup: Mapping[str | int, DirectRelationReference] | None = None,
 ) -> PropertyValueWrite:
     if (source_property, destination_container_property) in SPECIAL_CONVERTER_BY_SOURCE_DESTINATION:
         converter_cls = SPECIAL_CONVERTER_BY_SOURCE_DESTINATION[(source_property, destination_container_property)]
-        return converter_cls().convert(value)
+        if issubclass(converter_cls, _SourceConverter):
+            if direct_relation_lookup is None:
+                raise TypeError("Direct relation lookup must be provided for DirectRelationReference conversion")
+            return converter_cls(direct_relation_lookup=direct_relation_lookup).convert(value)
+        else:
+            return converter_cls().convert(value)
     else:
         # Fallback to the standard conversion
-        return convert_to_primary_property(value, type_, nullable)
+        return convert_to_primary_property(value, type_, nullable, direct_relation_lookup=direct_relation_lookup)
 
 
 def convert_to_primary_property(
-    value: str | int | float | bool | dict | list | None, type_: PropertyType, nullable: bool
+    value: str | int | float | bool | dict | list | None,
+    type_: PropertyType,
+    nullable: bool,
+    direct_relation_lookup: Mapping[str | int, DirectRelationReference] | None = None,
 ) -> PropertyValueWrite:
     """Convert a string value to the appropriate type based on the provided property type.
 
@@ -49,6 +69,8 @@ def convert_to_primary_property(
         value (str | int | float | bool): The value to convert.
         type_ (PropertyType): The type of the property to convert to.
         nullable (bool): Whether the property can be null.
+        direct_relation_lookup (Mapping[str | int, DirectRelationReference] | None): A mapping from external/internal
+            IDs to DirectRelationReference objects,
 
     Returns:
         PropertyValueWrite: The converted value as a PropertyValue, or None if is_nullable is True and the value is empty.
@@ -62,6 +84,10 @@ def convert_to_primary_property(
         converter: _ValueConverter = _EnumConverter(type_=type_, nullable=nullable)
     elif issubclass(converter_cls, _EnumConverter):
         raise TypeError(f"Unsupported property type {type_} for enum conversion")
+    elif issubclass(converter_cls, _DirectRelationshipConverter):
+        if direct_relation_lookup is None:
+            raise TypeError("Direct relation lookup must be provided for DirectRelationReference conversion")
+        converter = _DirectRelationshipConverter(direct_relation_lookup=direct_relation_lookup)
     else:
         converter = converter_cls(nullable)
 
@@ -209,7 +235,7 @@ class _ValueConverter(_Converter, ABC):
     schema_type: ClassVar[DataType | None] = None
     _handles_list: ClassVar[bool] = False
 
-    def __init__(self, nullable: bool):
+    def __init__(self, nullable: bool) -> None:
         self.nullable = nullable
 
     def convert(self, value: str | int | float | bool | dict | list | None) -> PropertyValueWrite:
@@ -231,7 +257,7 @@ class _ValueConverter(_Converter, ABC):
 class _SpecialCaseConverter(_Converter, ABC):
     """Abstract base class for converters handling special cases."""
 
-    source_property: ClassVar[tuple[AssetCentric, str]]
+    source_property: ClassVar[tuple[AssetCentricType, str]]
     destination_container_property: ClassVar[tuple[ContainerId, str]]
 
 
@@ -291,6 +317,37 @@ class _FileLabelConverter(_LabelConverter):
     source_property = ("file", "labels")
 
 
+class _SourceConverter(_SpecialCaseConverter, ABC):
+    destination_container_property = (ContainerId("cdf_cdm", "CogniteSourceable"), "source")
+
+    def __init__(self, direct_relation_lookup: Mapping[str | int, DirectRelationReference]) -> None:
+        self.direct_relation_lookup = direct_relation_lookup
+
+    def convert(self, value: str | int | float | bool | dict | list | None) -> DirectRelationReference:
+        if isinstance(value, str | int):
+            if value in self.direct_relation_lookup:
+                return self.direct_relation_lookup[value]
+            elif isinstance(value, str) and value.casefold() in self.direct_relation_lookup:
+                # The aggregate endpoint used to create source system lowercases all sources, thus
+                # we do a case-insensitive lookup as a fallback.
+                return self.direct_relation_lookup[value.casefold()]
+        raise ValueError(
+            f"Cannot convert {value!r} to DirectRelationReference. Invalid data type or missing in lookup."
+        )
+
+
+class _AssetSourceConverter(_SourceConverter):
+    source_property = ("asset", "source")
+
+
+class _FileSourceConverter(_SourceConverter):
+    source_property = ("file", "source")
+
+
+class _EventSourceConverter(_SourceConverter):
+    source_property = ("event", "source")
+
+
 class _TextConverter(_ValueConverter):
     type_str = "text"
     schema_type = "string"
@@ -422,7 +479,7 @@ class _TimestampConverter(_ValueConverter):
     type_str = "timestamp"
     schema_type = "timestamp"
 
-    def _convert(self, value: str | int | float | bool | dict) -> PropertyValueWrite:
+    def _convert(self, value: str | int | float | bool | dict) -> datetime:
         if isinstance(value, int | float):
             try:
                 return ms_to_datetime(value)
@@ -436,6 +493,23 @@ class _TimestampConverter(_ValueConverter):
         raise ValueError(f"Cannot convert {value} to timestamp.")
 
 
+class _EpochConverter(_ValueConverter):
+    type_str = "epoch"
+    schema_type = "epoch"
+
+    def _convert(self, value: str | int | float | bool | dict) -> int:
+        if isinstance(value, int | float):
+            return int(value)
+        elif isinstance(value, str):
+            try:
+                dt = parser.isoparse(value)
+                epoch = int(dt.timestamp() * 1000)
+                return epoch
+            except ValueError as e:
+                raise ValueError(f"Cannot convert {value} to epoch timestamp: {e}") from e
+        raise ValueError(f"Cannot convert {value} to epoch timestamp.")
+
+
 class _DateConverter(_ValueConverter):
     type_str = "date"
     schema_type = "date"
@@ -468,8 +542,16 @@ class _EnumConverter(_ValueConverter):
 class _DirectRelationshipConverter(_ValueConverter):
     type_str = "direct"
 
+    def __init__(self, direct_relation_lookup: Mapping[str | int, DirectRelationReference]) -> None:
+        super().__init__(nullable=True)
+        self.direct_relation_lookup = direct_relation_lookup
+
     def _convert(self, value: str | int | float | bool | dict) -> PropertyValueWrite:
-        raise ToolkitNotSupported("Direct relationship conversion is not supported.")
+        if isinstance(value, str | int) and value in self.direct_relation_lookup:
+            return self.direct_relation_lookup[value]
+        raise ValueError(
+            f"Cannot convert {value!r} to DirectRelationReference. Invalid data type or missing in lookup."
+        )
 
 
 class _TimeSeriesReferenceConverter(_ValueConverter):
@@ -498,7 +580,7 @@ CONVERTER_BY_DTYPE: Mapping[str, type[_ValueConverter]] = {
     for cls_ in _ValueConverter.__subclasses__()
 }
 SPECIAL_CONVERTER_BY_SOURCE_DESTINATION: Mapping[
-    tuple[tuple[AssetCentric, str], tuple[ContainerId, str]],
+    tuple[tuple[AssetCentricTypeExtended, str], tuple[ContainerId, str]],
     type[_SpecialCaseConverter],
 ] = {
     (subclass.source_property, subclass.destination_container_property): subclass

cognite_toolkit/_cdf_tk/utils/file.py

@@ -14,6 +14,7 @@ from contextlib import contextmanager
 from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Any, Literal, TypeVar, overload
+from zipfile import ZipFile
 
 import pandas as pd
 import yaml
@@ -22,8 +23,6 @@ from rich import print
 from cognite_toolkit._cdf_tk.cdf_toml import CDFToml
 from cognite_toolkit._cdf_tk.constants import ENV_VAR_PATTERN, HINT_LEAD_TEXT, URL
 from cognite_toolkit._cdf_tk.exceptions import (
-    ToolkitFileNotFoundError,
-    ToolkitNotADirectoryError,
     ToolkitValueError,
     ToolkitYAMLFormatError,
 )
@@ -146,6 +145,28 @@ _TOTAL_ELAPSED_TIME = 0.0
 _HAS_HINTED = False
 
 
+def keyvault_constructor(loader: yaml.Loader, node: yaml.Node) -> Any:
+    """Constructor to preserve the !keyvault tag along with the value."""
+    if isinstance(node, yaml.ScalarNode):
+        value = loader.construct_scalar(node)
+        return f"!keyvault {value}"
+    elif isinstance(node, yaml.MappingNode):
+        return loader.construct_mapping(node)
+    elif isinstance(node, yaml.SequenceNode):
+        return loader.construct_sequence(node)
+    else:
+        # This should never happen as YAML only has these three node types
+        # But we need this for type checking completeness
+        raise ValueError(f"Unexpected node type: {type(node)}")
+
+
+yaml.SafeLoader.add_constructor("!keyvault", keyvault_constructor)
+
+# Register for the fast CSafeLoader (used directly when libyaml is available)
+if hasattr(yaml, "CSafeLoader"):
+    yaml.CSafeLoader.add_constructor("!keyvault", keyvault_constructor)
+
+
 def read_yaml_content(content: str) -> dict[str, Any] | list[dict[str, Any]]:
     """Read a YAML string and return a dictionary
 
@@ -173,8 +194,8 @@ def read_yaml_content(content: str) -> dict[str, Any] | list[dict[str, Any]]:
 _ILLEGAL_CHARACTERS = re.compile(r"[<>:\"/\\|?*\s]")
 
 
-def to_directory_compatible(text: str) -> str:
-    """Convert a string to be compatible with directory names on all platforms"""
+def sanitize_filename(text: str) -> str:
+    """Convert a string to be a valid filename by replacing illegal characters with underscores."""
     cleaned = _ILLEGAL_CHARACTERS.sub("_", text)
     # Replace multiple underscores with a single one
     return re.sub(r"_+", "_", cleaned)
@@ -429,38 +450,31 @@ def get_table_columns(table: Path) -> list[str]:
         raise ToolkitValueError(f"The file {table.name} is not a supported table format (csv, parquet)")
 
 
-def find_adjacent_files(filepath: Path, suffix: str) -> list[Path]:
-    """Find files in the same directory as the given file that have the same
-    prefix and a specific suffix.
-
-    Args:
-        filepath (Path): The path to the file for which to find adjacent files.
-        suffix (str): The suffix to match for the adjacent files.
-
-    Returns:
-        list[Path]: A list of Paths to the adjacent files that match the criteria.
+@contextmanager
+def create_temporary_zip(directory: Path, zipname: str) -> typing.Generator[Path, None, None]:
     """
-    if not filepath.is_file():
-        raise ToolkitFileNotFoundError(f"The provided path {filepath} is not a file.")
-    return find_files_with_suffix_and_prefix(filepath.parent, filepath.name, suffix=suffix)
-
-
-def find_files_with_suffix_and_prefix(dirpath: Path, name: str, suffix: str) -> list[Path]:
-    """Find files in a directory that have the same prefix as the given name
+    Create a temporary zip file from a directory.
 
     Args:
-        dirpath: The directory in which to search for files.
-        name: The name of the file to match against, which should include the prefix.
-        suffix: The suffix to match for the files.
+        directory: The directory to zip.
+        zipname: The name of the zip file.
 
     Returns:
-        list[Path]: A list of Paths to the files that match the criteria.
+        A context manager that yields the path to the zip file.
+
     """
-    if not dirpath.is_dir():
-        raise ToolkitNotADirectoryError(f"The provided path {dirpath} is not a directory.")
-    found_files: list[Path] = []
-    for file in dirpath.glob(f"*{suffix}"):
-        filestem = file.name.removesuffix(suffix)
-        if file.is_file() and name.startswith(filestem):
-            found_files.append(file)
-    return found_files
+    with tempfile.TemporaryDirectory() as tmpdir:
+        zip_path = Path(tmpdir, zipname)
+        with ZipFile(zip_path, "w", strict_timestamps=False) as zf:
+            for root, _, files in os.walk(directory):
+                root_path = Path(root)
+                arc_root = root_path.relative_to(directory)
+
+                # Add directory entry to preserve directory structure, including empty ones.
+                zf.write(root_path, arcname=str(arc_root))
+
+                for filename in files:
+                    file_path = root_path / filename
+                    arcname = file_path.relative_to(directory)
+                    zf.write(file_path, arcname)
+        yield zip_path

cognite_toolkit/_cdf_tk/utils/fileio/__init__.py

@@ -8,9 +8,11 @@ from ._compression import (
 )
 from ._readers import (
     FILE_READ_CLS_BY_FORMAT,
+    TABLE_READ_CLS_BY_FORMAT,
     CSVReader,
     FailedParsing,
     FileReader,
+    MultiFileReader,
     NDJsonReader,
     ParquetReader,
     YAMLReader,
@@ -32,6 +34,7 @@ __all__ = [
     "COMPRESSION_BY_SUFFIX",
     "FILE_READ_CLS_BY_FORMAT",
     "FILE_WRITE_CLS_BY_FORMAT",
+    "TABLE_READ_CLS_BY_FORMAT",
     "TABLE_WRITE_CLS_BY_FORMAT",
     "CSVReader",
     "CSVWriter",
@@ -43,6 +46,7 @@ __all__ = [
     "FileReader",
     "FileWriter",
     "GzipCompression",
+    "MultiFileReader",
     "NDJsonReader",
     "NDJsonWriter",
     "ParquetReader",