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

cognite_toolkit/_cdf_tk/{loaders/_resource_loaders/extraction_pipeline_loaders.py → cruds/_resource_cruds/extraction_pipeline.py}

@@ -11,11 +11,9 @@
 # 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 re
+
 from collections.abc import Hashable, Iterable, Sequence
-from functools import lru_cache
 from pathlib import Path
 from typing import Any, final
 
@@ -27,27 +25,23 @@ from cognite.client.data_classes import (
 )
 from cognite.client.data_classes.capabilities import (
     Capability,
-    ExtractionConfigsAcl,
     ExtractionPipelinesAcl,
 )
 from cognite.client.data_classes.extractionpipelines import (
     ExtractionPipelineConfigList,
     ExtractionPipelineConfigWrite,
-    ExtractionPipelineConfigWriteList,
     ExtractionPipelineWrite,
-    ExtractionPipelineWriteList,
 )
 from cognite.client.exceptions import CogniteAPIError, CogniteDuplicatedError, CogniteNotFoundError
 from cognite.client.utils.useful_types import SequenceNotStr
-from rich import print
 
-from cognite_toolkit._cdf_tk._parameters import ANYTHING, ParameterSpec, ParameterSpecSet
 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 (
     ToolkitRequiredValueError,
 )
-from cognite_toolkit._cdf_tk.loaders._base_loaders import ResourceLoader
+from cognite_toolkit._cdf_tk.resource_classes import ExtractionPipelineConfigYAML, ExtractionPipelineYAML
 from cognite_toolkit._cdf_tk.tk_warnings import (
     HighSeverityWarning,
 )
@@ -59,25 +53,19 @@ from cognite_toolkit._cdf_tk.utils import (
 )
 from cognite_toolkit._cdf_tk.utils.diff_list import diff_list_force_hashable, diff_list_identifiable
 
-from .auth_loaders import GroupAllScopedLoader
-from .data_organization_loaders import DataSetsLoader
-from .raw_loaders import RawDatabaseLoader, RawTableLoader
+from .auth import GroupAllScopedCRUD
+from .data_organization import DataSetsCRUD
+from .raw import RawDatabaseCRUD, RawTableCRUD
 
 
 @final
-class ExtractionPipelineLoader(
-    ResourceLoader[
-        str, ExtractionPipelineWrite, ExtractionPipeline, ExtractionPipelineWriteList, ExtractionPipelineList
-    ]
-):
+class ExtractionPipelineCRUD(ResourceCRUD[str, ExtractionPipelineWrite, ExtractionPipeline]):
     folder_name = "extraction_pipelines"
-    filename_pattern = r"^(?:(?!\.config).)*$"  # Matches all yaml files except file names who's stem contain *.config.
     resource_cls = ExtractionPipeline
     resource_write_cls = ExtractionPipelineWrite
-    list_cls = ExtractionPipelineList
-    list_write_cls = ExtractionPipelineWriteList
     kind = "ExtractionPipeline"
-    dependencies = frozenset({DataSetsLoader, RawDatabaseLoader, RawTableLoader, GroupAllScopedLoader})
+    dependencies = frozenset({DataSetsCRUD, RawDatabaseCRUD, RawTableCRUD, GroupAllScopedCRUD})
+    yaml_cls = ExtractionPipelineYAML
     _doc_url = "Extraction-Pipelines/operation/createExtPipes"
 
     @property
@@ -104,10 +92,7 @@ class ExtractionPipelineLoader(
             if data_set_id := {item.data_set_id for item in items if item.data_set_id}:
                 scope = ExtractionPipelinesAcl.Scope.DataSet(list(data_set_id))
 
-        return ExtractionPipelinesAcl(
-            actions,
-            scope,  # type: ignore[arg-type]
-        )
+        return ExtractionPipelinesAcl(actions, scope)
 
     @classmethod
     def get_id(cls, item: ExtractionPipeline | ExtractionPipelineWrite | dict) -> str:
@@ -130,18 +115,18 @@ class ExtractionPipelineLoader(
         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]]:
         seen_databases: set[str] = set()
         if "dataSetExternalId" in item:
-            yield DataSetsLoader, item["dataSetExternalId"]
+            yield DataSetsCRUD, item["dataSetExternalId"]
         if "rawTables" in item:
             for entry in item["rawTables"]:
                 if db := entry.get("dbName"):
                     if db not in seen_databases:
                         seen_databases.add(db)
