From 5e97e644650f8dd1f344c889759b42102fbfd9ce Mon Sep 17 00:00:00 2001 From: devendra-lohar Date: Wed, 28 Jan 2026 13:42:23 +0530 Subject: [PATCH 1/3] modify SafeLoaderIgnoreUnknown to expand keyvalut env vars --- cognite/extractorutils/configtools/loaders.py | 21 ++---- tests/tests_unit/test_configtools.py | 64 ++++++++++++++++++- 2 files changed, 70 insertions(+), 15 deletions(-) diff --git a/cognite/extractorutils/configtools/loaders.py b/cognite/extractorutils/configtools/loaders.py index 9efb6ca6..ae7bb332 100644 --- a/cognite/extractorutils/configtools/loaders.py +++ b/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): diff --git a/tests/tests_unit/test_configtools.py b/tests/tests_unit/test_configtools.py index b9354be0..9cbcb866 100644 --- a/tests/tests_unit/test_configtools.py +++ b/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,64 @@ 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() From 6c68f1bea42e201fca3166a7e1155824b82cb796 Mon Sep 17 00:00:00 2001 From: devendra-lohar Date: Wed, 28 Jan 2026 14:28:51 +0530 Subject: [PATCH 2/3] add KEYVAULT_NAME secrets to github workflow --- .github/workflows/test_and_build.yml | 1 + tests/tests_integration/dummyconfig_keyvault_remote.yaml | 2 +- tests/tests_unit/dummyconfig_keyvault.yaml | 2 +- 3 files changed, 3 insertions(+), 2 deletions(-) diff --git a/.github/workflows/test_and_build.yml b/.github/workflows/test_and_build.yml index e3b4350e..bc0efbe0 100644 --- a/.github/workflows/test_and_build.yml +++ b/.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 diff --git a/tests/tests_integration/dummyconfig_keyvault_remote.yaml b/tests/tests_integration/dummyconfig_keyvault_remote.yaml index 27793a40..01bcc48b 100644 --- a/tests/tests_integration/dummyconfig_keyvault_remote.yaml +++ b/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} diff --git a/tests/tests_unit/dummyconfig_keyvault.yaml b/tests/tests_unit/dummyconfig_keyvault.yaml index 7a14a61e..2eb8047a 100644 --- a/tests/tests_unit/dummyconfig_keyvault.yaml +++ b/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 From 8b6e5daba0aee3d3b0280a5704cec45b37f71cf8 Mon Sep 17 00:00:00 2001 From: devendra-lohar Date: Thu, 29 Jan 2026 12:39:25 +0530 Subject: [PATCH 3/3] add testcases --- tests/tests_unit/test_configtools.py | 31 ++++++++++++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/tests/tests_unit/test_configtools.py b/tests/tests_unit/test_configtools.py index 9cbcb866..3bd2676d 100644 --- a/tests/tests_unit/test_configtools.py +++ b/tests/tests_unit/test_configtools.py @@ -812,3 +812,34 @@ def test_keyvault_config_env_var_expansion(monkeypatch: pytest.MonkeyPatch, auth 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"