diff --git a/cognite_toolkit/_cdf_tk/commands/pull.py b/cognite_toolkit/_cdf_tk/commands/pull.py index 80fcf9d0f..91cc15d90 100644 --- a/cognite_toolkit/_cdf_tk/commands/pull.py +++ b/cognite_toolkit/_cdf_tk/commands/pull.py @@ -19,6 +19,7 @@ from rich.panel import Panel from cognite_toolkit._cdf_tk.data_classes import ( + BuildVariable, BuiltFullResourceList, DeployResults, ModuleResources, @@ -562,7 +563,7 @@ def pull_resources( for source_file, resources in resources_by_file.items(): file_results = ResourceDeployResult(loader.display_name) has_changes = False - to_write: list[dict[str, Any]] = [] + to_write: dict[T_ID, dict[str, Any]] = {} for resource in resources: local_resource_dict = resource.load_resource_dict(env_vars, validate=True) loaded_any = loader.load_resource(local_resource_dict, ToolGlobals, skip_validation=False) @@ -579,14 +580,14 @@ def pull_resources( are_equal, local_dumped, cdf_dumped = loader.are_equal(loaded, cdf_resource, return_dumped=True) if are_equal: file_results.unchanged += 1 - to_write.append(local_dumped) + to_write[item_id] = local_dumped else: file_results.changed += 1 - to_write.append(cdf_dumped) + to_write[item_id] = cdf_dumped has_changes = True if has_changes and not dry_run: - new_content = self._to_write_content(source_file.read_text(), to_write, resources) + new_content = self._to_write_content(source_file.read_text(), to_write, resources, loader) # type: ignore[arg-type] with source_file.open("w", encoding=ENCODING, newline=NEWLINE) as f: f.write(new_content) @@ -596,9 +597,47 @@ def pull_resources( print(table) def _to_write_content( - self, source: str, to_write: list[dict[str, Any]], resources: BuiltFullResourceList[T_ID] + self, + source: str, + to_write: dict[T_ID, dict[str, Any]], + resources: BuiltFullResourceList[T_ID], + loader: ResourceLoader[ + T_ID, T_WriteClass, T_WritableCogniteResource, T_CogniteResourceList, T_WritableCogniteResourceList + ], ) -> str: - raise NotImplementedError() + # Need to keep comments + # Keep variables + # Keep order + # 1. Replace all variables + # 2. Load source and keep the comments + # 3. Replace all values with the to_write values. + # 4. Dump the yaml + # 5. Replace the variables back + variables = resources[0].build_variables + content, value_by_placeholder = variables.replace(source, use_placeholder=True) + replace_content = variables.replace(source) + _ = YAMLWithComments._extract_comments(content) + loaded = yaml.safe_load(content) + loaded_with_ids = yaml.safe_load(replace_content) + updated: dict[str, Any] | list[dict[str, Any]] + if isinstance(loaded_with_ids, dict) and isinstance(loaded, dict): + item_id = loader.get_id(loaded_with_ids) + if item_id not in to_write: + raise ToolkitMissingResourceError(f"Resource {item_id} not found in to_write.") + item_write = to_write[item_id] + updated = self._replace(loaded, item_write, value_by_placeholder) + elif isinstance(loaded_with_ids, list) and isinstance(loaded, list): + updated = [] + for i, item in enumerate(loaded_with_ids): + item_id = loader.get_id(item) + if item_id not in to_write: + raise ToolkitMissingResourceError(f"Resource {item_id} not found in to_write.") + item_write = to_write[item_id] + updated.append(self._replace(loaded[i], item_write, value_by_placeholder)) + + dumped = yaml.safe_dump(updated, sort_keys=False) + # return YAMLWithComments._dump_yaml_with_comments(dumped, comments, 2, False) + return dumped @staticmethod def _select_resource_ids( @@ -616,3 +655,30 @@ def _select_resource_ids( f"No {loader.display_name} with external id {id_} found in the current configuration in {organization_dir}." ) return BuiltFullResourceList([r for r in local_resources if r.identifier == id_]) + + @classmethod + def _replace( + cls, loaded: dict[str, Any], to_write: dict[str, Any], value_by_placeholder: dict[str, BuildVariable] + ) -> dict[str, Any]: + updated: dict[str, Any] = {} + for key, current_value in loaded.items(): + if key in to_write: + new_value = to_write[key] + if new_value == current_value: + updated[key] = current_value + continue + for placeholder, variable in value_by_placeholder.items(): + if placeholder in current_value: + new_value = new_value.replace(variable.value, f"{{{{ {variable.key} }}}}") + + updated[key] = new_value + elif isinstance(current_value, dict): + updated[key] = cls._replace(current_value, to_write, value_by_placeholder) + elif isinstance(current_value, list): + updated[key] = [cls._replace(item, to_write, value_by_placeholder) for item in current_value] + + for new_key in to_write: + if new_key not in loaded: + updated[new_key] = to_write[new_key] + + return updated diff --git a/cognite_toolkit/_cdf_tk/data_classes/_build_variables.py b/cognite_toolkit/_cdf_tk/data_classes/_build_variables.py index 1f54c6e21..07240f860 100644 --- a/cognite_toolkit/_cdf_tk/data_classes/_build_variables.py +++ b/cognite_toolkit/_cdf_tk/data_classes/_build_variables.py @@ -1,12 +1,13 @@ from __future__ import annotations import re +import uuid from collections import defaultdict from collections.abc import Collection, Iterator, Sequence from dataclasses import dataclass from functools import cached_property from pathlib import Path -from typing import Any, SupportsIndex, overload +from typing import Any, Literal, SupportsIndex, overload from cognite_toolkit._cdf_tk.exceptions import ToolkitValueError from cognite_toolkit._cdf_tk.feature_flags import Flags @@ -151,9 +152,25 @@ def get_module_variables(self, module: ModuleLocation) -> list[BuildVariables]: for variable_set in variable_sets ] - def replace(self, content: str, file_suffix: str = ".yaml") -> str: + @overload + def replace(self, content: str, file_suffix: str = ".yaml", use_placeholder: Literal[False] = False) -> str: ... + + @overload + def replace( + self, content: str, file_suffix: str = ".yaml", use_placeholder: Literal[True] = True + ) -> tuple[str, dict[str, BuildVariable]]: ... + + def replace( + self, content: str, file_suffix: str = ".yaml", use_placeholder: bool = False + ) -> str | tuple[str, dict[str, BuildVariable]]: + variable_by_placeholder: dict[str, BuildVariable] = {} for variable in self: - replace = variable.value_variable + if not use_placeholder: + replace = variable.value_variable + else: + replace = f"VARIABLE_{uuid.uuid4().hex[:8]}" + variable_by_placeholder[replace] = variable + _core_pattern = rf"{{{{\s*{variable.key}\s*}}}}" if file_suffix in {".yaml", ".yml", ".json"}: # Preserve data types @@ -166,8 +183,10 @@ def replace(self, content: str, file_suffix: str = ".yaml") -> str: content = re.sub(pattern, str(replace), content) else: content = re.sub(_core_pattern, str(replace), content) - - return content + if use_placeholder: + return content, variable_by_placeholder + else: + return content # Implemented to get correct type hints def __iter__(self) -> Iterator[BuildVariable]: diff --git a/tests/test_unit/test_cdf_tk/test_commands/test_pull.py b/tests/test_unit/test_cdf_tk/test_commands/test_pull.py index 77a0343b0..693419951 100644 --- a/tests/test_unit/test_cdf_tk/test_commands/test_pull.py +++ b/tests/test_unit/test_cdf_tk/test_commands/test_pull.py @@ -12,6 +12,8 @@ BuiltResourceFull, SourceLocationLazy, ) +from cognite_toolkit._cdf_tk.loaders import DataSetsLoader +from cognite_toolkit._cdf_tk.utils import CDFToolConfig def load_update_diffs_use_cases(): @@ -424,11 +426,11 @@ def test_load_update_changes_dump( def to_write_content_use_cases() -> Iterable: - source = """"name: 'Ingestion' + source = """name: Ingestion externalId: {{ dataset }} description: This dataset contains Transformations, Functions, and Workflows for ingesting data into Cognite Data Fusion. """ - to_write = [{"name": "Ingestion", "externalId": "ingestion", "description": "New description"}] + to_write = {"ingestion": {"name": "Ingestion", "externalId": "ingestion", "description": "New description"}} resources = BuiltFullResourceList( [ BuiltResourceFull( @@ -454,8 +456,8 @@ def to_write_content_use_cases() -> Iterable: ] ) - expected = """"name: 'Ingestion' -externalId: {{ dataset }} + expected = """name: Ingestion +externalId: '{{ dataset }}' description: New description """ @@ -470,12 +472,18 @@ class TestPullCommand: def test_to_write_content( self, source: str, - to_write: list[dict[str, Any]], + to_write: dict[str, [dict[str, Any]]], resources: BuiltFullResourceList, expected: str, + cdf_tool_mock: CDFToolConfig, ) -> None: cmd = PullCommand(silent=True, skip_tracking=True) - actual = cmd._to_write_content(source=source, to_write=to_write, resources=resources) + actual = cmd._to_write_content( + source=source, + to_write=to_write, + resources=resources, + loader=DataSetsLoader.create_loader(cdf_tool_mock, None), + ) assert actual == expected