cognite-toolkit 0.7.43__py3-none-any.whl → 0.7.44__py3-none-any.whl

cognite_toolkit/_cdf_tk/client/data_classes/workflow_trigger.py

@@ -0,0 +1,63 @@
+from typing import Annotated, Literal
+
+from pydantic import Field, JsonValue
+
+from cognite_toolkit._cdf_tk.client.data_classes.base import (
+    BaseModelObject,
+    RequestResource,
+    ResponseResource,
+)
+
+from .identifiers import ExternalId
+
+
+class TriggerRuleDefinition(BaseModelObject):
+    trigger_type: str
+
+
+class ScheduleTriggerRule(BaseModelObject):
+    trigger_type: Literal["schedule"] = "schedule"
+    cron_expression: str
+    timezone: str | None = None
+
+
+class DataModelingTriggerRule(BaseModelObject):
+    trigger_type: Literal["dataModeling"] = "dataModeling"
+    data_modeling_query: JsonValue
+    batch_size: int
+    batch_timeout: int
+
+
+TriggerRule = Annotated[
+    ScheduleTriggerRule | DataModelingTriggerRule,
+    Field(discriminator="trigger_type"),
+]
+
+
+class NonceCredentials(BaseModelObject):
+    nonce: str
+
+
+class WorkflowTrigger(BaseModelObject):
+    external_id: str
+    trigger_rule: TriggerRule
+    workflow_external_id: str
+    workflow_version: str
+    input: JsonValue | None = None
+    metadata: dict[str, str] | None = None
+
+    def as_id(self) -> ExternalId:
+        return ExternalId(external_id=self.external_id)
+
+
+class WorkflowTriggerRequest(WorkflowTrigger, RequestResource):
+    authentication: NonceCredentials
+
+
+class WorkflowTriggerResponse(WorkflowTrigger, ResponseResource[WorkflowTriggerRequest]):
+    created_time: int
+    last_updated_time: int
+    is_paused: bool
+
+    def as_request_resource(self) -> WorkflowTriggerRequest:
+        return WorkflowTriggerRequest.model_validate(self.dump(), extra="ignore")

cognite_toolkit/_cdf_tk/client/data_classes/workflow_version.py

