Support environment variable interpolation in keyvault config (#504)

* modify SafeLoaderIgnoreUnknown to expand keyvalut env vars

* add KEYVAULT_NAME secrets to github workflow

* add testcases

Author: devendra-lohar

.github/workflows/test_and_build.yml

@@ -42,6 +42,7 @@ jobs:
           KEYVAULT_CLIENT_ID: ${{ secrets.KEYVAULT_CLIENT_ID }}
           KEYVAULT_TENANT_ID: ${{ secrets.KEYVAULT_TENANT_ID }}
           KEYVAULT_CLIENT_SECRET: ${{ secrets.KEYVAULT_CLIENT_SECRET }}
+          KEYVAULT_NAME: ${{ secrets.KEYVAULT_NAME }}
           COGNITE_PROJECT: extractor-tests
           COGNITE_BASE_URL: https://greenfield.cognitedata.com
           COGNITE_DEV_PROJECT: extractor-aws-dub-dev-testing

cognite/extractorutils/configtools/loaders.py

@@ -78,8 +78,6 @@ def __init__(self, config: dict | None) -> None:
         self.client: SecretClient | None = None
 
     def _init_client(self) -> None:
-        from dotenv import load_dotenv
-
         if not self.config:
             raise InvalidConfigError(
                 "Attempted to load values from Azure key vault with no key vault configured. "
@@ -108,20 +106,11 @@ def _init_client(self) -> None:
 
             _logger.info("Using Azure ClientSecret credentials to access KeyVault")
 
-            env_file_found = load_dotenv("./.env", override=True)
-
-            if not env_file_found:
-                _logger.info(f"Local environment file not found at {Path.cwd() / '.env'}")
-
             if all(param in self.config for param in auth_parameters):
-                tenant_id = os.path.expandvars(self.config["tenant-id"])
-                client_id = os.path.expandvars(self.config["client-id"])
-                secret = os.path.expandvars(self.config["secret"])
-
                 credentials = ClientSecretCredential(
-                    tenant_id=tenant_id,
-                    client_id=client_id,
-                    client_secret=secret,
+                    tenant_id=self.config["tenant-id"],
+                    client_id=self.config["client-id"],
+                    client_secret=self.config["secret"],
                 )
             else:
                 raise InvalidConfigError(
@@ -184,6 +173,10 @@ def ignore_unknown(self, node: yaml.Node) -> None:
         # Ignoring types since the key can be None.
 
     SafeLoaderIgnoreUnknown.add_constructor(None, SafeLoaderIgnoreUnknown.ignore_unknown)  # type: ignore
+    if expand_envvars:
+        SafeLoaderIgnoreUnknown.add_implicit_resolver("!env", re.compile(r"\$\{([^}^{]+)\}"), None)
+        SafeLoaderIgnoreUnknown.add_constructor("!env", _env_constructor)
+
     initial_load = yaml.load(source, Loader=SafeLoaderIgnoreUnknown)  # noqa: S506
 
     if not isinstance(initial_load, dict):

tests/tests_integration/dummyconfig_keyvault_remote.yaml

@@ -5,7 +5,7 @@ azure-keyvault:
     client-id: ${KEYVAULT_CLIENT_ID}
     tenant-id: ${KEYVAULT_TENANT_ID}
     secret: ${KEYVAULT_CLIENT_SECRET}
-    keyvault-name: extractor-keyvault
+    keyvault-name: ${KEYVAULT_NAME}
 
 cognite:
     host: ${COGNITE_BASE_URL}

tests/tests_unit/dummyconfig_keyvault.yaml

@@ -9,7 +9,7 @@ azure-keyvault:
     client-id: ${KEYVAULT_CLIENT_ID}
     tenant-id: ${KEYVAULT_TENANT_ID}
     secret: ${KEYVAULT_CLIENT_SECRET}
-    keyvault-name: extractor-keyvault
+    keyvault-name: ${KEYVAULT_NAME}
 
 cognite:
   project: mathiaslohne-develop

tests/tests_unit/test_configtools.py

@@ -22,7 +22,7 @@
 from dataclasses import dataclass
 from pathlib import Path
 from typing import IO
-from unittest.mock import patch
+from unittest.mock import MagicMock, patch
 
 import pytest
 import yaml
@@ -51,6 +51,7 @@
 from cognite.extractorutils.configtools.loaders import (
     ConfigResolver,
     compile_patterns,
+    load_yaml_dict,
 )
 from cognite.extractorutils.configtools.validators import matches_pattern, matches_patterns
 from cognite.extractorutils.exceptions import InvalidConfigError
@@ -750,3 +751,95 @@ def test_configresolver_fallback_encoding(tmp_path: Path, caplog: pytest.LogCapt
         assert config.logger.file.path is not None
         assert "café" in config.logger.file.path
         assert any("Falling back to system default encoding." in r.message for r in caplog.records)
+
+
+@pytest.mark.parametrize("auth_method", ["default", "client-secret"])
+def test_keyvault_config_env_var_expansion(monkeypatch: pytest.MonkeyPatch, auth_method: str) -> None:
+    monkeypatch.setenv("MY_KEYVAULT_NAME", "test-keyvault-from-env")
+
+    if auth_method == "default":
+        yaml_config = """
+            azure-keyvault:
+                keyvault-name: ${MY_KEYVAULT_NAME}
+                authentication-method: default
+
+            database:
+                password: !keyvault db-password
+        """
+    else:
+        monkeypatch.setenv("KV_CLIENT_ID", "client-id-123")
+        monkeypatch.setenv("KV_TENANT_ID", "tenant-id-456")
+        monkeypatch.setenv("KV_SECRET", "secret-789")
+        yaml_config = """
+            azure-keyvault:
+                keyvault-name: ${MY_KEYVAULT_NAME}
+                authentication-method: client-secret
+                client-id: ${KV_CLIENT_ID}
+                tenant-id: ${KV_TENANT_ID}
+                secret: ${KV_SECRET}
+
+            database:
+                password: !keyvault db-password
+        """
+
+    with (
+        patch("cognite.extractorutils.configtools.loaders.DefaultAzureCredential") as mock_default_cred,
+        patch("cognite.extractorutils.configtools.loaders.ClientSecretCredential") as mock_client_cred,
+        patch("cognite.extractorutils.configtools.loaders.SecretClient") as mock_secret_client,
+    ):
+        mock_client_instance = MagicMock()
+        mock_client_instance.get_secret.return_value = MagicMock(value="secret-from-keyvault")
+        mock_secret_client.return_value = mock_client_instance
+        mock_default_cred.return_value = MagicMock()
+        mock_client_cred.return_value = MagicMock()
+
+        config = load_yaml_dict(yaml_config)
+
+        mock_secret_client.assert_called_once()
+        call_kwargs = mock_secret_client.call_args[1]
+        assert call_kwargs["vault_url"] == "https://test-keyvault-from-env.vault.azure.net"
+
+        assert config["database"]["password"] == "secret-from-keyvault"
+
+        if auth_method == "default":
+            mock_default_cred.assert_called_once()
+            mock_client_cred.assert_not_called()
+
+        if auth_method == "client-secret":
+            mock_client_cred.assert_called_once_with(
+                tenant_id="tenant-id-456",
+                client_id="client-id-123",
+                client_secret="secret-789",
+            )
+            mock_default_cred.assert_not_called()
+
+
+def test_keyvault_tag_without_config_raises_error() -> None:
+    yaml_config = """
+        database:
+            password: !keyvault db-password
+    """
+
+    with pytest.raises(InvalidConfigError) as e:
+        load_yaml_dict(yaml_config)
+    assert (
+        e.value.message
+        == "Attempted to load values from Azure key vault with no key vault configured. Include an `azure-keyvault` section in your config to use the !keyvault tag."
+    )
+
+
+def test_keyvault_client_secret_missing_raises_error() -> None:
+    yaml_config = """
+        azure-keyvault:
+            keyvault-name: test-keyvault-from-env
+            authentication-method: client-secret
+            client-id: client-id-123
+            tenant-id: tenant-id-456
+
+        database:
+            password: !keyvault db-password
+    """
+
+    with pytest.raises(InvalidConfigError) as e:
+        load_yaml_dict(yaml_config)
+    assert e.value.message == "Missing client secret parameters. client-id, tenant-id and client-secret are mandatory"