cognite-toolkit 0.7.55__py3-none-any.whl → 0.7.57__py3-none-any.whl

cognite_toolkit/_cdf_tk/commands/_base.py

@@ -5,6 +5,8 @@ from typing import Any
 from rich import print
 from rich.console import Console
 
+from cognite_toolkit._cdf_tk.client import ToolkitClient
+from cognite_toolkit._cdf_tk.client.http_client import ToolkitAPIError
 from cognite_toolkit._cdf_tk.data_classes import CommandTrackingInfo
 from cognite_toolkit._cdf_tk.tk_warnings import (
     ToolkitWarning,
@@ -16,12 +18,27 @@ _HAS_PRINTED_COLLECT_MESSAGE = False
 
 
 class ToolkitCommand:
-    def __init__(self, print_warning: bool = True, skip_tracking: bool = False, silent: bool = False):
+    def __init__(
+        self,
+        print_warning: bool = True,
+        skip_tracking: bool = False,
+        silent: bool = False,
+        client: ToolkitClient | None = None,
+    ):
         self._print_warning = print_warning
         self.silent = silent
         self.warning_list = WarningList[ToolkitWarning]()
         self.tracker = Tracker(skip_tracking)
         self._additional_tracking_info = CommandTrackingInfo()
+        if client is not None:
+            self._additional_tracking_info.cluster = client.config.cdf_cluster
+            try:
+                result = client.project.organization()
+            except (ToolkitAPIError, ValueError):
+                self._additional_tracking_info.project = client.config.project
+            else:
+                self._additional_tracking_info.organization = result.organization
+                self._additional_tracking_info.project = result.name
 
     @property
     def print_warning(self) -> bool:

cognite_toolkit/_cdf_tk/commands/_import_cmd.py

@@ -21,8 +21,10 @@ class ImportTransformationCLI(ToolkitCommand):
         get_client: Callable[[], ToolkitClient] | None = None,
         print_warning: bool = True,
         skip_tracking: bool = False,
+        silent: bool = False,
+        client: ToolkitClient | None = None,
     ):
-        super().__init__(print_warning, skip_tracking)
+        super().__init__(print_warning, skip_tracking, silent, get_client() if get_client is not None else client)
         self._dataset_external_id_by_id: dict[int, str] = {}
         # We only initialize the client if we need to look up dataset ids.
         self._client: ToolkitClient | None = None

cognite_toolkit/_cdf_tk/commands/_migrate/command.py

@@ -9,10 +9,10 @@ from rich.table import Table
 
 from cognite_toolkit._cdf_tk.client import ToolkitClient
 from cognite_toolkit._cdf_tk.client.http_client import (
-    FailedRequestItems,
-    FailedResponseItems,
     HTTPClient,
-    SuccessResponseItems,
+    ItemsFailedRequest,
+    ItemsFailedResponse,
+    ItemsSuccessResponse,
 )
 from cognite_toolkit._cdf_tk.commands._base import ToolkitCommand
 from cognite_toolkit._cdf_tk.commands._migrate.creators import MigrationCreator
@@ -187,20 +187,20 @@ class MigrationCommand(ToolkitCommand):
             # Todo: Move logging into the UploadableStorageIO class
             issues: list[WriteIssue] = []
             for item in responses:
-                if isinstance(item, SuccessResponseItems):
+                if isinstance(item, ItemsSuccessResponse):
                     target.logger.tracker.finalize_item(item.ids, "success")
                     continue
-                if isinstance(item, FailedResponseItems):
+                if isinstance(item, ItemsFailedResponse):
                     error = item.error
                     for id_ in item.ids:
                         issue = WriteIssue(id=str(id_), status_code=error.code, message=error.message)
                         issues.append(issue)
-                elif isinstance(item, FailedRequestItems):
+                elif isinstance(item, ItemsFailedRequest):
                     for id_ in item.ids:
-                        issue = WriteIssue(id=str(id_), status_code=0, message=item.error)
+                        issue = WriteIssue(id=str(id_), status_code=0, message=item.error_message)
                         issues.append(issue)
 
-                if isinstance(item, FailedResponseItems | FailedRequestItems):
+                if isinstance(item, ItemsFailedResponse | ItemsFailedRequest):
                     target.logger.tracker.finalize_item(item.ids, "failure")
             if issues:
                 target.logger.log(issues)

cognite_toolkit/_cdf_tk/commands/_migrate/migration_io.py

@@ -6,14 +6,16 @@ from cognite.client.data_classes.data_modeling import EdgeId, InstanceApply, Nod
 
 from cognite_toolkit._cdf_tk.client import ToolkitClient
 from cognite_toolkit._cdf_tk.client.http_client import (
-    FailedResponse,
     HTTPClient,
-    HTTPMessage,
-    ItemsRequest,
-    RequestMessage2,
-    SuccessResponseItems,
+    RequestMessage,
     ToolkitAPIError,
 )
+from cognite_toolkit._cdf_tk.client.http_client._item_classes import (
+    ItemsFailedResponse,
+    ItemsRequest,
+    ItemsResultList,
+    ItemsSuccessResponse,
+)
 from cognite_toolkit._cdf_tk.client.resource_classes.legacy.pending_instances_ids import PendingInstanceId
 from cognite_toolkit._cdf_tk.client.resource_classes.three_d import (
     AssetMappingClassicResponse,
@@ -170,20 +172,20 @@ class AssetCentricMigrationIO(
         data_chunk: Sequence[UploadItem[InstanceApply]],
         http_client: HTTPClient,
         selector: AssetCentricMigrationSelector | None = None,
-    ) -> Sequence[HTTPMessage]:
+    ) -> ItemsResultList:
         """Upload items by first linking them using files/set-pending-instance-ids and then uploading the instances."""
         if self.skip_linking:
-            return list(super().upload_items(data_chunk, http_client, None))
+            return super().upload_items(data_chunk, http_client, None)
         elif selector is None:
             raise ToolkitNotImplementedError(f"Selector must be provided for uploading {self.KIND} items.")
         elif selector.kind not in self.PENDING_INSTANCE_ID_ENDPOINT_BY_KIND:
-            return list(super().upload_items(data_chunk, http_client, None))
+            return super().upload_items(data_chunk, http_client, None)
 
         pending_instance_id_endpoint = self.PENDING_INSTANCE_ID_ENDPOINT_BY_KIND[selector.kind]
-        results: list[HTTPMessage] = []
+        results = ItemsResultList()
         to_upload = self.link_asset_centric(data_chunk, http_client, pending_instance_id_endpoint)
         if to_upload:
-            results.extend(list(super().upload_items(to_upload, http_client, None)))
+            results.extend(super().upload_items(to_upload, http_client, None))
         return results
 
     @classmethod