@@ -0,0 +1,155 @@
+from typing import Annotated, Any, Literal
+
+from pydantic import Field, JsonValue, field_validator
+from pydantic_core.core_schema import ValidationInfo
+
+from cognite_toolkit._cdf_tk.client.data_classes.base import (
+    BaseModelObject,
+    RequestResource,
+    ResponseResource,
+)
+
+from .identifiers import Identifier, WorkflowVersionId
+
+
+class TaskId(Identifier):
+    external_id: str
+
+    def __str__(self) -> str:
+        return f"externalId='{self.external_id}'"
+
+
+class TaskParameterDefinition(BaseModelObject):
+    type: str
+
+
+class CogniteFunctionRef(BaseModelObject):
+    external_id: str
+    data: JsonValue | str | None = None
+
+
+class FunctionTaskParameters(TaskParameterDefinition):
+    type: Literal["function"] = Field("function", exclude=True)
+    function: CogniteFunctionRef
+    is_async_complete: bool | None = None
+
+
+class TransformationRef(BaseModelObject):
+    external_id: str
+    concurrency_policy: Literal["fail", "waitForCurrent", "restartAfterCurrent"] | None = None
+    use_transformation_credentials: bool | None = None
+
+
+class TransformationTaskParameters(TaskParameterDefinition):
+    type: Literal["transformation"] = Field("transformation", exclude=True)
+    transformation: TransformationRef
+
+
+class CDFRequest(BaseModelObject):
+    resource_path: str
+    method: str
+    query_parameters: JsonValue | str | None = None
+    body: JsonValue | str | None = None
+    request_timeout_in_millis: float | str | None = None
+    cdf_version_header: str | None = None
+
+
+class CDFTaskParameters(TaskParameterDefinition):
+    type: Literal["cdf"] = Field("cdf", exclude=True)
+    cdf_request: CDFRequest
+
+
+class DynamicRef(BaseModelObject):
+    tasks: str
+
+
+class DynamicTaskParameters(TaskParameterDefinition):
+    type: Literal["dynamic"] = Field("dynamic", exclude=True)
+    dynamic: DynamicRef
+
+
+class SubworkflowRef(BaseModelObject):
+    tasks: "list[Task]"
+
+
+class SubworkflowTaskParameters(TaskParameterDefinition):
+    type: Literal["subworkflow"] = Field("subworkflow", exclude=True)
+    subworkflow: WorkflowVersionId | SubworkflowRef
+
+
+class SimulatorInputUnit(BaseModelObject):
+    name: str
+
+
+class SimulatorInput(BaseModelObject):
+    reference_id: str
+    value: str | int | float | list[str] | list[int] | list[float]
+    unit: SimulatorInputUnit | None = None
+
+
+class SimulationRef(BaseModelObject):
+    routine_external_id: str
+    run_time: int | None = None
+    inputs: list[SimulatorInput] | None = None
+
+
+class SimulationTaskParameters(TaskParameterDefinition):
+    type: Literal["simulation"] = Field("simulation", exclude=True)
+    simulation: SimulationRef
+
+
+Parameter = Annotated[
+    FunctionTaskParameters
+    | TransformationTaskParameters
+    | CDFTaskParameters
+    | DynamicTaskParameters
+    | SubworkflowTaskParameters
+    | SimulationTaskParameters,
+    Field(discriminator="type"),
+]
+
+
+class Task(BaseModelObject):
+    external_id: str
+    type: str
+    name: str | None = None
+    description: str | None = None
+    retries: int | None = None
+    timeout: int | None = None
+    on_failure: Literal["abortWorkflow", "skipTask"] = "abortWorkflow"
+    depends_on: list[TaskId] | None = None
+    parameters: Parameter | None = None
+
+    @field_validator("parameters", mode="before")
+    @classmethod
+    def move_type_to_field(cls, value: Any, info: ValidationInfo) -> Any:
+        if not isinstance(value, dict) or "type" not in info.data:
+            return value
+        value = dict(value)
+        value["type"] = info.data["type"]
+        return value
+
+
+class WorkflowDefinition(BaseModelObject):
+    description: str | None = None
+    tasks: list[Task]
+
+
+class WorkflowVersion(BaseModelObject):
+    workflow_external_id: str
+    version: str
+    workflow_definition: WorkflowDefinition
+
+    def as_id(self) -> WorkflowVersionId:
+        return WorkflowVersionId(workflow_external_id=self.workflow_external_id, version=self.version)
+
+
+class WorkflowVersionRequest(WorkflowVersion, RequestResource): ...
+
+
+class WorkflowVersionResponse(WorkflowVersion, ResponseResource[WorkflowVersionRequest]):
+    created_time: int
+    last_updated_time: int
+
+    def as_request_resource(self) -> WorkflowVersionRequest:
+        return WorkflowVersionRequest.model_validate(self.dump(), extra="ignore")

cognite_toolkit/_cdf_tk/client/testing.py

@@ -7,12 +7,14 @@ from cognite.client._api.datapoints import DatapointsAPI
 from cognite.client._api.datapoints_subscriptions import DatapointsSubscriptionAPI
 from cognite.client._api.functions import FunctionCallsAPI, FunctionSchedulesAPI
 from cognite.client._api.raw import RawDatabasesAPI, RawRowsAPI, RawTablesAPI
+from cognite.client._api.simulators import SimulatorModelsAPI, SimulatorsAPI
 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
 from cognite_toolkit._cdf_tk.client.api.assets import AssetsAPI
+from cognite_toolkit._cdf_tk.client.api.filemetadata import FileMetadataAPI
 from cognite_toolkit._cdf_tk.client.api.legacy.canvas import CanvasAPI, IndustrialCanvasAPI
 from cognite_toolkit._cdf_tk.client.api.legacy.charts import ChartsAPI
 from cognite_toolkit._cdf_tk.client.api.legacy.dml import DMLAPI
@@ -143,7 +145,10 @@ class ToolkitClientMock(CogniteClientMock):
         self.tool.three_d.models = MagicMock(spec_set=ThreeDModelAPI)
         self.tool.assets = MagicMock(spec_set=AssetsAPI)
         self.tool.timeseries = MagicMock(spec_set=TimeSeriesAPI)
+        self.tool.filemetadata = MagicMock(spec_set=FileMetadataAPI)
         self.tool.events = MagicMock(spec_set=EventsAPI)