-                        yield RawDatabaseLoader, RawDatabase(db_name=db)
+                        yield RawDatabaseCRUD, RawDatabase(db_name=db)
                     if "tableName" in entry:
-                        yield RawTableLoader, RawTable._load(entry)
+                        yield RawTableCRUD, RawTable._load(entry)
 
     def load_resource(self, resource: dict[str, Any], is_dry_run: bool = False) -> ExtractionPipelineWrite:
         if ds_external_id := resource.pop("dataSetExternalId", None):
@@ -185,7 +170,7 @@ class ExtractionPipelineLoader(
     def retrieve(self, ids: SequenceNotStr[str]) -> ExtractionPipelineList:
         return self.client.extraction_pipelines.retrieve_multiple(external_ids=ids, ignore_unknown_ids=True)
 
-    def update(self, items: ExtractionPipelineWriteList) -> ExtractionPipelineList:
+    def update(self, items: Sequence[ExtractionPipelineWrite]) -> ExtractionPipelineList:
         # Bug in SDK overload so need the ignore.
         return self.client.extraction_pipelines.update(items, mode="replace")  # type: ignore[call-overload]
 
@@ -219,36 +204,23 @@ class ExtractionPipelineLoader(
             if pipeline.data_set_id == data_set.id:
                 yield pipeline
 
-    @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))
-        # Set on deploy time by toolkit
-        spec.discard(ParameterSpec(("dataSetId",), frozenset({"int"}), is_required=True, _is_nullable=False))
-        return spec
-
 
 @final
