cognite-toolkit 0.5.44__py3-none-any.whl → 0.7.34__py3-none-any.whl

cognite_toolkit/_cdf_tk/{loaders/_resource_loaders/transformation_loaders.py → cruds/_resource_cruds/transformation.py}

@@ -24,13 +24,14 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from __future__ import annotations
 
+
+import random
+import time
 import warnings
 from collections import defaultdict
-from collections.abc import Hashable, Iterable, Sequence
+from collections.abc import Callable, Hashable, Iterable, Sequence
 from copy import deepcopy
-from functools import lru_cache
 from pathlib import Path
 from typing import Any, Literal, cast, final
 
@@ -44,9 +45,7 @@ from cognite.client.data_classes import (
     TransformationSchedule,
     TransformationScheduleList,
     TransformationScheduleWrite,
-    TransformationScheduleWriteList,
     TransformationWrite,
-    TransformationWriteList,
 )
 from cognite.client.data_classes.capabilities import (
     Capability,
@@ -59,17 +58,16 @@ from cognite.client.data_classes.data_modeling.ids import (
 from cognite.client.data_classes.transformations import NonceCredentials
 from cognite.client.data_classes.transformations.notifications import (
     TransformationNotificationWrite,
-    TransformationNotificationWriteList,
 )
 from cognite.client.exceptions import CogniteAPIError, CogniteAuthError, CogniteDuplicatedError, CogniteNotFoundError
 from cognite.client.utils.useful_types import SequenceNotStr
 from rich import print
 from rich.console import Console
 
-from cognite_toolkit._cdf_tk._parameters import ANY_INT, ParameterSpec, ParameterSpecSet
 from cognite_toolkit._cdf_tk.client import ToolkitClient
 from cognite_toolkit._cdf_tk.client.data_classes.raw import RawDatabase, RawTable
 from cognite_toolkit._cdf_tk.constants import BUILD_FOLDER_ENCODING
+from cognite_toolkit._cdf_tk.cruds._base_cruds import ResourceCRUD
 from cognite_toolkit._cdf_tk.exceptions import (
     ResourceCreationError,
     ToolkitFileNotFoundError,
@@ -78,8 +76,12 @@ from cognite_toolkit._cdf_tk.exceptions import (
     ToolkitRequiredValueError,
     ToolkitYAMLFormatError,
 )
-from cognite_toolkit._cdf_tk.loaders._base_loaders import ResourceLoader
-from cognite_toolkit._cdf_tk.resource_classes import TransformationYAML
+from cognite_toolkit._cdf_tk.resource_classes import (
+    TransformationNotificationYAML,
+    TransformationScheduleYAML,
+    TransformationYAML,
+)
+from cognite_toolkit._cdf_tk.tk_warnings import HighSeverityWarning, MediumSeverityWarning
 from cognite_toolkit._cdf_tk.utils import (
     calculate_secure_hash,
     humanize_collection,
@@ -89,52 +91,49 @@ from cognite_toolkit._cdf_tk.utils import (
     safe_read,
 )
 from cognite_toolkit._cdf_tk.utils.cdf import read_auth, try_find_error
+from cognite_toolkit._cdf_tk.utils.collection import chunker
 from cognite_toolkit._cdf_tk.utils.diff_list import diff_list_hashable
 
-from .auth_loaders import GroupAllScopedLoader
-from .data_organization_loaders import DataSetsLoader
-from .datamodel_loaders import DataModelLoader, SpaceLoader, ViewLoader
-from .group_scoped_loader import GroupResourceScopedLoader
-from .raw_loaders import RawDatabaseLoader, RawTableLoader
+from .auth import GroupAllScopedCRUD
+from .data_organization import DataSetsCRUD
+from .datamodel import DataModelCRUD, SpaceCRUD, ViewCRUD
+from .group_scoped import GroupResourceScopedCRUD
+from .raw import RawDatabaseCRUD, RawTableCRUD
 
 
 @final
-class TransformationLoader(
-    ResourceLoader[str, TransformationWrite, Transformation, TransformationWriteList, TransformationList]
-):
+class TransformationCRUD(ResourceCRUD[str, TransformationWrite, Transformation]):
     folder_name = "transformations"
-    filename_pattern = (
-        # Matches all yaml files except file names whose stem contain *.schedule. or .Notification
-        r"^(?!.*schedule.*|.*\.notification$).*$"
-    )
     resource_cls = Transformation
     resource_write_cls = TransformationWrite
-    list_cls = TransformationList
-    list_write_cls = TransformationWriteList
     kind = "Transformation"
     yaml_cls = TransformationYAML
     dependencies = frozenset(
         {
-            DataSetsLoader,
-            RawDatabaseLoader,
-            GroupAllScopedLoader,
-            SpaceLoader,
-            ViewLoader,
-            DataModelLoader,
-            RawTableLoader,
-            RawDatabaseLoader,
-            GroupResourceScopedLoader,
+            DataSetsCRUD,
+            RawDatabaseCRUD,
+            GroupAllScopedCRUD,
+            SpaceCRUD,
+            ViewCRUD,
+            DataModelCRUD,
+            RawTableCRUD,
+            RawDatabaseCRUD,
+            GroupResourceScopedCRUD,
         }
     )
     _doc_url = "Transformations/operation/createTransformations"
     _hash_key = "-- cdf-auth"
 
+    # The API supports creating/updating 1000 transformations in a single request,
+    # however, when the transformation has credentials, the session API times out on
+    # larger number of transformations. Thus, we use a conservative batch size.
+    _BATCH_SIZE = 20  # The maximum number of transformations to create in a single batch
+
     def __init__(self, client: ToolkitClient, build_dir: Path | None, console: Console | None = None):
         super().__init__(client, build_dir, console)
         self._authentication_by_id_operation: dict[
             tuple[str, Literal["read", "write"]], OidcCredentials | ClientCredentials
         ] = {}
-        self._nonce_cache: dict[str, NonceCredentials] = {}
 
     @property
     def display_name(self) -> str:
@@ -157,10 +156,7 @@ class TransformationLoader(
             if data_set_ids := {item.data_set_id for item in items if item.data_set_id}:
                 scope = TransformationsAcl.Scope.DataSet(list(data_set_ids))
 
-        return TransformationsAcl(
-            actions,
-            scope,  # type: ignore[arg-type]
-        )
+        return TransformationsAcl(actions, scope)
 
     @classmethod
     def get_id(cls, item: Transformation | TransformationWrite | dict) -> str:
@@ -183,28 +179,28 @@ class TransformationLoader(
         return {"externalId": id}
 
     @classmethod
-    def get_dependent_items(cls, item: dict) -> Iterable[tuple[type[ResourceLoader], Hashable]]:
+    def get_dependent_items(cls, item: dict) -> Iterable[tuple[type[ResourceCRUD], Hashable]]:
         if "dataSetExternalId" in item:
-            yield DataSetsLoader, item["dataSetExternalId"]
+            yield DataSetsCRUD, item["dataSetExternalId"]
         if destination := item.get("destination", {}):
             if not isinstance(destination, dict):
                 return
             if destination.get("type") == "raw" and in_dict(("database", "table"), destination):
-                yield RawDatabaseLoader, RawDatabase(destination["database"])
-                yield RawTableLoader, RawTable(destination["database"], destination["table"])
+                yield RawDatabaseCRUD, RawDatabase(destination["database"])
+                yield RawTableCRUD, RawTable(destination["database"], destination["table"])
             elif destination.get("type") in ("nodes", "edges") and (view := destination.get("view", {})):
                 if space := destination.get("instanceSpace"):
-                    yield SpaceLoader, space
+                    yield SpaceCRUD, space
                 if in_dict(("space", "externalId", "version"), view):
                     view["version"] = str(view["version"])
-                    yield ViewLoader, ViewId.load(view)
+                    yield ViewCRUD, ViewId.load(view)
             elif destination.get("type") == "instances":
                 if space := destination.get("instanceSpace"):
-                    yield SpaceLoader, space
+                    yield SpaceCRUD, space
                 if data_model := destination.get("dataModel"):
                     if in_dict(("space", "externalId", "version"), data_model):
                         data_model["version"] = str(data_model["version"])
-                        yield DataModelLoader, DataModelId.load(data_model)
+                        yield DataModelCRUD, DataModelId.load(data_model)
 
     def safe_read(self, filepath: Path | str) -> str:
         # If the destination is a DataModel or a View we need to ensure that the version is a string
@@ -231,13 +227,13 @@ class TransformationLoader(
             if query_file is None and "query" not in item:
                 if filepath is None:
                     raise ValueError("filepath must be set if query is not set")
-                raise ToolkitYAMLFormatError(
-                    f"query property or is missing. It can be inline or a separate file named {filepath.stem}.sql or {external_id}.sql",
-                    filepath,
+                warning = HighSeverityWarning(
+                    f"query property or is missing in {filepath.as_posix()!r}. It can be inline or a separate file named {filepath.stem}.sql or {external_id}.sql",
                 )
+                warning.print_warning(console=self.console)
             elif query_file and not query_file.exists():
                 # We checked above that filepath is not None
-                raise ToolkitFileNotFoundError(f"Query file {query_file.as_posix()} not found", filepath)  # type: ignore[union-attr]
+                raise ToolkitFileNotFoundError(f"Query file {query_file.as_posix()} not found", filepath)
             elif query_file and "query" in item:
                 raise ToolkitYAMLFormatError(
                     f"query property is ambiguously defined in both the yaml file and a separate file named {query_file}\n"
@@ -349,6 +345,16 @@ class TransformationLoader(
         if "authentication" in local:
             # The hash added to the beginning of the query detects the change in the authentication
             dumped["authentication"] = local["authentication"]
+        cdf_destination = dumped.get("destination", {})
+        local_destination = local.get("destination", {})
+        if isinstance(cdf_destination, dict) and isinstance(local_destination, dict):
+            if cdf_destination.get("instanceSpace") is None and "instanceSpace" not in local_destination:
+                cdf_destination.pop("instanceSpace", None)
+        if not dumped.get("query") and "query" not in local:
+            dumped.pop("query", None)
+        if dumped.get("conflictMode") == "upsert" and "conflictMode" not in local:
+            # Default set from server side.
+            dumped.pop("conflictMode", None)
         return dumped
 
     def split_resource(
@@ -367,17 +373,7 @@ class TransformationLoader(
         return super().diff_list(local, cdf, json_path)
 
     def create(self, items: Sequence[TransformationWrite]) -> TransformationList:
-        with warnings.catch_warnings():
-            warnings.simplefilter("ignore")
-            # Ignoring warnings from SDK about session unauthorized. Motivation is CDF is not fast enough to
-            # handle first a group that authorizes the session and then the transformation.
-            self._update_nonce(items)
-        try:
-            return self.client.transformations.create(items)
-        except CogniteAuthError as e:
-            if error := self._create_auth_creation_error(items):
-                raise error from e
-            raise e
+        return self._execute_in_batches(items, self.client.transformations.create)
 
     def retrieve(self, ids: SequenceNotStr[str | int]) -> TransformationList:
         internal_ids, external_ids = self._split_ids(ids)
@@ -386,17 +382,10 @@ class TransformationLoader(
         )
 
     def update(self, items: Sequence[TransformationWrite]) -> TransformationList:
-        with warnings.catch_warnings():
-            warnings.simplefilter("ignore")
-            # Ignoring warnings from SDK about session unauthorized. Motivation is CDF is not fast enough to
-            # handle first a group that authorizes the session and then the transformation.
-            self._update_nonce(items)
-            try:
-                return self.client.transformations.update(items, mode="replace")
-            except CogniteAuthError as e:
-                if error := self._create_auth_creation_error(items):
-                    raise error from e
-                raise e
+        def update(transformations: Sequence[TransformationWrite]) -> TransformationList:
+            return self.client.transformations.update(transformations, mode="replace")
+
+        return self._execute_in_batches(items, update)
 
     @staticmethod
     def _create_auth_creation_error(items: Sequence[TransformationWrite]) -> ResourceCreationError | None:
@@ -423,19 +412,89 @@ class TransformationLoader(
             self.client.transformations.delete(id=existing, ignore_unknown_ids=True)
         return len(existing)
 
+    def _execute_in_batches(
+        self,
+        items: Sequence[TransformationWrite],
+        api_call: Callable[[Sequence[TransformationWrite]], TransformationList],
+    ) -> TransformationList:
+        results = TransformationList([])
+        for chunk in chunker(items, self._BATCH_SIZE):
+            with warnings.catch_warnings():
+                warnings.simplefilter("ignore")
+                self._update_nonce(chunk)
+            try:
+                chunk_results = api_call(chunk)
+            except CogniteAuthError as e:
+                if error := self._create_auth_creation_error(chunk):
+                    raise error from e
+                raise e
+            except CogniteAPIError as e:
+                if "Failed to bind session using nonce" in e.message and len(chunk) > 1:
+                    results.extend(self._execute_one_by_one(chunk, api_call))
+                else:
+                    raise
+            else:
+                results.extend(chunk_results)
+        return results
+
+    def _execute_one_by_one(
+        self,
+        chunk: Sequence[TransformationWrite],
+        api_call: Callable[[Sequence[TransformationWrite]], TransformationList],
+    ) -> TransformationList:
+        MediumSeverityWarning(
+            f"Failed to create {len(chunk)} transformations in a batch due to nonce binding error. "
+            "Trying to recover by creating them one by one."
+        ).print_warning(console=self.client.console)
+        # Retry one by one
+        failed_ids: list[str] = []
+        success_count = 0
+        delay = 0.3
+        self._sleep_with_jitter(delay, delay + 0.3)
+        results = TransformationList([])
+        for item in chunk:
+            try:
+                recovered = api_call([item])
+            except CogniteAPIError as e:
+                if "Failed to bind session using nonce" in e.message:
+                    failed_ids.append(item.external_id or "<missing>")
+                    self._sleep_with_jitter(delay, delay + 0.3)
+                else:
+                    raise
+            else:
+                results.extend(recovered)
+                success_count += 1
+        message = f"  [bold]RECOVERY COMPLETE:[/] Successfully created {success_count:,} transformations"
+        if failed_ids:
+            message += f", failed to create {len(failed_ids):,} transformations: {humanize_collection(failed_ids)}"
+        else:
+            message += "."
+        if failed_ids:
+            HighSeverityWarning(message).print_warning(include_timestamp=True, console=self.client.console)
+        else:
+            self.client.console.print(message)
+        return results
+
+    @staticmethod
+    def _sleep_with_jitter(base_delay: float, max_delay: float) -> None:
+        """Sleeps for a random duration between base_delay and max_delay (inclusive)."""
+        sleep_time = random.uniform(base_delay, max_delay)
+        time.sleep(sleep_time)
+
     def _update_nonce(self, items: Sequence[TransformationWrite]) -> None:
         for item in items:
             if not item.external_id:
                 raise ToolkitRequiredValueError("Transformation must have external_id set.")
-            if read_credentials := self._authentication_by_id_operation.get((item.external_id, "read")):
+            if item.source_nonce is None and (
+                read_credentials := self._authentication_by_id_operation.get((item.external_id, "read"))
+            ):
                 item.source_nonce = self._create_nonce(read_credentials)
-            if write_credentials := self._authentication_by_id_operation.get((item.external_id, "write")):
+            if item.destination_nonce is None and (
+                write_credentials := self._authentication_by_id_operation.get((item.external_id, "write"))
+            ):
                 item.destination_nonce = self._create_nonce(write_credentials)
 
     def _create_nonce(self, credentials: OidcCredentials | ClientCredentials) -> NonceCredentials:
-        key = calculate_secure_hash(credentials.dump(), shorten=True)
-        if key in self._nonce_cache:
-            return self._nonce_cache[key]
         if isinstance(credentials, ClientCredentials):
             session = self.client.iam.sessions.create(credentials)
             nonce = NonceCredentials(session.id, session.nonce, self.client.config.project)
@@ -448,7 +507,6 @@ class TransformationLoader(
             nonce = NonceCredentials(session.id, session.nonce, credentials.cdf_project_name)
         else:
             raise ValueError(f"Error in TransformationLoader: {type(credentials)} is not a valid credentials type")
-        self._nonce_cache[key] = nonce
         return nonce
 
     def _iterate(
@@ -461,43 +519,6 @@ class TransformationLoader(
             self.client.transformations(data_set_external_ids=[data_set_external_id] if data_set_external_id else None)
         )
 
-    @classmethod
-    @lru_cache(maxsize=1)
-    def get_write_cls_parameter_spec(cls) -> ParameterSpecSet:
-        spec = super().get_write_cls_parameter_spec()
-        spec.update(
-            ParameterSpecSet(
-                {
-                    # Added by toolkit
-                    ParameterSpec(("dataSetExternalId",), frozenset({"str"}), is_required=False, _is_nullable=False),
-                    ParameterSpec(("authentication",), frozenset({"dict"}), is_required=False, _is_nullable=False),
-                    ParameterSpec(
-                        ("authentication", "clientId"), frozenset({"str"}), is_required=True, _is_nullable=False
-                    ),
-                    ParameterSpec(
-                        ("authentication", "clientSecret"), frozenset({"str"}), is_required=True, _is_nullable=False
-                    ),
-                    ParameterSpec(
-                        ("authentication", "scopes"), frozenset({"str"}), is_required=False, _is_nullable=False
-                    ),
-                    ParameterSpec(
-                        ("authentication", "scopes", ANY_INT), frozenset({"str"}), is_required=False, _is_nullable=False
-                    ),
-                    ParameterSpec(
-                        ("authentication", "tokenUri"), frozenset({"str"}), is_required=True, _is_nullable=False
-                    ),
-                    ParameterSpec(
-                        ("authentication", "cdfProjectName"), frozenset({"str"}), is_required=True, _is_nullable=False
-                    ),
-                    ParameterSpec(
-                        ("authentication", "audience"), frozenset({"str"}), is_required=False, _is_nullable=False
-                    ),
-                    ParameterSpec(("queryFile",), frozenset({"str"}), is_required=False, _is_nullable=False),
-                }
-            )
-        )
-        return spec
-
     def sensitive_strings(self, item: TransformationWrite) -> Iterable[str]:
         if item.source_oidc_credentials:
             yield item.source_oidc_credentials.client_secret
@@ -512,26 +533,21 @@ class TransformationLoader(
 
 
 @final
-class TransformationScheduleLoader(
-    ResourceLoader[
+class TransformationScheduleCRUD(
+    ResourceCRUD[
         str,
         TransformationScheduleWrite,
         TransformationSchedule,
-        TransformationScheduleWriteList,
-        TransformationScheduleList,
     ]
 ):
     folder_name = "transformations"
-    # Matches all yaml files whose stem contains *schedule or *TransformationSchedule.
-    filename_pattern = r"^.*schedule$"
     resource_cls = TransformationSchedule
     resource_write_cls = TransformationScheduleWrite
-    list_cls = TransformationScheduleList
-    list_write_cls = TransformationScheduleWriteList
     kind = "Schedule"
-    dependencies = frozenset({TransformationLoader})
+    yaml_cls = TransformationScheduleYAML
+    dependencies = frozenset({TransformationCRUD})
     _doc_url = "Transformation-Schedules/operation/createTransformationSchedules"
-    parent_resource = frozenset({TransformationLoader})
+    parent_resource = frozenset({TransformationCRUD})
 
     @property
     def display_name(self) -> str:
@@ -558,9 +574,9 @@ class TransformationScheduleLoader(
         return {"externalId": id}
 
     @classmethod
-    def get_dependent_items(cls, item: dict) -> Iterable[tuple[type[ResourceLoader], Hashable]]:
+    def get_dependent_items(cls, item: dict) -> Iterable[tuple[type[ResourceCRUD], Hashable]]:
         if "externalId" in item:
-            yield TransformationLoader, item["externalId"]
+            yield TransformationCRUD, item["externalId"]
 
     def create(self, items: Sequence[TransformationScheduleWrite]) -> TransformationScheduleList:
         try:
@@ -576,7 +592,7 @@ class TransformationScheduleLoader(
     def retrieve(self, ids: SequenceNotStr[str]) -> TransformationScheduleList:
         return self.client.transformations.schedules.retrieve_multiple(external_ids=ids, ignore_unknown_ids=True)
 
-    def update(self, items: TransformationScheduleWriteList) -> TransformationScheduleList:
+    def update(self, items: Sequence[TransformationScheduleWrite]) -> TransformationScheduleList:
         return self.client.transformations.schedules.update(items, mode="replace")
 
     def delete(self, ids: str | SequenceNotStr[str] | None) -> int:
@@ -609,27 +625,22 @@ class TransformationScheduleLoader(
 
 
 @final
-class TransformationNotificationLoader(
-    ResourceLoader[
+class TransformationNotificationCRUD(
+    ResourceCRUD[
         str,
         TransformationNotificationWrite,
         TransformationNotification,
-        TransformationNotificationWriteList,
-        TransformationNotificationList,
     ]
 ):
     folder_name = "transformations"
-    # Matches all yaml files whose stem ends with *Notification.
-    filename_pattern = r"^.*Notification$"
     resource_cls = TransformationNotification
     resource_write_cls = TransformationNotificationWrite
-    list_cls = TransformationNotificationList
-    list_write_cls = TransformationNotificationWriteList
     kind = "Notification"
-    dependencies = frozenset({TransformationLoader})
+    dependencies = frozenset({TransformationCRUD})
     _doc_url = "Transformation-Notifications/operation/createTransformationNotifications"
     _split_character = "@@@"
-    parent_resource = frozenset({TransformationLoader})
+    parent_resource = frozenset({TransformationCRUD})
+    yaml_cls = TransformationNotificationYAML
 
     @property
     def display_name(self) -> str:
@@ -670,8 +681,8 @@ class TransformationNotificationLoader(
             dumped["transformationExternalId"] = local["transformationExternalId"]
         return dumped
 
-    def create(self, items: TransformationNotificationWriteList) -> TransformationNotificationList:
-        return self.client.transformations.notifications.create(items)  # type: ignore[return-value]
+    def create(self, items: Sequence[TransformationNotificationWrite]) -> TransformationNotificationList:
+        return self.client.transformations.notifications.create(items)
 
     def retrieve(self, ids: SequenceNotStr[str]) -> TransformationNotificationList:
         retrieved = TransformationNotificationList([])
@@ -701,7 +712,7 @@ class TransformationNotificationLoader(
 
         return retrieved
 
-    def update(self, items: TransformationNotificationWriteList) -> TransformationNotificationList:
+    def update(self, items: Sequence[TransformationNotificationWrite]) -> TransformationNotificationList:
         # Note that since a notification is identified by the combination of transformationExternalId and destination,
         # which is the entire object, an update should never happen. However, implementing just in case.
         item_by_id = {self.get_id(item): item for item in items}
@@ -718,7 +729,7 @@ class TransformationNotificationLoader(
             else:
                 create.append(local_item)
             if existing_item:
-                delete.append(cast(int, existing_item.id))
+                delete.append(existing_item.id)
         if delete:
             self.client.transformations.notifications.delete(delete)
         updated_by_id: dict[str, TransformationNotification] = {}
@@ -736,7 +747,7 @@ class TransformationNotificationLoader(
         # while the toolkit uses the transformationExternalId + destination as the id. Thus, there could
         # be multiple notifications for the same transformationExternalId + destination.
         if existing := self.retrieve(ids):
-            self.client.transformations.notifications.delete([item.id for item in existing])  # type: ignore[misc]
+            self.client.transformations.notifications.delete([item.id for item in existing])
         return len(existing)
 
     def _iterate(
@@ -755,11 +766,11 @@ class TransformationNotificationLoader(
                     yield from self.client.transformations.notifications(transformation_id=transformation_id)
 
     @classmethod
-    def get_dependent_items(cls, item: dict) -> Iterable[tuple[type[ResourceLoader], Hashable]]:
+    def get_dependent_items(cls, item: dict) -> Iterable[tuple[type[ResourceCRUD], Hashable]]:
         """Returns all items that this item requires.
 
         For example, a TimeSeries requires a DataSet, so this method would return the
         DatasetLoader and identifier of that dataset.
         """
         if "transformationExternalId" in item:
-            yield TransformationLoader, item["transformationExternalId"]
+            yield TransformationCRUD, item["transformationExternalId"]