Skip to content
2 changes: 1 addition & 1 deletion airflow-core/src/airflow/models/variable.py
Original file line number Diff line number Diff line change
Expand Up @@ -179,7 +179,7 @@ def get(
if var_val is None:
if default_var is not cls.__NO_DEFAULT_SENTINEL:
return default_var
raise KeyError(f"Variable {key} does not exist")
raise KeyError(f"Variable {key} does not exist.")
if deserialize_json:
obj = json.loads(var_val)
mask_secret(obj, key)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -421,7 +421,9 @@ def _parse_secret_path(self, secret_path: str) -> tuple[str, str]:
if not self.mount_point:
split_secret_path = secret_path.split("/", 1)
if len(split_secret_path) < 2:
raise InvalidPath
raise InvalidPath(
"The variable path you have provided is invalid. Please provide a full path of the format: path/to/secret/variable"
)
return split_secret_path[0], split_secret_path[1]
return self.mount_point, secret_path

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -226,7 +226,13 @@ def get_variable(self, key: str, team_name: str | None = None) -> str | None:
response = self.vault_client.get_secret(
secret_path=(mount_point + "/" if mount_point else "") + secret_path
)
return response.get("value") if response else None
if not response:
return None
try:
return response["value"]
except KeyError:
self.log.warning('Vault secret %s fetched but does not have required key "value"', key)
return None

def get_config(self, key: str) -> str | None:
"""
Expand All @@ -245,4 +251,10 @@ def get_config(self, key: str) -> str | None:
response = self.vault_client.get_secret(
secret_path=(mount_point + "/" if mount_point else "") + secret_path
)
return response.get("value") if response else None
if not response:
return None
try:
return response["value"]
except KeyError:
self.log.warning('Vault config %s fetched but does not have required key "value"', key)
return None
58 changes: 58 additions & 0 deletions providers/hashicorp/tests/unit/hashicorp/secrets/test_vault.py
Original file line number Diff line number Diff line change
Expand Up @@ -263,6 +263,64 @@ def test_get_variable_value_non_existent_key(self, mock_hvac):
)
assert test_client.get_variable("hello") is None

@mock.patch("airflow.providers.hashicorp._internal_client.vault_client.hvac")
@mock.patch("airflow.utils.log.logging_mixin.logging")
def test_get_variable_does_not_contain_value_key(self, mock_hvac, mock_get_logger):
"""
Test that if the 'value' key is not present in Vault, _VaultClient.get_variable
should log a warning and return None
"""
mock_hvac.Client.return_value = mock.MagicMock()
mock_logger = mock.MagicMock()
mock_get_logger.getLogger.return_value = mock_logger

test_client = VaultBackend(
**{
"variables_path": "variables",
"mount_point": "airflow",
"auth_type": "token",
"url": "http://127.0.0.1:8200",
"token": "s.7AU0I51yv1Q1lxOIg1F3ZRAS",
}
)
test_client._log = mock_logger
test_client.vault_client.get_secret = mock.MagicMock(return_value={"test_key": "data"})
result = test_client.get_variable("hello")
assert result is None
mock_logger.warning.assert_called_with(
'Vault secret %s fetched but does not have required key "value"', "hello"
)

@mock.patch("airflow.providers.hashicorp._internal_client.vault_client.hvac")
@mock.patch("airflow.utils.log.logging_mixin.logging")
def test_get_config_does_not_contain_value_key(self, mock_hvac, mock_get_logger):
"""
Test that if the 'value' key is not present in Vault, _VaultClient.get_config
should log a warning and return None
"""
mock_client = mock.MagicMock()
mock_hvac.Client.return_value = mock_client
mock_logger = mock.MagicMock()
mock_get_logger.getLogger.return_value = mock_logger

kwargs = {
"variables_path": "variables",
"mount_point": "airflow",
"auth_type": "token",
"url": "http://127.0.0.1:8200",
"token": "s.7AU0I51yv1Q1lxOIg1F3ZRAS",
}
test_client = VaultBackend(**kwargs)
test_client._log = mock_logger
response = {"test_key": "data"}
test_client.vault_client.get_secret = mock.MagicMock(return_value=response)

returned_uri = test_client.get_config("sql_alchemy_conn")
assert returned_uri is None
mock_logger.warning.assert_called_with(
'Vault config %s fetched but does not have required key "value"', "sql_alchemy_conn"
)

@mock.patch("airflow.providers.hashicorp._internal_client.vault_client.hvac")
def test_auth_failure_raises_error(self, mock_hvac):
mock_client = mock.MagicMock()
Expand Down