-class ExtractionPipelineConfigLoader(
-    ResourceLoader[
+class ExtractionPipelineConfigCRUD(
+    ResourceCRUD[
         str,
         ExtractionPipelineConfigWrite,
         ExtractionPipelineConfig,
-        ExtractionPipelineConfigWriteList,
-        ExtractionPipelineConfigList,
     ]
 ):
     folder_name = "extraction_pipelines"
-    filename_pattern = r"^.*config$"
     resource_cls = ExtractionPipelineConfig
     resource_write_cls = ExtractionPipelineConfigWrite
-    list_cls = ExtractionPipelineConfigList
-    list_write_cls = ExtractionPipelineConfigWriteList
     kind = "Config"
-    dependencies = frozenset({ExtractionPipelineLoader})
+    dependencies = frozenset({ExtractionPipelineCRUD})
     _doc_url = "Extraction-Pipelines-Config/operation/createExtPipeConfig"
+    parent_resource = frozenset({ExtractionPipelineCRUD})
+    yaml_cls = ExtractionPipelineConfigYAML
 
     @property
     def display_name(self) -> str:
@@ -258,21 +230,8 @@ class ExtractionPipelineConfigLoader(
     def get_required_capability(
         cls, items: Sequence[ExtractionPipelineConfigWrite] | None, read_only: bool
     ) -> list[Capability]:
-        if not items and items is not None:
-            return []
-
-        actions = (
-            [ExtractionConfigsAcl.Action.Read]
-            if read_only
-            else [ExtractionConfigsAcl.Action.Read, ExtractionConfigsAcl.Action.Write]
-        )
-
-        return [
-            ExtractionConfigsAcl(
-                actions,
-                ExtractionConfigsAcl.Scope.All(),
-            )
-        ]
+        # We check the parent extraction pipeline permissions instead
+        return []
 
     @classmethod
     def get_id(cls, item: ExtractionPipelineConfig | ExtractionPipelineConfigWrite | dict) -> str:
@@ -287,9 +246,9 @@ class ExtractionPipelineConfigLoader(
         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 ExtractionPipelineLoader, item["externalId"]
+            yield ExtractionPipelineCRUD, item["externalId"]
 
     def safe_read(self, filepath: Path | str) -> str:
         # The config is expected to be a string that is parsed as a YAML on the server side.
@@ -323,21 +282,51 @@ class ExtractionPipelineConfigLoader(
     def load_resource(self, resource: dict[str, Any], is_dry_run: bool = False) -> ExtractionPipelineConfigWrite:
         config_raw = resource.get("config")
         if isinstance(config_raw, str):
-            # There might be keyvauls secrets in the config that would lead to parsing errors. The syntax
-            # for this is `connection-string: !keyvault secret`. This is not valid YAML, so we need to
-            # replace it with `connection-string: keyvault secret` to make it valid.
-            config_raw = re.sub(r": !(\w+)", r": \1", config_raw)
-            try:
-                read_yaml_content(config_raw)
-            except yaml.YAMLError as e:
-                print(
-                    HighSeverityWarning(
-                        f"Configuration for {resource.get('external_id', 'missing')} could not be parsed "
-                        f"as valid YAML, which is the recommended format. Error: {e}"
-                    ).get_message()
-                )
+            self._validate_config(config_raw, resource)
         return ExtractionPipelineConfigWrite._load(resource)
 
+    def _validate_config(self, config_raw: str, resource: dict[str, Any]) -> dict[str, Any]:
+        """The extraction pipeline API specifies config to be a string. However, it should be a valid YAML mapping,
+        so we validate that here and warn if it is not."""
+        try:
+            result = read_yaml_content(config_raw)
+        except yaml.YAMLError as e:
+            id_ = self._get_id(resource, default="missing")
+            HighSeverityWarning(
+                f"Configuration for {id_!r} could not be parsed "
+                f"as valid YAML, which is the recommended format. Error: {e}"
+            ).print_warning(console=self.console)
+        else:
+            if isinstance(result, dict):
+                return result
+            id_ = self._get_id(resource, default="missing")
+            HighSeverityWarning(
+                f"Configuration for {id_!r} is not a valid YAML mapping (dict). Got {type(result).__name__} instead."
+            ).print_warning(console=self.console)
+        return {}
+
+    def _get_id(self, resource: dict[str, Any], default: str) -> str:
+        try:
+            return self.get_id(resource)
+        except (ToolkitRequiredValueError, KeyError):
+            return default
+
+    def dump_resource(self, resource: ExtractionPipelineConfig, local: dict[str, Any] | None = None) -> dict[str, Any]:
+        dumped = resource.as_write().dump()
+        local = local or {}
+        if (
+            "config" in dumped
+            and isinstance(dumped["config"], str)
+            and ("config" not in local or isinstance(local["config"], dict))
+        ):
+            # When we dump a config from CDF, i.e., running `cdf dump extraction-pipeline`, we want to parse the config
+            # as YAML to make it easier to read and edit.
+            if dumped["config"].strip() == "":
+                dumped["config"] = {}
+            else:
+                dumped["config"] = self._validate_config(dumped["config"], dumped)
+        return dumped
+
     def diff_list(
         self, local: list[Any], cdf: list[Any], json_path: tuple[str | int, ...]
     ) -> tuple[dict[int, int], list[int]]:
@@ -346,18 +335,18 @@ class ExtractionPipelineConfigLoader(
             return diff_list_force_hashable(local, cdf)
         return super().diff_list(local, cdf, json_path)
 
-    def _upsert(self, items: ExtractionPipelineConfigWriteList) -> ExtractionPipelineConfigList:
+    def _upsert(self, items: Sequence[ExtractionPipelineConfigWrite]) -> ExtractionPipelineConfigList:
         upserted = ExtractionPipelineConfigList([])
         for item in items:
             created = self.client.extraction_pipelines.config.create(item)
             upserted.append(created)
         return upserted
 
-    def create(self, items: ExtractionPipelineConfigWriteList) -> ExtractionPipelineConfigList:
+    def create(self, items: Sequence[ExtractionPipelineConfigWrite]) -> ExtractionPipelineConfigList:
         return self._upsert(items)
 
     # configs cannot be updated, instead new revision is created
-    def update(self, items: ExtractionPipelineConfigWriteList) -> ExtractionPipelineConfigList:
+    def update(self, items: Sequence[ExtractionPipelineConfigWrite]) -> ExtractionPipelineConfigList:
         return self._upsert(items)
 
     def retrieve(self, ids: SequenceNotStr[str]) -> ExtractionPipelineConfigList:
@@ -418,11 +407,3 @@ class ExtractionPipelineConfigLoader(
             except CogniteAPIError as e:
                 if e.code == 404 and "There is no config stored" in e.message:
                     continue
-
-    @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(("config", ANYTHING), frozenset({"dict"}), is_required=True, _is_nullable=False))
-        return spec