Visitar URL original
Secrets backend failover by fhoda · Pull Request #16404 · apache/airflow · GitHub
Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
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
16 changes: 12 additions & 4 deletions airflow/configuration.py
Original file line number Diff line number Diff line change
Expand Up @@ -81,10 +81,18 @@ def run_command(command):

def _get_config_value_from_secret_backend(config_key):
"""Get Config option values from Secret Backend"""
secrets_client = get_custom_secret_backend()
if not secrets_client:
return None
return secrets_client.get_config(config_key)
try:
secrets_client = get_custom_secret_backend()
if not secrets_client:
return None
return secrets_client.get_config(config_key)
except Exception as e: # pylint: disable=broad-except
raise AirflowConfigException(
'Cannot retrieve config from alternative secrets backend. '
'Make sure it is configured properly and that the Backend '
'is accessible.\n'
f'{e}'
)


def _default_config_file_path(file_name: str):
Expand Down
17 changes: 14 additions & 3 deletions airflow/models/connection.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
# under the License.

import json
import logging
Comment thread
uranusjr marked this conversation as resolved.
Resolved
import warnings
from json import JSONDecodeError
from typing import Dict, Optional, Union
Expand All @@ -35,6 +36,8 @@
from airflow.utils.log.secrets_masker import mask_secret
from airflow.utils.module_loading import import_string

log = logging.getLogger(__name__)


def parse_netloc_to_hostname(*args, **kwargs):
"""This method is deprecated."""
Expand Down Expand Up @@ -389,7 +392,15 @@ def get_connection_from_secrets(cls, conn_id: str) -> 'Connection':
:return: connection
"""
for secrets_backend in ensure_secrets_loaded():
conn = secrets_backend.get_connection(conn_id=conn_id)
if conn:
return conn
try:
conn = secrets_backend.get_connection(conn_id=conn_id)
if conn:
return conn
except Exception: # pylint: disable=broad-except
log.exception(
'Unable to retrieve connection from secrets backend (%s). '
'Checking subsequent secrets backend.',
type(secrets_backend).__name__,
)

raise AirflowNotFoundException(f"The conn_id `{conn_id}` isn't defined")
15 changes: 11 additions & 4 deletions airflow/models/variable.py
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,7 @@
from airflow.utils.log.secrets_masker import mask_secret
from airflow.utils.session import provide_session

log = logging.getLogger()
log = logging.getLogger(__name__)


class Variable(Base, LoggingMixin):
Expand Down Expand Up @@ -201,7 +201,14 @@ def get_variable_from_secrets(key: str) -> Optional[str]:
:return: Variable Value
"""
for secrets_backend in ensure_secrets_loaded():
var_val = secrets_backend.get_variable(key=key)
if var_val is not None:
return var_val
try:
var_val = secrets_backend.get_variable(key=key)
if var_val is not None:
return var_val
except Exception: # pylint: disable=broad-except
log.exception(
'Unable to retrieve variable from secrets backend (%s). '
'Checking subsequent secrets backend.',
type(secrets_backend).__name__,
)
return None
37 changes: 37 additions & 0 deletions tests/core/test_configuration.py
Original file line number Diff line number Diff line change
Expand Up @@ -251,6 +251,43 @@ def test_config_from_secret_backend(self, mock_hvac):

assert 'sqlite:////Users/airflow/airflow/airflow.db' == test_conf.get('test', 'sql_alchemy_conn')

@mock.patch("airflow.providers.hashicorp._internal_client.vault_client.hvac")
@conf_vars(
{
("secrets", "backend"): "airflow.providers.hashicorp.secrets.vault.VaultBackend",
("secrets", "backend_kwargs"): '{"url": "http://127.0.0.1:8200", "token": "token"}',
}
)
def test_config_raise_exception_from_secret_backend_connection_error(self, mock_hvac):
"""Get Config Value from a Secret Backend"""

mock_client = mock.MagicMock()
# mock_client.side_effect = AirflowConfigException
mock_hvac.Client.return_value = mock_client
mock_client.secrets.kv.v2.read_secret_version.return_value = Exception

test_config = '''[test]
sql_alchemy_conn_secret = sql_alchemy_conn
'''
test_config_default = '''[test]
sql_alchemy_conn = airflow
'''
test_conf = AirflowConfigParser(default_config=parameterized_config(test_config_default))
test_conf.read_string(test_config)
test_conf.sensitive_config_values = test_conf.sensitive_config_values | {
('test', 'sql_alchemy_conn'),
}

with pytest.raises(
AirflowConfigException,
match=re.escape(
'Cannot retrieve config from alternative secrets backend. '
'Make sure it is configured properly and that the Backend '
'is accessible.'
),
):
test_conf.get('test', 'sql_alchemy_conn')

def test_getboolean(self):
"""Test AirflowConfigParser.getboolean"""
test_config = """
Expand Down