def check_for_write_conflict(key: str) -> None:
"""
Logs a warning if a variable exists outside of the metastore.
If we try to write a variable to the metastore while the same key
exists in an environment variable or custom secrets backend, then
subsequent reads will not read the set value.
:param key: Variable Key
"""
for secrets_backend in ensure_secrets_loaded():
if not isinstance(secrets_backend, MetastoreBackend):
try:
var_val = secrets_backend.get_variable(key=key)
if var_val is not None:
log.warning(
"The variable {key} is defined in the {cls} secrets backend, which takes "
"precedence over reading from the database. The value in the database will be "
"updated, but to read it you have to delete the conflicting variable "
"from {cls}".format(
key=key,
cls=secrets_backend.__class__.__name__))
return
except Exception:
log.exception(
'Unable to retrieve variable from secrets backend (%s). '
'Checking subsequent secrets backend.',
type(secrets_backend).__name__,
)
return None
def get_variable_from_secrets(key: str) -> Optional[str]:
"""
Get Airflow Variable by iterating over all Secret Backends.
:param key: Variable Key
: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
return None
def get_connection_from_secrets(cls, conn_id: str) -> 'Connection':
"""
Get connection by conn_id.
:param conn_id: connection id
:return: connection
"""
for secrets_backend in ensure_secrets_loaded():
conn = secrets_backend.get_connection(conn_id=conn_id)
if conn:
return conn
raise AirflowNotFoundException(f"The conn_id `{conn_id}` isn't defined")
def get_connections_from_secrets(cls, conn_id: str) -> List['Connection']:
"""
Get all connections as an iterable.
:param conn_id: connection id
:return: array of connections
"""
for secrets_backend in ensure_secrets_loaded():
conn_list = secrets_backend.get_connections(conn_id=conn_id)
if conn_list:
return list(conn_list)
raise AirflowNotFoundException("The conn_id `{0}` isn't defined".format(conn_id))
def test_backend_fallback_to_env_var(self, mock_get_uri):
mock_get_uri.return_value = None
backends = ensure_secrets_loaded()
backend_classes = [backend.__class__.__name__ for backend in backends]
self.assertIn('SystemsManagerParameterStoreBackend', backend_classes)
uri = Connection.get_connections_from_secrets(conn_id="test_mysql")
# Assert that SystemsManagerParameterStoreBackend.get_conn_uri was called
mock_get_uri.assert_called_once_with(conn_id='test_mysql')
self.assertEqual('mysql://*****:*****@host:5432/airflow', uri[0].get_uri())
def get_variable_from_secrets(key: str) -> Optional[str]:
"""
Get Airflow Variable by iterating over all Secret Backends.
:param key: Variable Key
:return: Variable Value
"""
for secrets_backend in ensure_secrets_loaded():
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
def get_connection_from_secrets(cls, conn_id: str) -> 'Connection':
"""
Get connection by conn_id.
:param conn_id: connection id
:return: connection
"""
for secrets_backend in ensure_secrets_loaded():
try:
conn = secrets_backend.get_connection(conn_id=conn_id)
if conn:
return conn
except Exception:
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")
