Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 7 additions & 14 deletions cognite/extractorutils/configtools/loaders.py
Original file line number Diff line number Diff line change
Expand Up @@ -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. "
Expand Down Expand Up @@ -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)
Comment thread
einarmo marked this conversation as resolved.

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"],
Comment thread
devendra-lohar marked this conversation as resolved.
)
else:
raise InvalidConfigError(
Expand Down Expand Up @@ -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):
Expand Down
64 changes: 63 additions & 1 deletion tests/tests_unit/test_configtools.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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:
Comment thread
devendra-lohar marked this conversation as resolved.
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()
Loading