You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by el...@apache.org on 2023/03/01 19:33:31 UTC
[airflow] branch main updated: Allow to specify which connection, variable or config are being looked up in the backend using *_lookup_pattern parameters (#29580)
This is an automated email from the ASF dual-hosted git repository.
eladkal pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 5de47910f3 Allow to specify which connection, variable or config are being looked up in the backend using *_lookup_pattern parameters (#29580)
5de47910f3 is described below
commit 5de47910f3ebd803453b8fb5ca6e4f26ad611375
Author: Vincent <97...@users.noreply.github.com>
AuthorDate: Wed Mar 1 14:33:17 2023 -0500
Allow to specify which connection, variable or config are being looked up in the backend using *_lookup_pattern parameters (#29580)
* Allow to specify which connection, variable or config are being looked up in the backend using *_lookup_pattern parameters
---
.../amazon/aws/secrets/secrets_manager.py | 51 +++++--
.../amazon/aws/secrets/systems_manager.py | 49 ++++++-
.../secrets-backends/aws-secrets-manager.rst | 43 +++++-
.../secrets-backends/aws-ssm-parameter-store.rst | 41 +++++-
.../security/secrets/secrets-backend/index.rst | 4 +-
.../amazon/aws/secrets/test_secrets_manager.py | 102 +++++++++++++-
.../amazon/aws/secrets/test_systems_manager.py | 150 +++++++++++++++++++++
7 files changed, 410 insertions(+), 30 deletions(-)
diff --git a/airflow/providers/amazon/aws/secrets/secrets_manager.py b/airflow/providers/amazon/aws/secrets/secrets_manager.py
index ff95e8651f..41181698df 100644
--- a/airflow/providers/amazon/aws/secrets/secrets_manager.py
+++ b/airflow/providers/amazon/aws/secrets/secrets_manager.py
@@ -19,6 +19,7 @@
from __future__ import annotations
import json
+import re
import warnings
from typing import Any
from urllib.parse import unquote
@@ -41,13 +42,16 @@ class SecretsManagerBackend(BaseSecretsBackend, LoggingMixin):
backend = airflow.providers.amazon.aws.secrets.secrets_manager.SecretsManagerBackend
backend_kwargs = {"connections_prefix": "airflow/connections"}
- For example, if secrets prefix is ``airflow/connections/smtp_default``, this would be accessible
- if you provide ``{"connections_prefix": "airflow/connections"}`` and request conn_id ``smtp_default``.
- If variables prefix is ``airflow/variables/hello``, this would be accessible
- if you provide ``{"variables_prefix": "airflow/variables"}`` and request variable key ``hello``.
- And if config_prefix is ``airflow/config/sql_alchemy_conn``, this would be accessible
- if you provide ``{"config_prefix": "airflow/config"}`` and request config
- key ``sql_alchemy_conn``.
+ For example, when ``{"connections_prefix": "airflow/connections"}`` is set, if a secret is defined with
+ the path ``airflow/connections/smtp_default``, the connection with conn_id ``smtp_default`` would be
+ accessible.
+
+ When ``{"variables_prefix": "airflow/variables"}`` is set, if a secret is defined with
+ the path ``airflow/variables/hello``, the variable with the name ``hello`` would be accessible.
+
+ When ``{"config_prefix": "airflow/config"}`` set, if a secret is defined with
+ the path ``airflow/config/sql_alchemy_conn``, the config with they ``sql_alchemy_conn`` would be
+ accessible.
You can also pass additional keyword arguments listed in AWS Connection Extra config
to this class, and they would be used for establishing a connection and passed on to Boto3 client.
@@ -84,12 +88,24 @@ class SecretsManagerBackend(BaseSecretsBackend, LoggingMixin):
:param connections_prefix: Specifies the prefix of the secret to read to get Connections.
If set to None (null value in the configuration), requests for connections will not be
sent to AWS Secrets Manager. If you don't want a connections_prefix, set it as an empty string
+ :param connections_lookup_pattern: Specifies a pattern the connection ID needs to match to be looked up in
+ AWS Secrets Manager. Applies only if `connections_prefix` is not None.
+ If set to None (null value in the configuration), all connections will be looked up first in
+ AWS Secrets Manager.
:param variables_prefix: Specifies the prefix of the secret to read to get Variables.
If set to None (null value in the configuration), requests for variables will not be sent to
AWS Secrets Manager. If you don't want a variables_prefix, set it as an empty string
+ :param variables_lookup_pattern: Specifies a pattern the variable key needs to match to be looked up in
+ AWS Secrets Manager. Applies only if `variables_prefix` is not None.
+ If set to None (null value in the configuration), all variables will be looked up first in
+ AWS Secrets Manager.
:param config_prefix: Specifies the prefix of the secret to read to get Configurations.
If set to None (null value in the configuration), requests for configurations will not be sent to
AWS Secrets Manager. If you don't want a config_prefix, set it as an empty string
+ :param config_lookup_pattern: Specifies a pattern the config key needs to match to be looked up in
+ AWS Secrets Manager. Applies only if `config_prefix` is not None.
+ If set to None (null value in the configuration), all config keys will be looked up first in
+ AWS Secrets Manager.
:param sep: separator used to concatenate secret_prefix and secret_id. Default: "/"
:param extra_conn_words: for using just when you set full_url_mode as false and store
the secrets in different fields of secrets manager. You can add more words for each connection
@@ -101,8 +117,11 @@ class SecretsManagerBackend(BaseSecretsBackend, LoggingMixin):
def __init__(
self,
connections_prefix: str = "airflow/connections",
+ connections_lookup_pattern: str | None = None,
variables_prefix: str = "airflow/variables",
+ variables_lookup_pattern: str | None = None,
config_prefix: str = "airflow/config",
+ config_lookup_pattern: str | None = None,
sep: str = "/",
extra_conn_words: dict[str, list[str]] | None = None,
**kwargs,
@@ -120,6 +139,9 @@ class SecretsManagerBackend(BaseSecretsBackend, LoggingMixin):
self.config_prefix = config_prefix.rstrip(sep)
else:
self.config_prefix = config_prefix
+ self.connections_lookup_pattern = connections_lookup_pattern
+ self.variables_lookup_pattern = variables_lookup_pattern
+ self.config_lookup_pattern = config_lookup_pattern
self.sep = sep
if kwargs.pop("full_url_mode", None) is not None:
@@ -223,7 +245,7 @@ class SecretsManagerBackend(BaseSecretsBackend, LoggingMixin):
if self.connections_prefix is None:
return None
- secret = self._get_secret(self.connections_prefix, conn_id)
+ secret = self._get_secret(self.connections_prefix, conn_id, self.connections_lookup_pattern)
if secret is not None and secret.strip().startswith("{"):
# Before Airflow 2.3, the AWS SecretsManagerBackend added support for JSON secrets.
@@ -264,14 +286,14 @@ class SecretsManagerBackend(BaseSecretsBackend, LoggingMixin):
def get_variable(self, key: str) -> str | None:
"""
- Get Airflow Variable from Environment Variable
+ Get Airflow Variable
:param key: Variable Key
:return: Variable Value
"""
if self.variables_prefix is None:
return None
- return self._get_secret(self.variables_prefix, key)
+ return self._get_secret(self.variables_prefix, key, self.variables_lookup_pattern)
def get_config(self, key: str) -> str | None:
"""
@@ -282,14 +304,19 @@ class SecretsManagerBackend(BaseSecretsBackend, LoggingMixin):
if self.config_prefix is None:
return None
- return self._get_secret(self.config_prefix, key)
+ return self._get_secret(self.config_prefix, key, self.config_lookup_pattern)
- def _get_secret(self, path_prefix, secret_id: str) -> str | None:
+ def _get_secret(self, path_prefix, secret_id: str, lookup_pattern: str | None) -> str | None:
"""
Get secret value from Secrets Manager
:param path_prefix: Prefix for the Path to get Secret
:param secret_id: Secret Key
+ :param lookup_pattern: If provided, `secret_id` must match this pattern to look up the secret in
+ Secrets Manager
"""
+ if lookup_pattern and not re.match(lookup_pattern, secret_id, re.IGNORECASE):
+ return None
+
error_msg = "An error occurred when calling the get_secret_value operation"
if path_prefix:
secrets_path = self.build_path(path_prefix, secret_id, self.sep)
diff --git a/airflow/providers/amazon/aws/secrets/systems_manager.py b/airflow/providers/amazon/aws/secrets/systems_manager.py
index af45cb5e98..35f168e899 100644
--- a/airflow/providers/amazon/aws/secrets/systems_manager.py
+++ b/airflow/providers/amazon/aws/secrets/systems_manager.py
@@ -18,6 +18,7 @@
"""Objects relating to sourcing connections from AWS SSM Parameter Store"""
from __future__ import annotations
+import re
import warnings
from airflow.compat.functools import cached_property
@@ -41,14 +42,26 @@ class SystemsManagerParameterStoreBackend(BaseSecretsBackend, LoggingMixin):
For example, if ssm path is ``/airflow/connections/smtp_default``, this would be accessible
if you provide ``{"connections_prefix": "/airflow/connections"}`` and request conn_id ``smtp_default``.
And if ssm path is ``/airflow/variables/hello``, this would be accessible
- if you provide ``{"variables_prefix": "/airflow/variables"}`` and request conn_id ``hello``.
+ if you provide ``{"variables_prefix": "/airflow/variables"}`` and variable key ``hello``.
:param connections_prefix: Specifies the prefix of the secret to read to get Connections.
If set to None (null), requests for connections will not be sent to AWS SSM Parameter Store.
+ :param connections_lookup_pattern: Specifies a pattern the connection ID needs to match to be looked up in
+ AWS Parameter Store. Applies only if `connections_prefix` is not None.
+ If set to None (null value in the configuration), all connections will be looked up first in
+ AWS Parameter Store.
:param variables_prefix: Specifies the prefix of the secret to read to get Variables.
If set to None (null), requests for variables will not be sent to AWS SSM Parameter Store.
+ :param variables_lookup_pattern: Specifies a pattern the variable key needs to match to be looked up in
+ AWS Parameter Store. Applies only if `variables_prefix` is not None.
+ If set to None (null value in the configuration), all variables will be looked up first in
+ AWS Parameter Store.
:param config_prefix: Specifies the prefix of the secret to read to get Variables.
If set to None (null), requests for configurations will not be sent to AWS SSM Parameter Store.
+ :param config_lookup_pattern: Specifies a pattern the config key needs to match to be looked up in
+ AWS Parameter Store. Applies only if `config_prefix` is not None.
+ If set to None (null value in the configuration), all config keys will be looked up first in
+ AWS Parameter Store.
You can also pass additional keyword arguments listed in AWS Connection Extra config
to this class, and they would be used for establish connection and passed on to Boto3 client.
@@ -67,8 +80,11 @@ class SystemsManagerParameterStoreBackend(BaseSecretsBackend, LoggingMixin):
def __init__(
self,
connections_prefix: str = "/airflow/connections",
+ connections_lookup_pattern: str | None = None,
variables_prefix: str = "/airflow/variables",
+ variables_lookup_pattern: str | None = None,
config_prefix: str = "/airflow/config",
+ config_lookup_pattern: str | None = None,
**kwargs,
):
super().__init__()
@@ -85,6 +101,9 @@ class SystemsManagerParameterStoreBackend(BaseSecretsBackend, LoggingMixin):
else:
self.config_prefix = config_prefix
+ self.connections_lookup_pattern = connections_lookup_pattern
+ self.variables_lookup_pattern = variables_lookup_pattern
+ self.config_lookup_pattern = config_lookup_pattern
self.profile_name = kwargs.get("profile_name", None)
# Remove client specific arguments from kwargs
self.api_version = kwargs.pop("api_version", None)
@@ -122,7 +141,7 @@ class SystemsManagerParameterStoreBackend(BaseSecretsBackend, LoggingMixin):
if self.connections_prefix is None:
return None
- return self._get_secret(self.connections_prefix, conn_id)
+ return self._get_secret(self.connections_prefix, conn_id, self.connections_lookup_pattern)
def get_conn_uri(self, conn_id: str) -> str | None:
"""
@@ -146,7 +165,7 @@ class SystemsManagerParameterStoreBackend(BaseSecretsBackend, LoggingMixin):
def get_variable(self, key: str) -> str | None:
"""
- Get Airflow Variable from Environment Variable
+ Get Airflow Variable
:param key: Variable Key
:return: Variable Value
@@ -154,7 +173,7 @@ class SystemsManagerParameterStoreBackend(BaseSecretsBackend, LoggingMixin):
if self.variables_prefix is None:
return None
- return self._get_secret(self.variables_prefix, key)
+ return self._get_secret(self.variables_prefix, key, self.variables_lookup_pattern)
def get_config(self, key: str) -> str | None:
"""
@@ -166,19 +185,37 @@ class SystemsManagerParameterStoreBackend(BaseSecretsBackend, LoggingMixin):
if self.config_prefix is None:
return None
- return self._get_secret(self.config_prefix, key)
+ return self._get_secret(self.config_prefix, key, self.config_lookup_pattern)
- def _get_secret(self, path_prefix: str, secret_id: str) -> str | None:
+ def _get_secret(self, path_prefix: str, secret_id: str, lookup_pattern: str | None) -> str | None:
"""
Get secret value from Parameter Store.
:param path_prefix: Prefix for the Path to get Secret
:param secret_id: Secret Key
+ :param lookup_pattern: If provided, `secret_id` must match this pattern to look up the secret in
+ Systems Manager
"""
+ if lookup_pattern and not re.match(lookup_pattern, secret_id, re.IGNORECASE):
+ return None
+
ssm_path = self.build_path(path_prefix, secret_id)
+ ssm_path = self._ensure_leading_slash(ssm_path)
+
try:
response = self.client.get_parameter(Name=ssm_path, WithDecryption=True)
return response["Parameter"]["Value"]
except self.client.exceptions.ParameterNotFound:
self.log.debug("Parameter %s not found.", ssm_path)
return None
+
+ def _ensure_leading_slash(self, ssm_path: str):
+ """
+ AWS Systems Manager mandate to have a leading "/". Adding it dynamically if not there to the SSM path
+
+ :param ssm_path: SSM parameter path
+ """
+ if not ssm_path.startswith("/"):
+ ssm_path = f"/{ssm_path}"
+
+ return ssm_path
diff --git a/docs/apache-airflow-providers-amazon/secrets-backends/aws-secrets-manager.rst b/docs/apache-airflow-providers-amazon/secrets-backends/aws-secrets-manager.rst
index 7d1190a904..4a0b32f968 100644
--- a/docs/apache-airflow-providers-amazon/secrets-backends/aws-secrets-manager.rst
+++ b/docs/apache-airflow-providers-amazon/secrets-backends/aws-secrets-manager.rst
@@ -28,7 +28,15 @@ Here is a sample configuration:
[secrets]
backend = airflow.providers.amazon.aws.secrets.secrets_manager.SecretsManagerBackend
- backend_kwargs = {"connections_prefix": "airflow/connections", "variables_prefix": "airflow/variables", "profile_name": "default"}
+ backend_kwargs = {
+ "connections_prefix": "airflow/connections",
+ "connections_lookup_pattern": null,
+ "variables_prefix": "airflow/variables",
+ "variables_lookup_pattern": null,
+ "config_prefix": "airflow/config",
+ "config_lookup_pattern": null,
+ "profile_name": "default"
+ }
To authenticate you can either supply arguments listed in
:ref:`Amazon Webservices Connection Extra config <howto/connection:aws:configuring-the-connection>` or set
@@ -38,7 +46,12 @@ To authenticate you can either supply arguments listed in
[secrets]
backend = airflow.providers.amazon.aws.secrets.secrets_manager.SecretsManagerBackend
- backend_kwargs = {"connections_prefix": "airflow/connections", "variables_prefix": "airflow/variables", "role_arn": "arn:aws:iam::123456789098:role/role-name"}
+ backend_kwargs = {
+ "connections_prefix": "airflow/connections",
+ "variables_prefix": "airflow/variables",
+ "config_prefix": "airflow/config",
+ "role_arn": "arn:aws:iam::123456789098:role/role-name"
+ }
Storing and Retrieving Connections
@@ -107,7 +120,7 @@ If you don't want to use any ``connections_prefix`` for retrieving connections,
Storing and Retrieving Variables
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-If you have set ``variables_prefix`` as ``airflow/variables``, then for an Variable key of ``hello``,
+If you have set ``variables_prefix`` as ``airflow/variables``, then for a Variable key of ``hello``,
you would want to store your Variable at ``airflow/variables/hello``.
Optional lookup
@@ -118,13 +131,33 @@ This will prevent requests being sent to AWS Secrets Manager for the excluded ty
If you want to look up some and not others in AWS Secrets Manager you may do so by setting the relevant ``*_prefix`` parameter of the ones to be excluded as ``null``.
-For example, if you want to set parameter ``connections_prefix`` to ``"airflow/connections"`` and not look up variables, your configuration file should look like this:
+For example, if you want to set parameter ``connections_prefix`` to ``"airflow/connections"`` and not look up variables and config, your configuration file should look like this:
.. code-block:: ini
[secrets]
backend = airflow.providers.amazon.aws.secrets.secrets_manager.SecretsManagerBackend
- backend_kwargs = {"connections_prefix": "airflow/connections", "variables_prefix": null, "profile_name": "default"}
+ backend_kwargs = {
+ "connections_prefix": "airflow/connections",
+ "variables_prefix": null,
+ "config_prefix": null,
+ "profile_name": "default"
+ }
+
+If you want to only lookup a specific subset of connections, variables or config in AWS Secrets Manager, you may do so by setting the relevant ``*_lookup_pattern`` parameter.
+This parameter takes a Regex as a string as value.
+
+For example, if you want to only lookup connections starting by "m" in AWS Secrets Manager, your configuration file should look like this:
+
+.. code-block:: ini
+
+ [secrets]
+ backend = airflow.providers.amazon.aws.secrets.secrets_manager.SecretsManagerBackend
+ backend_kwargs = {
+ "connections_prefix": "airflow/connections",
+ "connections_lookup_pattern": "^m",
+ "profile_name": "default"
+ }
Example of storing Google Secrets in AWS Secrets Manager
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
diff --git a/docs/apache-airflow-providers-amazon/secrets-backends/aws-ssm-parameter-store.rst b/docs/apache-airflow-providers-amazon/secrets-backends/aws-ssm-parameter-store.rst
index 82aabe92d8..bb4d2be625 100644
--- a/docs/apache-airflow-providers-amazon/secrets-backends/aws-ssm-parameter-store.rst
+++ b/docs/apache-airflow-providers-amazon/secrets-backends/aws-ssm-parameter-store.rst
@@ -29,7 +29,15 @@ Here is a sample configuration:
[secrets]
backend = airflow.providers.amazon.aws.secrets.systems_manager.SystemsManagerParameterStoreBackend
- backend_kwargs = {"connections_prefix": "/airflow/connections", "variables_prefix": "/airflow/variables", "profile_name": "default"}
+ backend_kwargs = {
+ "connections_prefix": "airflow/connections",
+ "connections_lookup_pattern": null,
+ "variables_prefix": "airflow/variables",
+ "variables_lookup_pattern": null,
+ "config_prefix": "airflow/config",
+ "config_lookup_pattern": null,
+ "profile_name": "default"
+ }
To authenticate you can either supply arguments listed in
:ref:`Amazon Webservices Connection Extra config <howto/connection:aws:configuring-the-connection>` or set
@@ -39,7 +47,12 @@ To authenticate you can either supply arguments listed in
[secrets]
backend = airflow.providers.amazon.aws.secrets.systems_manager.SystemsManagerParameterStoreBackend
- backend_kwargs = {"connections_prefix": "airflow/connections", "variables_prefix": "airflow/variables", "role_arn": "arn:aws:iam::123456789098:role/role-name"}
+ backend_kwargs = {
+ "connections_prefix": "airflow/connections",
+ "variables_prefix": "airflow/variables",
+ "config_prefix": "airflow/config",
+ "role_arn": "arn:aws:iam::123456789098:role/role-name"
+ }
Optional lookup
@@ -50,13 +63,33 @@ This will prevent requests being sent to AWS SSM Parameter Store for the exclude
If you want to look up some and not others in AWS SSM Parameter Store you may do so by setting the relevant ``*_prefix`` parameter of the ones to be excluded as ``null``.
-For example, if you want to set parameter ``connections_prefix`` to ``"/airflow/connections"`` and not look up variables, your configuration file should look like this:
+For example, if you want to set parameter ``connections_prefix`` to ``"airflow/connections"`` and not look up variables and config, your configuration file should look like this:
.. code-block:: ini
[secrets]
backend = airflow.providers.amazon.aws.secrets.systems_manager.SystemsManagerParameterStoreBackend
- backend_kwargs = {"connections_prefix": "/airflow/connections", "variables_prefix": null, "profile_name": "default"}
+ backend_kwargs = {
+ "connections_prefix": "airflow/connections",
+ "variables_prefix": null,
+ "config_prefix": null,
+ "profile_name": "default"
+ }
+
+If you want to only lookup a specific subset of connections, variables or config in AWS Secrets Manager, you may do so by setting the relevant ``*_lookup_pattern`` parameter.
+This parameter takes a Regex as a string as value.
+
+For example, if you want to only lookup connections starting by "m" in AWS Secrets Manager, your configuration file should look like this:
+
+.. code-block:: ini
+
+ [secrets]
+ backend = airflow.providers.amazon.aws.secrets.secrets_manager.SecretsManagerBackend
+ backend_kwargs = {
+ "connections_prefix": "airflow/connections",
+ "connections_lookup_pattern": "^m",
+ "profile_name": "default"
+ }
Storing and Retrieving Connections
""""""""""""""""""""""""""""""""""
diff --git a/docs/apache-airflow/administration-and-deployment/security/secrets/secrets-backend/index.rst b/docs/apache-airflow/administration-and-deployment/security/secrets/secrets-backend/index.rst
index 06b820a373..76f50a0424 100644
--- a/docs/apache-airflow/administration-and-deployment/security/secrets/secrets-backend/index.rst
+++ b/docs/apache-airflow/administration-and-deployment/security/secrets/secrets-backend/index.rst
@@ -39,7 +39,9 @@ When looking up a connection/variable, by default Airflow will search environmen
database second.
If you enable an alternative secrets backend, it will be searched first, followed by environment variables,
-then metastore. This search ordering is not configurable.
+then metastore. This search ordering is not configurable. Though, in some alternative secrets backend you might have
+the option to filter which connection/variable/config is searched in the secret backend. Please look at the
+documentation of the secret backend you are using to see if such option is available.
.. warning::
diff --git a/tests/providers/amazon/aws/secrets/test_secrets_manager.py b/tests/providers/amazon/aws/secrets/test_secrets_manager.py
index 5d875b69d8..82037f4978 100644
--- a/tests/providers/amazon/aws/secrets/test_secrets_manager.py
+++ b/tests/providers/amazon/aws/secrets/test_secrets_manager.py
@@ -223,10 +223,28 @@ class TestSecretsManagerBackend:
assert secrets_manager_backend.get_variable("test_mysql") is None
+ @mock_secretsmanager
+ def test_get_config_non_existent_key(self):
+ """
+ Test that if Config key is not present,
+ SystemsManagerParameterStoreBackend.get_config should return None
+ """
+ secret_id = "airflow/config/hello"
+ create_param = {
+ "Name": secret_id,
+ }
+ param = {"SecretId": secret_id, "SecretString": "world"}
+
+ secrets_manager_backend = SecretsManagerBackend()
+ secrets_manager_backend.client.create_secret(**create_param)
+ secrets_manager_backend.client.put_secret_value(**param)
+
+ assert secrets_manager_backend.get_config("test") is None
+
@mock.patch("airflow.providers.amazon.aws.secrets.secrets_manager.SecretsManagerBackend._get_secret")
def test_connection_prefix_none_value(self, mock_get_secret):
"""
- Test that if Variable key is not present in AWS Secrets Manager,
+ Test that if Connection ID is not present in AWS Secrets Manager,
SecretsManagerBackend.get_conn_value should return None,
SecretsManagerBackend._get_secret should not be called
"""
@@ -235,6 +253,7 @@ class TestSecretsManagerBackend:
secrets_manager_backend = SecretsManagerBackend(**kwargs)
assert secrets_manager_backend.get_conn_value("test_mysql") is None
+ mock_get_secret.assert_not_called()
@mock.patch("airflow.providers.amazon.aws.secrets.secrets_manager.SecretsManagerBackend._get_secret")
def test_variable_prefix_none_value(self, mock_get_secret):
@@ -253,7 +272,7 @@ class TestSecretsManagerBackend:
@mock.patch("airflow.providers.amazon.aws.secrets.secrets_manager.SecretsManagerBackend._get_secret")
def test_config_prefix_none_value(self, mock_get_secret):
"""
- Test that if Variable key is not present in AWS Secrets Manager,
+ Test that if Config key is not present in AWS Secrets Manager,
SecretsManagerBackend.get_config should return None,
SecretsManagerBackend._get_secret should not be called
"""
@@ -264,6 +283,85 @@ class TestSecretsManagerBackend:
assert secrets_manager_backend.get_config("config") is None
mock_get_secret.assert_not_called()
+ @mock.patch(
+ "airflow.providers.amazon.aws.secrets.secrets_manager.SecretsManagerBackend.client",
+ new_callable=mock.PropertyMock,
+ )
+ @pytest.mark.parametrize(
+ "connection_id, connections_lookup_pattern, num_client_calls",
+ [
+ ("test", "test", 1),
+ ("test", ".*", 1),
+ ("test", "T.*", 1),
+ ("test", "dummy-pattern", 0),
+ ("test", None, 1),
+ ],
+ )
+ def test_connection_lookup_pattern(
+ self, mock_client, connection_id, connections_lookup_pattern, num_client_calls
+ ):
+ """
+ Test that if Connection ID is looked up in AWS Secrets Manager
+ """
+ mock_client().get_secret_value.return_value = {"SecretString": None}
+ kwargs = {"connections_lookup_pattern": connections_lookup_pattern}
+
+ secrets_manager_backend = SecretsManagerBackend(**kwargs)
+ secrets_manager_backend.get_conn_value(connection_id)
+ assert mock_client().get_secret_value.call_count == num_client_calls
+
+ @mock.patch(
+ "airflow.providers.amazon.aws.secrets.secrets_manager.SecretsManagerBackend.client",
+ new_callable=mock.PropertyMock,
+ )
+ @pytest.mark.parametrize(
+ "variable_key, variables_lookup_pattern, num_client_calls",
+ [
+ ("test", "test", 1),
+ ("test", ".*", 1),
+ ("test", "T.*", 1),
+ ("test", "dummy-pattern", 0),
+ ("test", None, 1),
+ ],
+ )
+ def test_variable_lookup_pattern(
+ self, mock_client, variable_key, variables_lookup_pattern, num_client_calls
+ ):
+ """
+ Test that if Variable key is looked up in AWS Secrets Manager
+ """
+ mock_client().get_secret_value.return_value = {"SecretString": None}
+ kwargs = {"variables_lookup_pattern": variables_lookup_pattern}
+
+ secrets_manager_backend = SecretsManagerBackend(**kwargs)
+ secrets_manager_backend.get_variable(variable_key)
+ assert mock_client().get_secret_value.call_count == num_client_calls
+
+ @mock.patch(
+ "airflow.providers.amazon.aws.secrets.secrets_manager.SecretsManagerBackend.client",
+ new_callable=mock.PropertyMock,
+ )
+ @pytest.mark.parametrize(
+ "config_key, config_lookup_pattern, num_client_calls",
+ [
+ ("test", "test", 1),
+ ("test", ".*", 1),
+ ("test", "T.*", 1),
+ ("test", "dummy-pattern", 0),
+ ("test", None, 1),
+ ],
+ )
+ def test_config_lookup_pattern(self, mock_client, config_key, config_lookup_pattern, num_client_calls):
+ """
+ Test that if Variable key is looked up in AWS Secrets Manager
+ """
+ mock_client().get_secret_value.return_value = {"SecretString": None}
+ kwargs = {"config_lookup_pattern": config_lookup_pattern}
+
+ secrets_manager_backend = SecretsManagerBackend(**kwargs)
+ secrets_manager_backend.get_config(config_key)
+ assert mock_client().get_secret_value.call_count == num_client_calls
+
@mock.patch("airflow.providers.amazon.aws.hooks.base_aws.SessionFactory")
def test_passing_client_kwargs(self, mock_session_factory):
secrets_manager_backend = SecretsManagerBackend(
diff --git a/tests/providers/amazon/aws/secrets/test_systems_manager.py b/tests/providers/amazon/aws/secrets/test_systems_manager.py
index c10cc3944a..76c9e2a413 100644
--- a/tests/providers/amazon/aws/secrets/test_systems_manager.py
+++ b/tests/providers/amazon/aws/secrets/test_systems_manager.py
@@ -253,3 +253,153 @@ class TestSsmSecrets:
assert ssm_backend.get_config("config") is None
mock_get_secret.assert_not_called()
+
+ @mock.patch(
+ "airflow.providers.amazon.aws.secrets.systems_manager.SystemsManagerParameterStoreBackend.client",
+ new_callable=mock.PropertyMock,
+ )
+ @pytest.mark.parametrize(
+ "connection_id, connections_lookup_pattern, num_client_calls",
+ [
+ ("test", "test", 1),
+ ("test", ".*", 1),
+ ("test", "T.*", 1),
+ ("test", "dummy-pattern", 0),
+ ("test", None, 1),
+ ],
+ )
+ def test_connection_lookup_pattern(
+ self, mock_client, connection_id, connections_lookup_pattern, num_client_calls
+ ):
+ """
+ Test that if Connection ID is looked up in AWS Parameter Store
+ """
+ mock_client().get_parameter.return_value = {
+ "Parameter": {
+ "Value": None,
+ },
+ }
+ kwargs = {"connections_lookup_pattern": connections_lookup_pattern}
+
+ secrets_manager_backend = SystemsManagerParameterStoreBackend(**kwargs)
+ secrets_manager_backend.get_conn_value(connection_id)
+ assert mock_client().get_parameter.call_count == num_client_calls
+
+ @mock.patch(
+ "airflow.providers.amazon.aws.secrets.systems_manager.SystemsManagerParameterStoreBackend.client",
+ new_callable=mock.PropertyMock,
+ )
+ @pytest.mark.parametrize(
+ "variable_key, variables_lookup_pattern, num_client_calls",
+ [
+ ("test", "test", 1),
+ ("test", ".*", 1),
+ ("test", "T.*", 1),
+ ("test", "dummy-pattern", 0),
+ ("test", None, 1),
+ ],
+ )
+ def test_variable_lookup_pattern(
+ self, mock_client, variable_key, variables_lookup_pattern, num_client_calls
+ ):
+ """
+ Test that if Variable key is looked up in AWS Parameter Store
+ """
+ mock_client().get_parameter.return_value = {
+ "Parameter": {
+ "Value": None,
+ },
+ }
+ kwargs = {"variables_lookup_pattern": variables_lookup_pattern}
+
+ secrets_manager_backend = SystemsManagerParameterStoreBackend(**kwargs)
+ secrets_manager_backend.get_variable(variable_key)
+ assert mock_client().get_parameter.call_count == num_client_calls
+
+ @mock.patch(
+ "airflow.providers.amazon.aws.secrets.systems_manager.SystemsManagerParameterStoreBackend.client",
+ new_callable=mock.PropertyMock,
+ )
+ @pytest.mark.parametrize(
+ "config_key, config_lookup_pattern, num_client_calls",
+ [
+ ("test", "test", 1),
+ ("test", ".*", 1),
+ ("test", "T.*", 1),
+ ("test", "dummy-pattern", 0),
+ ("test", None, 1),
+ ],
+ )
+ def test_config_lookup_pattern(self, mock_client, config_key, config_lookup_pattern, num_client_calls):
+ """
+ Test that if Variable key is looked up in AWS Parameter Store
+ """
+ mock_client().get_parameter.return_value = {
+ "Parameter": {
+ "Value": None,
+ },
+ }
+ kwargs = {"config_lookup_pattern": config_lookup_pattern}
+
+ secrets_manager_backend = SystemsManagerParameterStoreBackend(**kwargs)
+ secrets_manager_backend.get_config(config_key)
+ assert mock_client().get_parameter.call_count == num_client_calls
+
+ @mock.patch(
+ "airflow.providers.amazon.aws.secrets.systems_manager.SystemsManagerParameterStoreBackend.client",
+ new_callable=mock.PropertyMock,
+ )
+ def test_connection_prefix_with_no_leading_slash(self, mock_client):
+ """
+ Test that if Connection ID is looked up in AWS Parameter Store with the added leading "/"
+ """
+ mock_client().get_parameter.return_value = {
+ "Parameter": {
+ "Value": None,
+ },
+ }
+ kwargs = {"connections_prefix": "airflow/connections"}
+
+ secrets_manager_backend = SystemsManagerParameterStoreBackend(**kwargs)
+ secrets_manager_backend.get_conn_value("test_mysql")
+ mock_client().get_parameter.assert_called_with(
+ Name="/airflow/connections/test_mysql", WithDecryption=True
+ )
+
+ @mock.patch(
+ "airflow.providers.amazon.aws.secrets.systems_manager.SystemsManagerParameterStoreBackend.client",
+ new_callable=mock.PropertyMock,
+ )
+ def test_variable_prefix_with_no_leading_slash(self, mock_client):
+ """
+ Test that if Variable key is looked up in AWS Parameter Store with the added leading "/"
+ """
+ mock_client().get_parameter.return_value = {
+ "Parameter": {
+ "Value": None,
+ },
+ }
+ kwargs = {"variables_prefix": "airflow/variables"}
+
+ secrets_manager_backend = SystemsManagerParameterStoreBackend(**kwargs)
+ secrets_manager_backend.get_variable("hello")
+ mock_client().get_parameter.assert_called_with(Name="/airflow/variables/hello", WithDecryption=True)
+
+ @mock.patch(
+ "airflow.providers.amazon.aws.secrets.systems_manager.SystemsManagerParameterStoreBackend.client",
+ new_callable=mock.PropertyMock,
+ )
+ def test_config_prefix_with_no_leading_slash(self, mock_client):
+ """
+ Test that if Config key is looked up in AWS Parameter Store with the added leading "/"
+ """
+ mock_client().get_parameter.return_value = {
+ "Parameter": {
+ "Value": None,
+ },
+ }
+ kwargs = {"config_prefix": "airflow/config"}
+
+ secrets_manager_backend = SystemsManagerParameterStoreBackend(**kwargs)
+ secrets_manager_backend.get_config("config")
+ mock_client().get_parameter.assert_called_with(Name="/airflow/config/config", WithDecryption=True)