Fix loading keyvault parameters in remote config (#332)

* Fix loading keyvault parameters in remote config * Don't bump version yet
cognitedata · Jun 12, 2024 · 3ff518a · 3ff518a
1 parent 3ddf6ac
commit 3ff518a
Show file tree

Hide file tree

Showing 4 changed files with 125 additions and 11 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -12,6 +12,12 @@ Changes are grouped as follows
 - `Fixed` for any bug fixes.
 - `Security` in case of vulnerabilities.
 
+## Next Release
+
+### Fixed
+
+ * Fixed using the keyvault tag in remote config.
+
 ## 7.2.0
 
 ### Fixed

diff --git a/cognite/extractorutils/configtools/loaders.py b/cognite/extractorutils/configtools/loaders.py
@@ -139,11 +139,10 @@ def _env_constructor(_: yaml.SafeLoader, node: yaml.Node) -> bool:
     return bool_values.get(expanded_value.lower(), expanded_value)
 
 
-def _load_yaml_dict(
+def _load_yaml_dict_raw(
     source: Union[TextIO, str],
-    case_style: str = "hyphen",
     expand_envvars: bool = True,
-    dict_manipulator: Callable[[Dict[str, Any]], Dict[str, Any]] = lambda x: x,
+    keyvault_loader: Optional[KeyVaultLoader] = None,
 ) -> Dict[str, Any]:
     loader = _EnvLoader if expand_envvars else yaml.SafeLoader
 
@@ -159,11 +158,14 @@ def ignore_unknown(self, node: yaml.Node) -> None:
     if not isinstance(source, str):
         source.seek(0)
 
-    keyvault_config = initial_load.get("azure-keyvault", initial_load.get("key-vault"))
+    if keyvault_loader:
+        _EnvLoader.add_constructor("!keyvault", keyvault_loader)
+    else:
+        keyvault_config = initial_load.get("azure-keyvault", initial_load.get("key-vault"))
+        _EnvLoader.add_constructor("!keyvault", KeyVaultLoader(keyvault_config))
 
     _EnvLoader.add_implicit_resolver("!env", re.compile(r"\$\{([^}^{]+)\}"), None)
     _EnvLoader.add_constructor("!env", _env_constructor)
-    _EnvLoader.add_constructor("!keyvault", KeyVaultLoader(keyvault_config))
 
     try:
         config_dict = yaml.load(source, Loader=loader)  # noqa: S506
@@ -173,12 +175,26 @@ def ignore_unknown(self, node: yaml.Node) -> None:
         cause = e.problem or e.context
         raise InvalidConfigError(f"Invalid YAML{formatted_location}: {cause or ''}") from e
 
-    if "azure-keyvault" in config_dict:
-        config_dict.pop("azure-keyvault")
+    return config_dict
+
+
+def _load_yaml_dict(
+    source: Union[TextIO, str],
+    case_style: str = "hyphen",
+    expand_envvars: bool = True,
+    dict_manipulator: Callable[[Dict[str, Any]], Dict[str, Any]] = lambda x: x,
+    keyvault_loader: Optional[KeyVaultLoader] = None,
+) -> Dict[str, Any]:
+    config_dict = _load_yaml_dict_raw(source, expand_envvars, keyvault_loader)
 
     config_dict = dict_manipulator(config_dict)
     config_dict = _to_snake_case(config_dict, case_style)
 
+    if "azure_keyvault" in config_dict:
+        config_dict.pop("azure_keyvault")
+    if "key_vault" in config_dict:
+        config_dict.pop("key_vault")
+
     return config_dict
 
 
@@ -188,9 +204,14 @@ def _load_yaml(
     case_style: str = "hyphen",
     expand_envvars: bool = True,
     dict_manipulator: Callable[[Dict[str, Any]], Dict[str, Any]] = lambda x: x,
+    keyvault_loader: Optional[KeyVaultLoader] = None,
 ) -> CustomConfigClass:
     config_dict = _load_yaml_dict(
-        source, case_style=case_style, expand_envvars=expand_envvars, dict_manipulator=dict_manipulator
+        source,
+        case_style=case_style,
+        expand_envvars=expand_envvars,
+        dict_manipulator=dict_manipulator,
+        keyvault_loader=keyvault_loader,
     )
 
     try:
@@ -239,6 +260,7 @@ def load_yaml(
     config_type: Type[CustomConfigClass],
     case_style: str = "hyphen",
     expand_envvars: bool = True,
+    keyvault_loader: Optional[KeyVaultLoader] = None,
 ) -> CustomConfigClass:
     """
     Read a YAML file, and create a config object based on its contents.
@@ -249,20 +271,28 @@ def load_yaml(
         case_style: Casing convention of config file. Valid options are 'snake', 'hyphen' or 'camel'. Should be
             'hyphen'.
         expand_envvars: Substitute values with the pattern ${VAR} with the content of the environment variable VAR
+        keyvault_loader: Pre-built loader for keyvault tags. Will be loaded from config if not set.
 
     Returns:
         An initialized config object.
 
     Raises:
         InvalidConfigError: If any config field is given as an invalid type, is missing or is unknown
     """
-    return _load_yaml(source=source, config_type=config_type, case_style=case_style, expand_envvars=expand_envvars)
+    return _load_yaml(
+        source=source,
+        config_type=config_type,
+        case_style=case_style,
+        expand_envvars=expand_envvars,
+        keyvault_loader=keyvault_loader,
+    )
 
 
 def load_yaml_dict(
     source: Union[TextIO, str],
     case_style: str = "hyphen",
     expand_envvars: bool = True,
+    keyvault_loader: Optional[KeyVaultLoader] = None,
 ) -> Dict[str, Any]:
     """
     Read a YAML file and return a dictionary from its contents
@@ -272,14 +302,17 @@ def load_yaml_dict(
         case_style: Casing convention of config file. Valid options are 'snake', 'hyphen' or 'camel'. Should be
             'hyphen'.
         expand_envvars: Substitute values with the pattern ${VAR} with the content of the environment variable VAR
+        keyvault_loader: Pre-built loader for keyvault tags. Will be loaded from config if not set.
 
     Returns:
         A raw dict with the contents of the config file.
 
     Raises:
         InvalidConfigError: If any config field is given as an invalid type, is missing or is unknown
     """
-    return _load_yaml_dict(source=source, case_style=case_style, expand_envvars=expand_envvars)
+    return _load_yaml_dict(
+        source=source, case_style=case_style, expand_envvars=expand_envvars, keyvault_loader=keyvault_loader
+    )
 
 
 class ConfigResolver(Generic[CustomConfigClass]):
@@ -374,6 +407,10 @@ def _use_cached_cognite_client(self, tmp_config: _BaseConfig) -> bool:
             and tmp_config.cognite.idp_authentication == self._config.cognite.idp_authentication
         )
 
+    def _get_keyvault_loader(self) -> KeyVaultLoader:
+        temp_config = _load_yaml_dict_raw(self._config_text)
+        return KeyVaultLoader(temp_config.get("azure-keyvault", temp_config.get("key-vault")))
+
     def _resolve_config(self) -> None:
         self._reload_file()
 
@@ -400,6 +437,7 @@ def _resolve_config(self) -> None:
                 source=response.config,
                 config_type=self.config_type,
                 dict_manipulator=lambda d: self._inject_cognite(tmp_config, d),
+                keyvault_loader=self._get_keyvault_loader(),
             )
 
         else:

diff --git a/tests/tests_integration/dummyconfig_keyvault_remote.yaml b/tests/tests_integration/dummyconfig_keyvault_remote.yaml
@@ -0,0 +1,20 @@
+type: remote
+
+azure-keyvault:
+    authentication-method: client-secret
+    client-id: ${KEYVAULT_CLIENT_ID}
+    tenant-id: ${KEYVAULT_TENANT_ID}
+    secret: ${KEYVAULT_CLIENT_SECRET}
+    keyvault-name: extractor-keyvault
+
+cognite:
+    host: ${COGNITE_BASE_URL}
+    project: ${COGNITE_PROJECT}
+    idp-authentication:
+        client-id: ${COGNITE_CLIENT_ID}
+        secret: ${COGNITE_CLIENT_SECRET}
+        token-url: ${COGNITE_TOKEN_URL}
+        scopes:
+          - ${COGNITE_BASE_URL}/.default
+    extraction-pipeline:
+        external-id: utils-test-keyvault-remote
diff --git a/tests/tests_integration/test_configtools_integration.py b/tests/tests_integration/test_configtools_integration.py
@@ -12,12 +12,20 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
+from dataclasses import dataclass
 from io import StringIO
 
 from cognite.client import CogniteClient
-from cognite.client.data_classes import DataSet
+from cognite.client.data_classes import DataSet, ExtractionPipelineConfigWrite, ExtractionPipelineWrite
 from cognite.client.exceptions import CogniteDuplicatedError
 from cognite.extractorutils.configtools import BaseConfig, load_yaml
+from cognite.extractorutils.configtools.loaders import ConfigResolver
+
+
+@dataclass
+class DummyConfig(BaseConfig):
+    dummy_id: str
+    dummy_secret: str
 
 
 def test_dataset_resolve(set_client: CogniteClient):
@@ -82,3 +90,45 @@ def test_dataset_resolve(set_client: CogniteClient):
     config2: BaseConfig = load_yaml(config_file_extid, BaseConfig)
     assert config2.cognite.get_data_set(client).id == data_set_id
     assert config2.cognite.get_data_set(client).name == data_set_name
+
+
+def test_keyvault_and_remote(set_client: CogniteClient):
+    # Set up extraction pipeline
+    data_set_name: str = "Extractor Utils Test Data Set"
+    data_set_extid: str = "extractorUtils-testdataset"
+
+    data_set_id: int
+
+    try:
+        data_set_id = set_client.data_sets.create(DataSet(name=data_set_name, external_id=data_set_extid)).id
+    except CogniteDuplicatedError:
+        data_set_id = set_client.data_sets.retrieve(external_id=data_set_extid).id
+
+    try:
+        set_client.extraction_pipelines.create(
+            ExtractionPipelineWrite(
+                external_id="utils-test-keyvault-remote",
+                name="Utils test keyvault remote",
+                data_set_id=data_set_id,
+            )
+        )
+    except CogniteDuplicatedError:
+        pass
+
+    set_client.extraction_pipelines.config.create(
+        ExtractionPipelineConfigWrite(
+            external_id="utils-test-keyvault-remote",
+            config="""
+                logger:
+                    console:
+                        level: INFO
+                dummy-id: !keyvault test-id
+                dummy-secret: !keyvault test-secret
+            """,
+        )
+    )
+
+    resolver = ConfigResolver("tests/tests_integration/dummyconfig_keyvault_remote.yaml", config_type=DummyConfig)
+    config: DummyConfig = resolver.config
+    assert config.dummy_id == "12345"
+    assert config.dummy_secret == "abcde"