@@ -197,7 +199,7 @@ class AssetCentricMigrationIO(
         config = http_client.config
         successful_linked: set[str] = set()
         for batch in chunker_sequence(data_chunk, cls.CHUNK_SIZE):
-            batch_results = http_client.request_with_retries(
+            batch_results = http_client.request_items_retries(
                 message=ItemsRequest(
                     endpoint_url=config.create_api_url(pending_instance_id_endpoint),
                     method="POST",
@@ -209,7 +211,7 @@ class AssetCentricMigrationIO(
                 )
             )
             for res in batch_results:
-                if isinstance(res, SuccessResponseItems):
+                if isinstance(res, ItemsSuccessResponse):
                     successful_linked.update(res.ids)
         to_upload = [item for item in data_chunk if item.source_id in successful_linked]
         return to_upload
@@ -441,31 +443,31 @@ class ThreeDMigrationIO(UploadableStorageIO[ThreeDSelector, ThreeDModelResponse,
         data_chunk: Sequence[UploadItem[ThreeDMigrationRequest]],
         http_client: HTTPClient,
         selector: ThreeDSelector | None = None,
-    ) -> Sequence[HTTPMessage]:
+    ) -> ItemsResultList:
         """Migrate 3D models by uploading them to the migrate/models endpoint."""
         if len(data_chunk) > self.CHUNK_SIZE:
             raise RuntimeError(f"Uploading more than {self.CHUNK_SIZE} 3D models at a time is not supported.")
 
-        results: list[HTTPMessage] = []
-        responses = http_client.request_with_retries(
+        results = ItemsResultList()
+        responses = http_client.request_items_retries(
             message=ItemsRequest(
                 endpoint_url=self.client.config.create_api_url(self.UPLOAD_ENDPOINT),
                 method="POST",
-                items=list(data_chunk),
+                items=data_chunk,
             )
         )
         if (
-            failed_response := next((res for res in responses if isinstance(res, FailedResponse)), None)
+            failed_response := next((res for res in responses if isinstance(res, ItemsFailedResponse)), None)
         ) and failed_response.status_code == 400:
             raise ToolkitAPIError("3D model migration failed. You need to enable the 3D migration alpha feature flag.")
 
         results.extend(responses)
-        success_ids = {id for res in responses if isinstance(res, SuccessResponseItems) for id in res.ids}
+        success_ids = {id for res in responses if isinstance(res, ItemsSuccessResponse) for id in res.ids}
         for data in data_chunk:
             if data.source_id not in success_ids:
                 continue
             revision = http_client.request_single_retries(
-                message=RequestMessage2(
+                message=RequestMessage(
                     endpoint_url=self.client.config.create_api_url(self.REVISION_ENDPOINT),
                     method="POST",
                     body_content={"items": [data.item.revision.dump(camel_case=True)]},
@@ -535,21 +537,21 @@ class ThreeDAssetMappingMigrationIO(
         data_chunk: Sequence[UploadItem[AssetMappingDMRequest]],
         http_client: HTTPClient,
         selector: T_Selector | None = None,
-    ) -> Sequence[HTTPMessage]:
+    ) -> ItemsResultList:
         """Migrate 3D asset mappings by uploading them to the migrate/asset-mappings endpoint."""
         if not data_chunk:
-            return []
+            return ItemsResultList()
         # Assume all items in the chunk belong to the same model and revision, they should
         # if the .stream_data method is used for downloading.
         first = data_chunk[0]
         model_id = first.item.model_id
         revision_id = first.item.revision_id
         endpoint = self.UPLOAD_ENDPOINT.format(modelId=model_id, revisionId=revision_id)
-        responses = http_client.request_with_retries(
+        return http_client.request_items_retries(
             ItemsRequest(
                 endpoint_url=self.client.config.create_api_url(endpoint),
                 method="POST",
-                items=list(data_chunk),
+                items=data_chunk,
                 extra_body_fields={
                     "dmsContextualizationConfig": {
                         "object3DSpace": self.object_3D_space,
@@ -558,7 +560,6 @@ class ThreeDAssetMappingMigrationIO(
                 },
             )
         )
-        return responses
 
     def json_to_resource(self, item_json: dict[str, JsonVal]) -> AssetMappingDMRequest:
         raise NotImplementedError("Deserializing 3D Asset Mappings from JSON is not supported.")

cognite_toolkit/_cdf_tk/commands/_profile.py

@@ -71,8 +71,9 @@ class ProfileCommand(ToolkitCommand, ABC, Generic[T_Index]):
         print_warning: bool = True,
         skip_tracking: bool = False,
         silent: bool = False,
+        client: ToolkitClient | None = None,
     ) -> None:
-        super().__init__(print_warning, skip_tracking, silent)
+        super().__init__(print_warning, skip_tracking, silent, client)
         self.table_title = self.__class__.__name__.removesuffix("Command")
         self.output_spreadsheet: Path | None = output_spreadsheet
         if output_spreadsheet is not None:
@@ -347,8 +348,9 @@ class ProfileAssetCommand(ProfileCommand[AssetIndex]):
         print_warning: bool = True,
         skip_tracking: bool = False,
         silent: bool = False,
+        client: ToolkitClient | None = None,
     ) -> None:
-        super().__init__(output_spreadsheet, print_warning, skip_tracking, silent)
+        super().__init__(output_spreadsheet, print_warning, skip_tracking, silent, client)
         self.table_title = "Asset Profile for Hierarchy"
         self.hierarchy: str | None = None
         self.aggregators: dict[str, MetadataAggregator] = {}
@@ -641,8 +643,9 @@ class ProfileAssetCentricCommand(ProfileCommand[str]):
         print_warning: bool = True,
         skip_tracking: bool = False,
         silent: bool = False,
+        client: ToolkitClient | None = None,
     ) -> None:
-        super().__init__(output_spreadsheet, print_warning, skip_tracking, silent)
+        super().__init__(output_spreadsheet, print_warning, skip_tracking, silent, client)
         self.hierarchy: str | None = None
         self.table_title = "Asset Centric Profile"
         self.aggregators: dict[str, AssetCentricAggregator] = {}
@@ -742,8 +745,9 @@ class ProfileTransformationCommand(ProfileCommand[str]):
         print_warning: bool = True,
         skip_tracking: bool = False,
         silent: bool = False,
+        client: ToolkitClient | None = None,
     ) -> None:
-        super().__init__(output_spreadsheet, print_warning, skip_tracking, silent)
+        super().__init__(output_spreadsheet, print_warning, skip_tracking, silent, client)
         self.table_title = "Transformation Profile"
         self.destination_type: AssetCentricDestinationType | None = None
 
@@ -826,8 +830,9 @@ class ProfileRawCommand(ProfileCommand[RawProfileIndex]):
         print_warning: bool = True,
         skip_tracking: bool = False,
         silent: bool = False,
+        client: ToolkitClient | None = None,
     ) -> None:
-        super().__init__(output_spreadsheet, print_warning, skip_tracking, silent)
+        super().__init__(output_spreadsheet, print_warning, skip_tracking, silent, client)
         self.table_title = "RAW Profile"
         self.destination_type: AssetCentricDestinationType | None = None
         self.client: ToolkitClient | None = None

cognite_toolkit/_cdf_tk/commands/_purge.py

@@ -1,9 +1,9 @@
 import uuid
 from abc import ABC, abstractmethod
-from collections.abc import Callable, Hashable, Iterable, Sequence
+from collections.abc import Callable, Iterable, Sequence
 from dataclasses import dataclass
 from functools import partial
-from typing import Literal, cast
+from typing import Any, Literal, cast
 
 import questionary
 from cognite.client.data_classes import DataSetUpdate
@@ -11,17 +11,17 @@ from cognite.client.data_classes.data_modeling import Edge, NodeId
 from cognite.client.data_classes.data_modeling.statistics import SpaceStatistics
 from cognite.client.exceptions import CogniteAPIError
 from cognite.client.utils._identifier import InstanceId
+from pydantic import JsonValue
 from rich import print
 from rich.console import Console
 from rich.panel import Panel
 
 from cognite_toolkit._cdf_tk.client import ToolkitClient
+from cognite_toolkit._cdf_tk.client._resource_base import RequestItem
 from cognite_toolkit._cdf_tk.client.http_client import (
     HTTPClient,
     ItemsRequest,
-    ItemsRequest2,
-    ItemsSuccessResponse2,
-    SuccessResponseItems,
+    ItemsSuccessResponse,
 )
 from cognite_toolkit._cdf_tk.client.resource_classes.instance_api import TypedInstanceIdentifier
 from cognite_toolkit._cdf_tk.cruds import (
@@ -79,15 +79,14 @@ class DeleteResults:
     failed: int = 0
 
 
-@dataclass
-class DeleteItem:
-    item: JsonVal
-    as_id_fun: Callable[[JsonVal], Hashable]
+class DeleteItem(RequestItem):
+    item: JsonValue
+    as_id_fun: Callable[[JsonValue], str]
 
-    def dump(self) -> JsonVal:
-        return self.item
+    def dump(self, camel_case: bool = True, exclude_extra: bool = False) -> dict[str, Any]:
+        return self.item  # type: ignore[return-value]
 
-    def as_id(self) -> Hashable:
+    def __str__(self) -> str:
         return self.as_id_fun(self.item)
 
 
@@ -222,7 +221,7 @@ class PurgeCommand(ToolkitCommand):
         if not dry_run and not auto_yes:
             confirm = questionary.confirm(
                 f"Are you really sure you want to purge the {selected_space!r} space?", default=False
-            ).unsafe_ask()
+            ).ask()
             if not confirm:
                 return DeployResults([], "purge", dry_run=dry_run)
 
@@ -364,15 +363,15 @@ class PurgeCommand(ToolkitCommand):
     ) -> Callable[[list[JsonVal]], None]:
         crud = delete_item.crud
 
-        def as_id(item: JsonVal) -> Hashable:
+        def as_id(item: JsonVal) -> str:
             try:
-                return crud.get_id(item)
+                return str(crud.get_id(item))
             except KeyError:
                 # Fallback to internal ID
-                return crud.get_internal_id(item)
+                return str(crud.get_internal_id(item))
 
         def process(items: list[JsonVal]) -> None:
-            responses = delete_client.request_with_retries(
+            responses = delete_client.request_items_retries(
                 ItemsRequest(
                     endpoint_url=delete_url,
                     method="POST",
@@ -381,10 +380,10 @@ class PurgeCommand(ToolkitCommand):
                 )
             )
             for response in responses:
-                if isinstance(response, SuccessResponseItems):
+                if isinstance(response, ItemsSuccessResponse):
                     result.deleted += len(response.ids)
                 else:
-                    result.unchanged += len(items)
+                    result.unchanged += len(response.ids)
 
         return process
 
@@ -423,7 +422,7 @@ class PurgeCommand(ToolkitCommand):
         if not dry_run and not auto_yes:
             confirm = questionary.confirm(
                 f"Are you really sure you want to purge the {selected_data_set_external_id!r} dataSet?", default=False
-            ).unsafe_ask()
+            ).ask()
             if not confirm:
                 return DeployResults([], "purge", dry_run=dry_run)
 
@@ -595,7 +594,7 @@ class PurgeCommand(ToolkitCommand):
                 confirm = questionary.confirm(
                     f"Are you really sure you want to purge all {total:,} instances in {selector!s}?",
                     default=False,
-                ).unsafe_ask()
+                ).ask()
                 if not confirm:
                     return DeleteResults()
 
@@ -716,14 +715,14 @@ class PurgeCommand(ToolkitCommand):
             return
 
         responses = delete_client.request_items_retries(
-            ItemsRequest2(
+            ItemsRequest(
                 endpoint_url=delete_client.config.create_api_url("/models/instances/delete"),
                 method="POST",
                 items=[TypedInstanceIdentifier._load(item) for item in items],
             )
         )
         for response in responses:
-            if isinstance(response, ItemsSuccessResponse2):
+            if isinstance(response, ItemsSuccessResponse):
                 results.deleted += len(response.ids)
             else:
                 results.failed += len(response.ids)
@@ -735,17 +734,13 @@ class PurgeCommand(ToolkitCommand):
         node_ids = [instance for instance in instances if isinstance(instance, NodeId)]
         if node_ids:
             timeseries = client.time_series.retrieve_multiple(instance_ids=node_ids, ignore_unknown_ids=True)
-            migrated_timeseries_ids = [
-                ts.id
-                for ts in timeseries
-                if ts.instance_id and ts.pending_instance_id  # type: ignore[attr-defined]
-            ]
             if not dry_run and timeseries:
+                migrated_timeseries_ids = [ts.id for ts in timeseries if ts.instance_id and ts.pending_instance_id]  # type: ignore[attr-defined]
                 client.time_series.unlink_instance_ids(id=migrated_timeseries_ids)
                 if verbose:
                     console.print(f"Unlinked {len(migrated_timeseries_ids)} timeseries from datapoints.")
             elif verbose and timeseries:
-                console.print(f"Would have unlinked {len(migrated_timeseries_ids)} timeseries from datapoints.")
+                console.print(f"Would have unlinked {len(timeseries)} timeseries from datapoints.")
         return instances
 
     @staticmethod
@@ -755,15 +750,15 @@ class PurgeCommand(ToolkitCommand):
         file_ids = [instance for instance in instances if isinstance(instance, NodeId)]
         if file_ids:
             files = client.files.retrieve_multiple(instance_ids=file_ids, ignore_unknown_ids=True)
-            migrated_file_ids = [
-                file.id
-                for file in files
-                if file.instance_id and file.pending_instance_id and file.id is not None  # type: ignore[attr-defined]
-            ]
             if not dry_run and files:
+                migrated_file_ids = [
+                    file.id
+                    for file in files
+                    if file.instance_id and file.pending_instance_id and file.id is not None  # type: ignore[attr-defined]
+                ]
                 client.files.unlink_instance_ids(id=migrated_file_ids)
                 if verbose:
                     console.print(f"Unlinked {len(migrated_file_ids)} files from nodes.")
             elif verbose and files:
-                console.print(f"Would have unlinked {len(migrated_file_ids)} files from their blob content.")
+                console.print(f"Would have unlinked {len(files)} files from their blob content.")
         return instances

cognite_toolkit/_cdf_tk/commands/_upload.py

@@ -3,13 +3,11 @@ from collections.abc import Sequence
 from functools import partial
 from pathlib import Path
 
-from cognite.client.data_classes.data_modeling import (
-    ViewId,
-)
+from cognite.client.data_classes.data_modeling import ViewId
 from rich.console import Console
 
 from cognite_toolkit._cdf_tk.client import ToolkitClient
-from cognite_toolkit._cdf_tk.client.http_client import HTTPClient, ItemMessage, SuccessResponseItems
+from cognite_toolkit._cdf_tk.client.http_client import HTTPClient, ItemsResultMessage, ItemsSuccessResponse
 from cognite_toolkit._cdf_tk.constants import DATA_MANIFEST_SUFFIX, DATA_RESOURCE_DIR
 from cognite_toolkit._cdf_tk.cruds import ViewCRUD
 from cognite_toolkit._cdf_tk.exceptions import ToolkitValueError
@@ -279,10 +277,10 @@ class UploadCommand(ToolkitCommand):
             return
         results = io.upload_items(data_chunk, upload_client, selector)
         for message in results:
-            if isinstance(message, SuccessResponseItems):
+            if isinstance(message, ItemsSuccessResponse):
                 for id_ in message.ids:
                     tracker.set_progress(id_, step=cls._UPLOAD, status="success")
-            elif isinstance(message, ItemMessage):
+            elif isinstance(message, ItemsResultMessage):
                 for id_ in message.ids:
                     tracker.set_progress(id_, step=cls._UPLOAD, status="failed")
             else:

cognite_toolkit/_cdf_tk/commands/build_cmd.py

@@ -99,8 +99,9 @@ class BuildCommand(ToolkitCommand):
         print_warning: bool = True,
         skip_tracking: bool = False,
         silent: bool = False,
+        client: ToolkitClient | None = None,
     ) -> None:
-        super().__init__(print_warning, skip_tracking, silent)
+        super().__init__(print_warning, skip_tracking, silent, client)
         self.existing_resources_by_loader: dict[type[ResourceCRUD], set[Hashable]] = defaultdict(set)
         self.instantiated_loaders: dict[type[ResourceCRUD], ResourceCRUD] = {}
 

cognite_toolkit/_cdf_tk/commands/build_v2/build_cmd.py

@@ -21,8 +21,14 @@ from cognite_toolkit._version import __version__
 
 
 class BuildCommand(ToolkitCommand):
-    def __init__(self, print_warning: bool = True, skip_tracking: bool = False, silent: bool = False) -> None:
-        super().__init__(print_warning, skip_tracking, silent)
+    def __init__(
+        self,
+        print_warning: bool = True,
+        skip_tracking: bool = False,
+        silent: bool = False,
+        client: ToolkitClient | None = None,
+    ) -> None:
+        super().__init__(print_warning, skip_tracking, silent, client)
         self.issues = IssueList()
 
     def execute(

cognite_toolkit/_cdf_tk/commands/deploy.py

@@ -62,8 +62,14 @@ from ._utils import _print_ids_or_length
 
 
 class DeployCommand(ToolkitCommand):
-    def __init__(self, print_warning: bool = True, skip_tracking: bool = False, silent: bool = False) -> None:
-        super().__init__(print_warning, skip_tracking, silent)
+    def __init__(
+        self,
+        print_warning: bool = True,
+        skip_tracking: bool = False,
+        silent: bool = False,
+        client: ToolkitClient | None = None,
+    ) -> None:
+        super().__init__(print_warning, skip_tracking, silent, client)
         self._clean_command = CleanCommand(print_warning, skip_tracking=True)
 
     def deploy_build_directory(

cognite_toolkit/_cdf_tk/commands/init.py

@@ -13,6 +13,7 @@ from rich import print
 from rich.panel import Panel
 
 from cognite_toolkit._cdf_tk.cdf_toml import CDFToml
+from cognite_toolkit._cdf_tk.client import ToolkitClient
 from cognite_toolkit._cdf_tk.commands._base import ToolkitCommand
 from cognite_toolkit._cdf_tk.commands.auth import AuthCommand
 from cognite_toolkit._cdf_tk.commands.collect import CollectCommand
@@ -54,8 +55,14 @@ class InitChecklistItem:
 class InitCommand(ToolkitCommand):
     organization_dir: Path | None
 
-    def __init__(self, print_warning: bool = True, skip_tracking: bool = False, silent: bool = False) -> None:
-        super().__init__(print_warning, skip_tracking, silent)
+    def __init__(
+        self,
+        print_warning: bool = True,
+        skip_tracking: bool = False,
+        silent: bool = False,
+        client: ToolkitClient | None = None,
+    ) -> None:
+        super().__init__(print_warning, skip_tracking, silent, client)
         self.organization_dir = None
 
     def execute(self, dry_run: bool = False) -> None:

cognite_toolkit/_cdf_tk/commands/modules.py

@@ -23,6 +23,7 @@ from rich.table import Table
 from rich.tree import Tree
 
 from cognite_toolkit._cdf_tk.cdf_toml import CDFToml, Library
+from cognite_toolkit._cdf_tk.client import ToolkitClient
 from cognite_toolkit._cdf_tk.commands import _cli_commands as CLICommands
 from cognite_toolkit._cdf_tk.commands._base import ToolkitCommand
 from cognite_toolkit._cdf_tk.commands._changes import (
@@ -81,10 +82,11 @@ class ModulesCommand(ToolkitCommand):
         print_warning: bool = True,
         skip_tracking: bool = False,
         silent: bool = False,
+        client: ToolkitClient | None = None,
         temp_dir_suffix: str | None = None,
         module_source_dir: Path | None = None,
     ):
-        super().__init__(print_warning, skip_tracking, silent)
+        super().__init__(print_warning, skip_tracking, silent, client)
         self._module_source_dir: Path | None = module_source_dir
         # Use suffix to make temp directory unique (useful for parallel test execution)
         modules_dir_name = f"{MODULES}.{temp_dir_suffix}" if temp_dir_suffix else MODULES

cognite_toolkit/_cdf_tk/commands/pull.py

@@ -411,8 +411,14 @@ class TextFileDifference(UserList):
 
 
 class PullCommand(ToolkitCommand):
-    def __init__(self, print_warning: bool = True, skip_tracking: bool = False, silent: bool = False) -> None:
-        super().__init__(print_warning, skip_tracking, silent)
+    def __init__(
+        self,
+        print_warning: bool = True,
+        skip_tracking: bool = False,
+        silent: bool = False,
+        client: ToolkitClient | None = None,
+    ) -> None:
+        super().__init__(print_warning, skip_tracking, silent, client)
         self._clean_command = CleanCommand(print_warning, skip_tracking=True)
 
     def pull_module(

cognite_toolkit/_cdf_tk/commands/repo.py

@@ -5,6 +5,7 @@ from pathlib import Path
 import questionary
 
 import cognite_toolkit
+from cognite_toolkit._cdf_tk.client import ToolkitClient
 from cognite_toolkit._cdf_tk.constants import REPO_FILES_DIR
 from cognite_toolkit._cdf_tk.exceptions import ToolkitValueError
 from cognite_toolkit._cdf_tk.tk_warnings import LowSeverityWarning, MediumSeverityWarning
@@ -27,8 +28,9 @@ class RepoCommand(ToolkitCommand):
         skip_tracking: bool = False,
         silent: bool = False,
         skip_git_verify: bool = False,
+        client: ToolkitClient | None = None,
     ) -> None:
-        super().__init__(print_warning=print_warning, skip_tracking=skip_tracking, silent=silent)
+        super().__init__(print_warning=print_warning, skip_tracking=skip_tracking, silent=silent, client=client)
         self._repo_files = Path(resources.files(cognite_toolkit.__name__)) / REPO_FILES_DIR  # type: ignore [arg-type]
         self.skip_git_verify = skip_git_verify
 

cognite_toolkit/_cdf_tk/commands/resources.py

@@ -17,9 +17,6 @@ from cognite_toolkit._cdf_tk.utils.file import yaml_safe_dump
 
 
 class ResourcesCommand(ToolkitCommand):
-    def __init__(self, print_warning: bool = True, skip_tracking: bool = False, silent: bool = False) -> None:
-        super().__init__(print_warning, skip_tracking, silent)
-
     def _get_or_prompt_module_path(self, module: str | None, organization_dir: Path, verbose: bool) -> Path:
         """
         Check if the module exists in the organization directory and return the module path.

cognite_toolkit/_cdf_tk/data_classes/_tracking_info.py

@@ -25,6 +25,7 @@ class CommandTrackingInfo(BaseModel):
 
     project: str | None = Field(default=None)
     cluster: str | None = Field(default=None)
+    organization: str | None = Field(default=None)
     module_ids: set[str] = Field(default_factory=set, alias="moduleIds")
     package_ids: set[str] = Field(default_factory=set, alias="packageIds")
     installed_module_ids: set[str] = Field(default_factory=set, alias="installedModuleIds")

cognite_toolkit/_cdf_tk/protocols.py

@@ -1,6 +1,6 @@
 import sys
 from collections.abc import Collection, Iterator
-from typing import Any, Generic, Protocol, TypeVar
+from typing import Any, Generic, Protocol, TypeVar, runtime_checkable
 
 from cognite.client import CogniteClient
 
@@ -10,6 +10,7 @@ else:
     from typing_extensions import Self
 
 
+@runtime_checkable
 class ResourceRequestProtocol(Protocol):
     @classmethod
     def _load(cls, data: dict[str, Any]) -> Self: ...
@@ -17,6 +18,7 @@ class ResourceRequestProtocol(Protocol):
     def dump(self, camel_case: bool = True) -> dict[str, Any]: ...
 
 
+@runtime_checkable
 class ResourceResponseProtocol(Protocol):
     def as_write(self) -> ResourceRequestProtocol: ...