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

cognite_toolkit/_cdf_tk/{loaders/_resource_loaders/function_loaders.py → cruds/_resource_cruds/function.py}

@@ -1,9 +1,9 @@
 import time
 from collections import defaultdict
 from collections.abc import Hashable, Iterable, Sequence
-from functools import lru_cache
+from functools import cached_property
 from pathlib import Path
-from typing import Any, cast, final
+from typing import Any, Literal, cast, final
 
 from cognite.client.data_classes import (
     ClientCredentials,
@@ -12,9 +12,7 @@ from cognite.client.data_classes import (
     FunctionSchedule,
     FunctionSchedulesList,
     FunctionScheduleWrite,
-    FunctionScheduleWriteList,
     FunctionWrite,
-    FunctionWriteList,
 )
 from cognite.client.data_classes.capabilities import (
     AllScope,
@@ -24,49 +22,46 @@ from cognite.client.data_classes.capabilities import (
     FunctionsAcl,
     SessionsAcl,
 )
+from cognite.client.data_classes.data_modeling.cdm.v1 import CogniteFileApply
 from cognite.client.data_classes.functions import HANDLER_FILE_NAME
-from cognite.client.exceptions import CogniteAuthError
+from cognite.client.exceptions import CogniteAPIError
 from cognite.client.utils.useful_types import SequenceNotStr
 from rich import print
 from rich.console import Console
 
-from cognite_toolkit._cdf_tk._parameters import ParameterSpec, ParameterSpecSet
 from cognite_toolkit._cdf_tk.client import ToolkitClient
 from cognite_toolkit._cdf_tk.client.data_classes.functions import FunctionScheduleID
+from cognite_toolkit._cdf_tk.cruds._base_cruds import ResourceCRUD
 from cognite_toolkit._cdf_tk.exceptions import (
     ResourceCreationError,
     ToolkitRequiredValueError,
 )
-from cognite_toolkit._cdf_tk.loaders._base_loaders import ResourceLoader
 from cognite_toolkit._cdf_tk.resource_classes import FunctionScheduleYAML, FunctionsYAML
 from cognite_toolkit._cdf_tk.tk_warnings import HighSeverityWarning, LowSeverityWarning
 from cognite_toolkit._cdf_tk.utils import (
     calculate_directory_hash,
     calculate_hash,
     calculate_secure_hash,
+    humanize_collection,
 )
 from cognite_toolkit._cdf_tk.utils.cdf import read_auth, try_find_error
+from cognite_toolkit._cdf_tk.utils.file import create_temporary_zip, sanitize_filename
 from cognite_toolkit._cdf_tk.utils.text import suffix_description
 
-from .auth_loaders import GroupAllScopedLoader
-from .data_organization_loaders import DataSetsLoader
-from .group_scoped_loader import GroupResourceScopedLoader
+from .auth import GroupAllScopedCRUD
+from .data_organization import DataSetsCRUD
+from .group_scoped import GroupResourceScopedCRUD
 
 
 @final
-class FunctionLoader(ResourceLoader[str, FunctionWrite, Function, FunctionWriteList, FunctionList]):
+class FunctionCRUD(ResourceCRUD[str, FunctionWrite, Function]):
     support_drop = True
     folder_name = "functions"
-    filename_pattern = (
-        r"^(?:(?!schedule).)*$"  # Matches all yaml files except file names who's stem contain *.schedule.
-    )
     resource_cls = Function
     resource_write_cls = FunctionWrite
-    list_cls = FunctionList
-    list_write_cls = FunctionWriteList
     kind = "Function"
     yaml_cls = FunctionsYAML
-    dependencies = frozenset({DataSetsLoader, GroupAllScopedLoader})
+    dependencies = frozenset({DataSetsCRUD, GroupAllScopedCRUD})
     _doc_url = "Functions/operation/postFunctions"
     metadata_value_limit = 512
     support_update = False
@@ -78,8 +73,14 @@ class FunctionLoader(ResourceLoader[str, FunctionWrite, Function, FunctionWriteL
     def __init__(self, client: ToolkitClient, build_path: Path | None, console: Console | None):
         super().__init__(client, build_path, console)
         self.data_set_id_by_external_id: dict[str, int] = {}
+        self.space_by_external_id: dict[str, str] = {}
         self.function_dir_by_external_id: dict[str, Path] = {}
 
+    @cached_property
+    def project_data_modeling_status(self) -> Literal["HYBRID", "DATA_MODELING_ONLY"]:
+        status = self.client.project.status()
+        return status.this_project.data_modeling_status
+
     @property
     def display_name(self) -> str:
         return "functions"
@@ -114,9 +115,9 @@ class FunctionLoader(ResourceLoader[str, FunctionWrite, Function, FunctionWriteL
         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"]
 
     def load_resource_file(
         self, filepath: Path, environment_variables: dict[str, str | None] | None = None
@@ -205,6 +206,8 @@ class FunctionLoader(ResourceLoader[str, FunctionWrite, Function, FunctionWriteL
         item_id = self.get_id(resource)
         if ds_external_id := resource.pop("dataSetExternalId", None):
             self.data_set_id_by_external_id[item_id] = self.client.lookup.data_sets.id(ds_external_id, is_dry_run)
+        if space := resource.pop("space", None):
+            self.space_by_external_id[item_id] = space
         if "fileId" not in resource:
             # The fileID is required for the function to be created, but in the `.create` method
             # we first create that file and then set the fileID.
@@ -253,11 +256,22 @@ class FunctionLoader(ResourceLoader[str, FunctionWrite, Function, FunctionWriteL
             function_zip_file = self.client.files.retrieve(id=file_id)
             if function_zip_file and (data_set_id := function_zip_file.data_set_id):
                 dumped["dataSetExternalId"] = self.client.lookup.data_sets.external_id(data_set_id)
+            if function_zip_file and function_zip_file.instance_id is not None:
+                dumped["space"] = function_zip_file.instance_id.space
 
         if dumped.get("functionPath") == HANDLER_FILE_NAME and "functionPath" not in local:
             # Remove the default value of the functionPath
             dumped.pop("functionPath", None)
 
+        # The Function API removes none-like values when processing the request.
+        # Ref https://github.com/cognitedata/infrastructure/blob/7459a8f6959af89240c3250f9133c079a19d57c5/services/context/context_api/utils/schema_helpers.py#L8
+        # This ensures that if the user sets owner or description to None or empty string,
+        # it will not cause the comparison with the server side to evaluate as different,
+        # even though the function is unchanged.
+        for key in ["owner", "description"]:
+            if key in local and local[key] in {None, ""} and key not in dumped:
+                dumped[key] = local[key]
+
         return dumped
 
     def _is_activated(self, action: str) -> bool:
@@ -280,7 +294,7 @@ class FunctionLoader(ResourceLoader[str, FunctionWrite, Function, FunctionWriteL
             self.client.functions.activate()
         return False
 
-    def create(self, items: FunctionWriteList) -> FunctionList:
+    def create(self, items: Sequence[FunctionWrite]) -> FunctionList:
         created = FunctionList([], cognite_client=self.client)
         if not self._is_activated("create"):
             return created
@@ -288,18 +302,12 @@ class FunctionLoader(ResourceLoader[str, FunctionWrite, Function, FunctionWriteL
             raise ValueError("build_path must be set to compare functions as function code must be compared.")
         for item in items:
             external_id = item.external_id or item.name
-            function_rootdir = self.function_dir_by_external_id[external_id]
-            file_id = self.client.functions._zip_and_upload_folder(
-                name=item.name,
-                folder=str(function_rootdir),
-                external_id=external_id,
-                data_set_id=self.data_set_id_by_external_id.get(external_id),
-            )
+            file_id = self._upload_function_code(external_id, item)
             # Wait until the files is available
             t0 = time.perf_counter()
             sleep_time = 1.0  # seconds
             for i in range(6):
-                file = self.client.files.retrieve(external_id=external_id)
+                file = self.client.files.retrieve(id=file_id)
                 if file and file.uploaded:
                     break
                 time.sleep(sleep_time)
@@ -312,11 +320,55 @@ class FunctionLoader(ResourceLoader[str, FunctionWrite, Function, FunctionWriteL
                     " problem persists, please contact Cognite support."
                 )
             item.file_id = file_id
-            created_item = self.client.functions.create(item)
+            created_item = self.client.functions.create_with_429_retry(item)
             self._warn_if_cpu_or_memory_changed(created_item, item)
             created.append(created_item)
         return created
 
+    def _upload_function_code(self, external_id: str, item: FunctionWrite) -> int:
+        """Uploads the function code to CDF.
+
+        It will either upload the code to a CogniteFile if a space is provided and the feature flag is enabled,
+        or to a classic file if not.
+
+        Args:
+            external_id (str): The external ID of the function.
+            item (FunctionWrite): The function write object.
+
+        Returns:
+            int: The ID of the created file metadata object.
+        """
+        function_rootdir = self.function_dir_by_external_id[external_id]
+        data_set_id = self.data_set_id_by_external_id.get(external_id)
+        space = self.space_by_external_id.get(external_id)
+        if space is None and self.project_data_modeling_status.casefold() == "data_modeling_only":
+            raise ResourceCreationError(
+                f"Function {external_id!r} must have a space set when the project is in DATA_MODELING_ONLY mode. "
+                "This is used to set the NodeId of the CogniteFile created for the function code."
+            )
+
+        with create_temporary_zip(function_rootdir, "function.zip") as zip_path:
+            if space:
+                cognite_file = CogniteFileApply(
+                    space=space,
+                    external_id=external_id,
+                    name=f"{sanitize_filename(item.name)}.zip",
+                    mime_type="application/zip",
+                )
+                _ = self.client.data_modeling.instances.apply(cognite_file, replace=True)
+                upload_file = self.client.files.upload_content_bytes(
+                    zip_path.read_bytes(), instance_id=cognite_file.as_id()
+                )
+            else:
+                upload_file = self.client.files.upload_bytes(
+                    zip_path.read_bytes(),
+                    name=f"{sanitize_filename(item.name)}.zip",
+                    external_id=external_id,
+                    overwrite=True,
+                    data_set_id=data_set_id,
+                )
+        return upload_file.id
+
     @staticmethod
     def _warn_if_cpu_or_memory_changed(created_item: Function, item: FunctionWrite) -> None:
         is_cpu_increased = (
@@ -352,7 +404,7 @@ class FunctionLoader(ResourceLoader[str, FunctionWrite, Function, FunctionWriteL
     def delete(self, ids: SequenceNotStr[str]) -> int:
         functions = self.retrieve(ids)
 
-        self.client.functions.delete(external_id=ids)
+        self.client.functions.delete_with_429_retry(external_id=ids, ignore_unknown_ids=True)
         file_ids = {func.file_id for func in functions if func.file_id}
         self.client.files.delete(id=list(file_ids), ignore_unknown_ids=True)
         return len(ids)
@@ -365,34 +417,17 @@ class FunctionLoader(ResourceLoader[str, FunctionWrite, Function, FunctionWriteL
     ) -> Iterable[Function]:
         return iter(self.client.functions)
 
-    @classmethod
-    @lru_cache(maxsize=1)
-    def get_write_cls_parameter_spec(cls) -> ParameterSpecSet:
-        spec = super().get_write_cls_parameter_spec()
-        # Added by toolkit
-        spec.add(ParameterSpec(("dataSetExternalId",), frozenset({"str"}), is_required=False, _is_nullable=False))
-        # Replaced by the toolkit
-        spec.discard(ParameterSpec(("fileId",), frozenset({"int"}), is_required=True, _is_nullable=False))
-        return spec
-
 
 @final
-class FunctionScheduleLoader(
-    ResourceLoader[
-        FunctionScheduleID, FunctionScheduleWrite, FunctionSchedule, FunctionScheduleWriteList, FunctionSchedulesList
-    ]
-):
+class FunctionScheduleCRUD(ResourceCRUD[FunctionScheduleID, FunctionScheduleWrite, FunctionSchedule]):
     folder_name = "functions"
-    filename_pattern = r"^.*schedule.*$"  # Matches all yaml files who's stem contain *.schedule
     resource_cls = FunctionSchedule
     resource_write_cls = FunctionScheduleWrite
-    list_cls = FunctionSchedulesList
-    list_write_cls = FunctionScheduleWriteList
     kind = "Schedule"
     yaml_cls = FunctionScheduleYAML
-    dependencies = frozenset({FunctionLoader, GroupResourceScopedLoader, GroupAllScopedLoader})
+    dependencies = frozenset({FunctionCRUD, GroupResourceScopedCRUD, GroupAllScopedCRUD})
     _doc_url = "Function-schedules/operation/postFunctionSchedules"
-    parent_resource = frozenset({FunctionLoader})
+    parent_resource = frozenset({FunctionCRUD})
     support_update = False
 
     _hash_key = "cdf-auth"
@@ -416,16 +451,12 @@ class FunctionScheduleLoader(
         function_actions = (
             [FunctionsAcl.Action.Read] if read_only else [FunctionsAcl.Action.Read, FunctionsAcl.Action.Write]
         )
-        session_actions = (
-            [SessionsAcl.Action.List]
-            if read_only
-            else [SessionsAcl.Action.List, SessionsAcl.Action.Create, SessionsAcl.Action.Delete]
-        )
-
-        return [
+        required_capabilities: list[Capability] = [
             FunctionsAcl(function_actions, FunctionsAcl.Scope.All()),
-            SessionsAcl(session_actions, SessionsAcl.Scope.All()),
         ]
+        if not read_only:
+            required_capabilities.append(SessionsAcl([SessionsAcl.Action.Create], SessionsAcl.Scope.All()))
+        return required_capabilities
 
     @classmethod
     def dump_id(cls, id: FunctionScheduleID) -> dict[str, Any]:
@@ -444,9 +475,9 @@ class FunctionScheduleLoader(
         return FunctionScheduleID(item.function_external_id, item.name)
 
     @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 "functionExternalId" in item:
-            yield FunctionLoader, item["functionExternalId"]
+            yield FunctionCRUD, item["functionExternalId"]
 
     def load_resource_file(
         self, filepath: Path, environment_variables: dict[str, str | None] | None = None
@@ -496,21 +527,11 @@ class FunctionScheduleLoader(
             dumped["authentication"] = local["authentication"]
         return dumped
 
-    def _resolve_functions_ext_id(self, items: FunctionScheduleWriteList) -> FunctionScheduleWriteList:
-        functions = FunctionLoader(self.client, None, None).retrieve(
-            list(set([item.function_external_id for item in items]))
-        )
-        for item in items:
-            for func in functions:
-                if func.external_id == item.function_external_id:
-                    item.function_id = func.id  # type: ignore[assignment]
-        return items
-
     def retrieve(self, ids: SequenceNotStr[FunctionScheduleID]) -> FunctionSchedulesList:
         names_by_function: dict[str, set[str]] = defaultdict(set)
         for id_ in ids:
             names_by_function[id_.function_external_id].add(id_.name)
-        functions = FunctionLoader(self.client, None, None).retrieve(list(names_by_function))
+        functions = FunctionCRUD(self.client, None, None).retrieve(list(names_by_function))
         schedules = FunctionSchedulesList([])
         for func in functions:
             func_external_id = cast(str, func.external_id)
@@ -522,27 +543,57 @@ class FunctionScheduleLoader(
             )
         return schedules
 
-    def create(self, items: FunctionScheduleWriteList) -> FunctionSchedulesList:
+    def create(self, items: Sequence[FunctionScheduleWrite]) -> FunctionSchedulesList:
         created_list = FunctionSchedulesList([], cognite_client=self.client)
+        function_id_by_external_id = self._get_function_ids_by_external_id(items)
+
         for item in items:
             id_ = self.get_id(item)
             if id_ not in self.authentication_by_id:
                 raise ToolkitRequiredValueError(f"Authentication is missing for schedule {id_!r}")
             client_credentials = self.authentication_by_id[id_]
             try:
-                created = self.client.functions.schedules.create(item, client_credentials=client_credentials)
-            except CogniteAuthError as e:
+                session = self.client.iam.sessions.create(client_credentials, session_type="CLIENT_CREDENTIALS")
+            except CogniteAPIError as e:
                 if hint := try_find_error(client_credentials):
                     raise ResourceCreationError(f"Failed to create Function Schedule {id_}: {hint}") from e
                 raise e
 
-            # The PySDK mutates the input object, such that function_id is set and function_external_id is None.
-            # If we call .get_id on the returned object, it will raise an error we require the function_external_id
-            # to be set.
+            # Serialization to create a copy as we are mutating the object.
+            to_create = FunctionScheduleWrite._load(item.dump())
+            to_create.nonce = session.nonce
+            to_create.function_id = function_id_by_external_id[id_.function_external_id]
+            to_create.function_external_id = None
+
+            created = self.client.functions.schedules.create(to_create)
+
             created.function_external_id = id_.function_external_id
+
             created_list.append(created)
         return created_list
 
+    def _get_function_ids_by_external_id(self, items: Sequence[FunctionScheduleWrite]) -> dict[str, int]:
+        functions_to_lookup = list({item.function_external_id for item in items if item.function_external_id})
+        if not functions_to_lookup:
+            return {}
+        function_ids = self.client.lookup.functions.id(functions_to_lookup)
+
+        if len(function_ids) == len(functions_to_lookup):
+            return dict(zip(functions_to_lookup, function_ids))
+        # The lookup API is cached, so it is cheap to do individual lookups to find the missing ones.
+        lookup_pair = (
+            (function_external_id, self.client.lookup.functions.id(function_external_id))
+            for function_external_id in functions_to_lookup
+        )
+        missing_functions = {func for func, func_id in lookup_pair if func_id is None}
+        failed_schedules = [self.get_id(item) for item in items if item.function_external_id in missing_functions]
+        plural_schedules = "s" if len(failed_schedules) > 1 else ""
+        plural_fun = "s" if len(missing_functions) > 1 else ""
+        raise ResourceCreationError(
+            f"Failed to create function schedule{plural_schedules} {humanize_collection(failed_schedules)}. "
+            f"Could not find function{plural_fun} {humanize_collection(missing_functions)!r}"
+        )
+
     def delete(self, ids: SequenceNotStr[FunctionScheduleID]) -> int:
         schedules = self.retrieve(ids)
         count = 0
@@ -566,20 +617,6 @@ class FunctionScheduleLoader(
                     continue
                 yield from self.client.functions.schedules(function_external_id=parent_id)
 
-    @classmethod
-    @lru_cache(maxsize=1)
-    def get_write_cls_parameter_spec(cls) -> ParameterSpecSet:
-        spec = super().get_write_cls_parameter_spec()
-        # Added by toolkit
-        spec.add(ParameterSpec(("authentication",), frozenset({"dict"}), is_required=False, _is_nullable=False))
-        spec.add(
-            ParameterSpec(("authentication", "clientId"), frozenset({"str"}), is_required=False, _is_nullable=False)
-        )
-        spec.add(
-            ParameterSpec(("authentication", "clientSecret"), frozenset({"str"}), is_required=False, _is_nullable=False)
-        )
-        return spec
-
     def sensitive_strings(self, item: FunctionScheduleWrite) -> Iterable[str]:
         id_ = self.get_id(item)
         if id_ in self.authentication_by_id:

cognite_toolkit/_cdf_tk/{loaders/_resource_loaders/group_scoped_loader.py → cruds/_resource_cruds/group_scoped.py}

@@ -20,29 +20,29 @@ from rich.console import Console
 
 from cognite_toolkit._cdf_tk.client import ToolkitClient
 
-from .auth_loaders import GroupLoader, SecurityCategoryLoader
-from .classic_loaders import AssetLoader
-from .data_organization_loaders import DataSetsLoader
-from .datamodel_loaders import SpaceLoader
-from .extraction_pipeline_loaders import ExtractionPipelineLoader
-from .location_loaders import LocationFilterLoader
-from .raw_loaders import RawDatabaseLoader, RawTableLoader
-from .timeseries_loaders import TimeSeriesLoader
+from .auth import GroupCRUD, SecurityCategoryCRUD
+from .classic import AssetCRUD
+from .data_organization import DataSetsCRUD
+from .datamodel import SpaceCRUD
+from .extraction_pipeline import ExtractionPipelineCRUD
+from .location import LocationFilterCRUD
+from .raw import RawDatabaseCRUD, RawTableCRUD
+from .timeseries import TimeSeriesCRUD
 
 
 @final
-class GroupResourceScopedLoader(GroupLoader):
+class GroupResourceScopedCRUD(GroupCRUD):
     dependencies = frozenset(
         {
-            SpaceLoader,
-            DataSetsLoader,
-            ExtractionPipelineLoader,
-            TimeSeriesLoader,
-            SecurityCategoryLoader,
-            LocationFilterLoader,
-            AssetLoader,
-            RawDatabaseLoader,
-            RawTableLoader,
+            SpaceCRUD,
+            DataSetsCRUD,
+            ExtractionPipelineCRUD,
+            TimeSeriesCRUD,
+            SecurityCategoryCRUD,
+            LocationFilterCRUD,
+            AssetCRUD,
+            RawDatabaseCRUD,
+            RawTableCRUD,
         }
     )