+        self.tool.simulators = MagicMock(spec=SimulatorsAPI)
+        self.tool.simulators.models = MagicMock(spec_set=SimulatorModelsAPI)
 
         self.streams = MagicMock(spec=StreamsAPI)
 

cognite_toolkit/_cdf_tk/commands/_migrate/conversion.py

@@ -1,7 +1,7 @@
 from collections.abc import Iterable, Mapping, Set
 from typing import Any, ClassVar, cast
 
-from cognite.client.data_classes import Annotation, FileMetadata
+from cognite.client.data_classes import Annotation
 from cognite.client.data_classes.data_modeling import (
     DirectRelation,
     DirectRelationReference,
@@ -18,6 +18,7 @@ from cognite.client.utils._identifier import InstanceId
 from cognite_toolkit._cdf_tk.client import ToolkitClient
 from cognite_toolkit._cdf_tk.client.data_classes.asset import AssetResponse
 from cognite_toolkit._cdf_tk.client.data_classes.event import EventResponse
+from cognite_toolkit._cdf_tk.client.data_classes.filemetadata import FileMetadataResponse
 from cognite_toolkit._cdf_tk.client.data_classes.legacy.migration import (
     AssetCentricId,
     ResourceViewMappingApply,
@@ -117,7 +118,7 @@ class DirectRelationCache:
                     source_ids.add(resource.source)
                 if resource.parent_id is not None:
                     asset_ids.add(resource.parent_id)
-            elif isinstance(resource, FileMetadata):
+            elif isinstance(resource, FileMetadataResponse):
                 if resource.source:
                     source_ids.add(resource.source)
                 if resource.asset_ids:
@@ -248,7 +249,7 @@ def asset_centric_to_dm(
 def _lookup_resource_type(resource_type: AssetCentricResourceExtended) -> AssetCentricTypeExtended:
     if isinstance(resource_type, AssetResponse):
         return "asset"
-    elif isinstance(resource_type, FileMetadata):
+    elif isinstance(resource_type, FileMetadataResponse):
         return "file"
     elif isinstance(resource_type, EventResponse):
         return "event"

cognite_toolkit/_cdf_tk/cruds/_data_cruds.py

@@ -6,6 +6,7 @@ from typing import TYPE_CHECKING, cast, final
 import pandas as pd
 from cognite.client.data_classes import FileMetadataWrite
 
+from cognite_toolkit._cdf_tk.client.data_classes.identifiers import ExternalId
 from cognite_toolkit._cdf_tk.client.data_classes.legacy.extendable_cognite_file import ExtendableCogniteFileApply
 from cognite_toolkit._cdf_tk.client.data_classes.legacy.raw import RawTable
 from cognite_toolkit._cdf_tk.constants import BUILD_FOLDER_ENCODING
@@ -119,11 +120,14 @@ class FileCRUD(DataCRUD):
                 if not datafile.exists():
                     continue
 
+                identifier = resource.identifier
+                if isinstance(identifier, ExternalId):
+                    identifier = identifier.external_id
                 if dry_run:
-                    yield f" Would upload file '{datafile!s}' to file with {id_name}={resource.identifier!r}", 1
+                    yield f" Would upload file '{datafile!s}' to file with {id_name}={identifier!r}", 1
                 else:
-                    self.client.files.upload_content(path=str(datafile), **{id_name: resource.identifier})
-                    yield f" Uploaded file '{datafile!s}' to file with {id_name}={resource.identifier!r}", 1
+                    self.client.files.upload_content(path=str(datafile), **{id_name: identifier})
+                    yield f" Uploaded file '{datafile!s}' to file with {id_name}={identifier!r}", 1
 
     @staticmethod
     def _read_metadata(

cognite_toolkit/_cdf_tk/cruds/_resource_cruds/file.py

@@ -15,13 +15,8 @@
 
 from collections.abc import Hashable, Iterable, Sequence
 from datetime import date, datetime
-from typing import Any, cast, final
+from typing import Any, final
 
-from cognite.client.data_classes import (
-    FileMetadata,
-    FileMetadataList,
-    FileMetadataWrite,
-)
 from cognite.client.data_classes.capabilities import (
     Capability,
     DataModelInstancesAcl,
@@ -31,9 +26,9 @@ from cognite.client.data_classes.data_modeling import NodeApplyResultList, NodeI
 from cognite.client.exceptions import CogniteAPIError
 from cognite.client.utils._time import convert_data_modelling_timestamp
 from cognite.client.utils.useful_types import SequenceNotStr
-from rich import print
 
-from cognite_toolkit._cdf_tk.client.data_classes.identifiers import ExternalId
+from cognite_toolkit._cdf_tk.client.data_classes.filemetadata import FileMetadataRequest, FileMetadataResponse
+from cognite_toolkit._cdf_tk.client.data_classes.identifiers import ExternalId, InternalOrExternalId
 from cognite_toolkit._cdf_tk.client.data_classes.legacy.extendable_cognite_file import (
     ExtendableCogniteFile,
     ExtendableCogniteFileApply,
@@ -56,11 +51,11 @@ from .datamodel import SpaceCRUD, ViewCRUD
 
 
 @final
-class FileMetadataCRUD(ResourceContainerCRUD[str, FileMetadataWrite, FileMetadata]):
+class FileMetadataCRUD(ResourceContainerCRUD[ExternalId, FileMetadataRequest, FileMetadataResponse]):
     item_name = "file contents"
     folder_name = "files"
-    resource_cls = FileMetadata
-    resource_write_cls = FileMetadataWrite
+    resource_cls = FileMetadataResponse
+    resource_write_cls = FileMetadataRequest
     yaml_cls = FileMetadataYAML
     kind = "FileMetadata"
     dependencies = frozenset({DataSetsCRUD, GroupAllScopedCRUD, LabelCRUD, AssetCRUD})
@@ -73,7 +68,7 @@ class FileMetadataCRUD(ResourceContainerCRUD[str, FileMetadataWrite, FileMetadat
 
     @classmethod
     def get_required_capability(
-        cls, items: Sequence[FileMetadataWrite] | None, read_only: bool
+        cls, items: Sequence[FileMetadataRequest] | None, read_only: bool
     ) -> Capability | list[Capability]:
         if not items and items is not None:
             return []
@@ -88,22 +83,22 @@ class FileMetadataCRUD(ResourceContainerCRUD[str, FileMetadataWrite, FileMetadat
         return FilesAcl(actions, scope)
 
     @classmethod
-    def get_id(cls, item: FileMetadata | FileMetadataWrite | dict) -> str:
+    def get_id(cls, item: FileMetadataRequest | FileMetadataResponse | dict) -> ExternalId:
         if isinstance(item, dict):
-            return item["externalId"]
+            return ExternalId(external_id=item["externalId"])
         if item.external_id is None:
             raise ToolkitRequiredValueError("FileMetadata must have external_id set.")
-        return item.external_id
+        return ExternalId(external_id=item.external_id)
 
     @classmethod
-    def get_internal_id(cls, item: FileMetadata | dict) -> int:
+    def get_internal_id(cls, item: FileMetadataResponse | dict) -> int:
         if isinstance(item, dict):
             return item["id"]
         return item.id
 
     @classmethod
-    def dump_id(cls, id: str) -> dict[str, Any]:
-        return {"externalId": id}
+    def dump_id(cls, id: ExternalId) -> dict[str, Any]:
+        return id.dump()
 
     @classmethod
     def get_dependent_items(cls, item: dict) -> Iterable[tuple[type[ResourceCRUD], Hashable]]:
@@ -121,69 +116,71 @@ class FileMetadataCRUD(ResourceContainerCRUD[str, FileMetadataWrite, FileMetadat
         for asset_external_id in item.get("assetExternalIds", []):
             yield AssetCRUD, ExternalId(external_id=asset_external_id)
 
-    def load_resource(self, resource: dict[str, Any], is_dry_run: bool = False) -> FileMetadataWrite:
-        if resource.get("dataSetExternalId") is not None:
-            ds_external_id = resource.pop("dataSetExternalId")
+    def load_resource(self, resource: dict[str, Any], is_dry_run: bool = False) -> FileMetadataRequest:
+        if ds_external_id := resource.pop("dataSetExternalId", None):
             resource["dataSetId"] = self.client.lookup.data_sets.id(ds_external_id, is_dry_run)
         if security_categories_names := resource.pop("securityCategoryNames", []):
-            security_categories = self.client.lookup.security_categories.id(security_categories_names, is_dry_run)
-            resource["securityCategories"] = security_categories
-        if "assetExternalIds" in resource:
-            resource["assetIds"] = self.client.lookup.assets.id(resource["assetExternalIds"], is_dry_run)
-        return FileMetadataWrite._load(resource)
-
-    def dump_resource(self, resource: FileMetadata, local: dict[str, Any] | None = None) -> dict[str, Any]:
-        dumped = resource.as_write().dump()
-        if ds_id := dumped.pop("dataSetId", None):
-            dumped["dataSetExternalId"] = self.client.lookup.data_sets.external_id(ds_id)
+            resource["securityCategories"] = self.client.lookup.security_categories.id(
+                security_categories_names, is_dry_run
+            )
+        if asset_external_ids := resource.pop("assetExternalIds", None):
+            resource["assetIds"] = self.client.lookup.assets.id(asset_external_ids, is_dry_run)
+        return FileMetadataRequest.model_validate(resource)
+
+    def dump_resource(self, resource: FileMetadataResponse, local: dict[str, Any] | None = None) -> dict[str, Any]:
+        dumped = resource.as_request_resource().dump()
+        if data_set_id := dumped.pop("dataSetId", None):
+            dumped["dataSetExternalId"] = self.client.lookup.data_sets.external_id(data_set_id)
         if security_categories := dumped.pop("securityCategories", []):
             dumped["securityCategoryNames"] = self.client.lookup.security_categories.external_id(security_categories)
         if asset_ids := dumped.pop("assetIds", []):
             dumped["assetExternalIds"] = self.client.lookup.assets.external_id(asset_ids)
         return dumped
 
-    def create(self, items: Sequence[FileMetadataWrite]) -> FileMetadataList:
-        created = FileMetadataList([])
-        for meta in items:
-            try:
-                created.append(self.client.files.create(meta))
-            except CogniteAPIError as e:
-                if e.code == 409:
-                    print(f"  [bold yellow]WARNING:[/] File {meta.external_id} already exists, skipping upload.")
-        return created
+    def create(self, items: Sequence[FileMetadataRequest]) -> list[FileMetadataResponse]:
+        return self.client.tool.filemetadata.create(items, overwrite=True)
 
-    def retrieve(self, ids: SequenceNotStr[str]) -> FileMetadataList:
-        return self.client.files.retrieve_multiple(external_ids=ids, ignore_unknown_ids=True)
+    def retrieve(self, ids: SequenceNotStr[ExternalId]) -> list[FileMetadataResponse]:
+        return self.client.tool.filemetadata.retrieve(list(ids), ignore_unknown_ids=True)
 
-    def update(self, items: Sequence[FileMetadataWrite]) -> FileMetadataList:
-        return self.client.files.update(items, mode="replace")
+    def update(self, items: Sequence[FileMetadataRequest]) -> list[FileMetadataResponse]:
+        return self.client.tool.filemetadata.update(items, mode="replace")
 
-    def delete(self, ids: str | int | SequenceNotStr[str | int] | None) -> int:
-        internal_ids, external_ids = self._split_ids(ids)
-        self.client.files.delete(id=internal_ids, external_id=external_ids)
-        return len(cast(SequenceNotStr[str], ids))
+    def delete(self, ids: SequenceNotStr[InternalOrExternalId]) -> int:
+        if not ids:
+            return 0
+        self.client.tool.filemetadata.delete(list(ids), ignore_unknown_ids=True)
+        return len(ids)
 
     def _iterate(
         self,
         data_set_external_id: str | None = None,
         space: str | None = None,
         parent_ids: list[Hashable] | None = None,
-    ) -> Iterable[FileMetadata]:
-        return iter(self.client.files(data_set_external_ids=[data_set_external_id] if data_set_external_id else None))
-
-    def count(self, ids: SequenceNotStr[str]) -> int:
+    ) -> Iterable[FileMetadataResponse]:
+        cursor: str | None = None
+        while True:
+            page = self.client.tool.filemetadata.iterate(
+                data_set_external_ids=[data_set_external_id] if data_set_external_id else None,
+                limit=1000,
+                cursor=cursor,
+            )
+            yield from page.items
+            if not page.next_cursor or not page.items:
+                break
+            cursor = page.next_cursor
+
+    def count(self, ids: SequenceNotStr[ExternalId]) -> int:
         return sum(
-            1
-            for meta in self.client.files.retrieve_multiple(external_ids=list(ids), ignore_unknown_ids=True)
-            if meta.uploaded
+            1 for meta in self.client.tool.filemetadata.retrieve(list(ids), ignore_unknown_ids=True) if meta.uploaded
         )
 
-    def drop_data(self, ids: SequenceNotStr[str]) -> int:
-        existing = self.client.files.retrieve_multiple(external_ids=list(ids), ignore_unknown_ids=True)
+    def drop_data(self, ids: SequenceNotStr[ExternalId]) -> int:
+        existing = self.client.tool.filemetadata.retrieve(list(ids), ignore_unknown_ids=True)
         # File and FileMetadata is tightly coupled, so we need to delete the metadata and recreate it
         # without the source set to delete the file.
-        deleted_files = self.delete(existing.as_external_ids())
-        self.create(existing.as_write())
+        deleted_files = self.delete([meta.as_id() for meta in existing])
+        self.create([meta.as_request_resource() for meta in existing])
         return deleted_files
 
 

cognite_toolkit/_cdf_tk/cruds/_resource_cruds/relationship.py

@@ -130,7 +130,7 @@ class RelationshipCRUD(ResourceCRUD[str, RelationshipWrite, Relationship]):
                     elif type_value == "timeseries":
                         yield TimeSeriesCRUD, ExternalId(external_id=id_value)
                     elif type_value == "file":
-                        yield FileMetadataCRUD, id_value
+                        yield FileMetadataCRUD, ExternalId(external_id=id_value)
                     elif type_value == "event":
                         yield EventCRUD, ExternalId(external_id=id_value)
 

cognite_toolkit/_cdf_tk/storageio/_asset_centric.py

@@ -3,16 +3,12 @@ from collections import defaultdict
 from collections.abc import Iterable, Sequence
 from typing import Any, ClassVar, Generic
 
-from cognite.client.data_classes import (
-    FileMetadata,
-    FileMetadataList,
-    Label,
-    LabelDefinition,
-)
+from cognite.client.data_classes import Label, LabelDefinition
 
 from cognite_toolkit._cdf_tk.client import ToolkitClient
 from cognite_toolkit._cdf_tk.client.data_classes.asset import AssetAggregateItem, AssetRequest, AssetResponse
 from cognite_toolkit._cdf_tk.client.data_classes.event import EventRequest, EventResponse
+from cognite_toolkit._cdf_tk.client.data_classes.filemetadata import FileMetadataResponse
 from cognite_toolkit._cdf_tk.client.data_classes.identifiers import InternalId
 from cognite_toolkit._cdf_tk.client.data_classes.timeseries import TimeSeriesRequest, TimeSeriesResponse
 from cognite_toolkit._cdf_tk.cruds import (
@@ -107,13 +103,16 @@ class AssetCentricIO(
 
     def _collect_dependencies(
         self,
-        resources: Sequence[AssetResponse] | FileMetadataList | Sequence[TimeSeriesResponse] | Sequence[EventResponse],
+        resources: Sequence[AssetResponse]
+        | Sequence[FileMetadataResponse]
+        | Sequence[TimeSeriesResponse]
+        | Sequence[EventResponse],
         selector: AssetCentricSelector,
     ) -> None:
         for resource in resources:
             if resource.data_set_id:
                 self._downloaded_data_sets_by_selector[selector].add(resource.data_set_id)
-            if isinstance(resource, AssetResponse | FileMetadata):
+            if isinstance(resource, AssetResponse | FileMetadataResponse):
                 for label in resource.labels or []:
                     if isinstance(label, str):
                         self._downloaded_labels_by_selector[selector].add(label)
@@ -146,18 +145,18 @@ class AssetCentricIO(
         return f"INTERNAL_ID_project_{project}_{internal_id!s}"
 
     def _populate_data_set_id_cache(
-        self, chunk: Sequence[AssetResponse | FileMetadata | TimeSeriesResponse | EventResponse]
+        self, chunk: Sequence[AssetResponse | FileMetadataResponse | TimeSeriesResponse | EventResponse]
     ) -> None:
         data_set_ids = {item.data_set_id for item in chunk if item.data_set_id is not None}
         self.client.lookup.data_sets.external_id(list(data_set_ids))
 
-    def _populate_security_category_cache(self, chunk: Sequence[FileMetadata | TimeSeriesResponse]) -> None:
+    def _populate_security_category_cache(self, chunk: Sequence[FileMetadataResponse | TimeSeriesResponse]) -> None:
         security_category_ids: set[int] = set()
         for item in chunk:
             security_category_ids.update(item.security_categories or [])
         self.client.lookup.security_categories.external_id(list(security_category_ids))
 
-    def _populate_asset_id_cache(self, chunk: Sequence[FileMetadata | EventResponse]) -> None:
+    def _populate_asset_id_cache(self, chunk: Sequence[FileMetadataResponse | EventResponse]) -> None:
         asset_ids: set[int] = set()
         for item in chunk:
             asset_ids.update(item.asset_ids or [])
@@ -379,7 +378,7 @@ class AssetIO(UploadableAssetCentricIO[AssetResponse, AssetRequest]):
             current_depth += 1
 
 
-class FileMetadataIO(AssetCentricIO[FileMetadata]):
+class FileMetadataIO(AssetCentricIO[FileMetadataResponse]):
     KIND = "FileMetadata"
     RESOURCE_TYPE = "file"
     SUPPORTED_DOWNLOAD_FORMATS = frozenset({".parquet", ".csv", ".ndjson"})
@@ -391,7 +390,7 @@ class FileMetadataIO(AssetCentricIO[FileMetadata]):
         super().__init__(client)
         self._crud = FileMetadataCRUD.create_loader(self.client)
 
-    def as_id(self, item: FileMetadata) -> str:
+    def as_id(self, item: FileMetadataResponse) -> str:
         return item.external_id if item.external_id is not None else self._create_identifier(item.id)
 
     def _get_aggregator(self) -> AssetCentricAggregator:
@@ -434,25 +433,31 @@ class FileMetadataIO(AssetCentricIO[FileMetadata]):
         ]
         return file_schema + metadata_schema
 
-    def stream_data(self, selector: AssetCentricSelector, limit: int | None = None) -> Iterable[Page[FileMetadata]]:
+    def stream_data(
+        self, selector: AssetCentricSelector, limit: int | None = None
+    ) -> Iterable[Page[FileMetadataResponse]]:
         asset_subtree_external_ids, data_set_external_ids = self._get_hierarchy_dataset_pair(selector)
-        for file_list in self.client.files(
-            chunk_size=self.CHUNK_SIZE,
-            limit=limit,
-            asset_subtree_external_ids=asset_subtree_external_ids,
-            data_set_external_ids=data_set_external_ids,
-            # We cannot use partitions here as it is not thread safe. This spawn multiple threads
-            # that are not shut down until all data is downloaded. We need to be able to abort.
-            partitions=None,
-        ):
-            self._collect_dependencies(file_list, selector)
-            yield Page(worker_id="main", items=file_list)
-
-    def retrieve(self, ids: Sequence[int]) -> FileMetadataList:
-        return self.client.files.retrieve_multiple(ids)
+        cursor: str | None = None
+        total_count = 0
+        while True:
+            page = self.client.tool.filemetadata.iterate(
+                data_set_external_ids=data_set_external_ids,
+                asset_subtree_external_ids=asset_subtree_external_ids,
+                limit=self.CHUNK_SIZE,
+                cursor=cursor,
+            )
+            self._collect_dependencies(page.items, selector)
+            yield Page(worker_id="main", items=page.items)
+            total_count += len(page.items)
+            if page.next_cursor is None or (limit is not None and total_count >= limit):
+                break
+            cursor = page.next_cursor
+
+    def retrieve(self, ids: Sequence[int]) -> list[FileMetadataResponse]:
+        return self.client.tool.filemetadata.retrieve(InternalId.from_ids(ids))
 
     def data_to_json_chunk(
-        self, data_chunk: Sequence[FileMetadata], selector: AssetCentricSelector | None = None
+        self, data_chunk: Sequence[FileMetadataResponse], selector: AssetCentricSelector | None = None
     ) -> list[dict[str, JsonVal]]:
         # Ensure data sets/assets/security-categories are looked up to populate cache.
         # This is to avoid looking up each data set id individually in the .dump_resource call