Skip to content

Commit f77a11d

Browse files
kaxilkaxil
authored
Add Secrets backend for Microsoft Azure Key Vault (apache#10898)
1 parent 76dc7ed commit f77a11d

File tree

6 files changed

+317
-0
lines changed

6 files changed

+317
-0
lines changed

‎airflow/providers/microsoft/azure/secrets/__init__.py‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one
2+
# or more contributor license agreements. See the NOTICE file
3+
# distributed with this work for additional information
4+
# regarding copyright ownership. The ASF licenses this file
5+
# to you under the Apache License, Version 2.0 (the
6+
# "License"); you may not use this file except in compliance
7+
# with the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing,
12+
# software distributed under the License is distributed on an
13+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
# KIND, either express or implied. See the License for the
15+
# specific language governing permissions and limitations
16+
# under the License.

‎airflow/providers/microsoft/azure/secrets/azure_key_vault.py‎

Lines changed: 143 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,143 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one
2+
# or more contributor license agreements. See the NOTICE file
3+
# distributed with this work for additional information
4+
# regarding copyright ownership. The ASF licenses this file
5+
# to you under the Apache License, Version 2.0 (the
6+
# "License"); you may not use this file except in compliance
7+
# with the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing,
12+
# software distributed under the License is distributed on an
13+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
# KIND, either express or implied. See the License for the
15+
# specific language governing permissions and limitations
16+
# under the License.
17+
from typing import Optional
18+
19+
from azure.core.exceptions import ResourceNotFoundError
20+
from azure.identity import DefaultAzureCredential
21+
from azure.keyvault.secrets import SecretClient
22+
from cached_property import cached_property
23+
24+
from airflow.secrets import BaseSecretsBackend
25+
from airflow.utils.log.logging_mixin import LoggingMixin
26+
27+
28+
class AzureKeyVaultBackend(BaseSecretsBackend, LoggingMixin):
29+
"""
30+
Retrieves Airflow Connections or Variables from Azure Key Vault secrets.
31+
32+
The Azure Key Vault can be configured as a secrets backend in the ``airflow.cfg``:
33+
34+
.. code-block:: ini
35+
36+
[secrets]
37+
backend = airflow.providers.microsoft.azure.secrets.azure_key_vault.AzureKeyVaultBackend
38+
backend_kwargs = {"connections_prefix": "airflow-connections", "vault_url": "<azure_key_vault_uri>"}
39+
40+
For example, if the secrets prefix is ``airflow-connections-smtp-default``, this would be accessible
41+
if you provide ``{"connections_prefix": "airflow-connections"}`` and request conn_id ``smtp-default``.
42+
And if variables prefix is ``airflow-variables-hello``, this would be accessible
43+
if you provide ``{"variables_prefix": "airflow-variables"}`` and request variable key ``hello``.
44+
45+
:param connections_prefix: Specifies the prefix of the secret to read to get Connections
46+
:type connections_prefix: str
47+
:param variables_prefix: Specifies the prefix of the secret to read to get Variables
48+
:type variables_prefix: str
49+
:param config_prefix: Specifies the prefix of the secret to read to get Variables.
50+
:type config_prefix: str
51+
:param vault_url: The URL of an Azure Key Vault to use
52+
:type vault_url: str
53+
:param sep: separator used to concatenate secret_prefix and secret_id. Default: "-"
54+
:type sep: str
55+
"""
56+
57+
def __init__(
58+
self,
59+
connections_prefix: str = 'airflow-connections',
60+
variables_prefix: str = 'airflow-variables',
61+
config_prefix: str = 'airflow-config',
62+
vault_url: str = '',
63+
sep: str = '-',
64+
**kwargs,
65+
):
66+
super().__init__()
67+
self.vault_url = vault_url
68+
self.connections_prefix = connections_prefix.rstrip(sep)
69+
self.variables_prefix = variables_prefix.rstrip(sep)
70+
self.config_prefix = config_prefix.rstrip(sep)
71+
self.sep = sep
72+
self.kwargs = kwargs
73+
74+
@cached_property
75+
def client(self):
76+
"""
77+
Create a Azure Key Vault client.
78+
"""
79+
credential = DefaultAzureCredential()
80+
client = SecretClient(vault_url=self.vault_url, credential=credential, **self.kwargs)
81+
return client
82+
83+
def get_conn_uri(self, conn_id: str) -> Optional[str]:
84+
"""
85+
Get an Airflow Connection URI from an Azure Key Vault secret
86+
87+
:param conn_id: The Airflow connection id to retrieve
88+
:type conn_id: str
89+
"""
90+
return self._get_secret(self.connections_prefix, conn_id)
91+
92+
def get_variable(self, key: str) -> Optional[str]:
93+
"""
94+
Get an Airflow Variable from an Azure Key Vault secret.
95+
96+
:param key: Variable Key
97+
:type key: str
98+
:return: Variable Value
99+
"""
100+
return self._get_secret(self.variables_prefix, key)
101+
102+
def get_config(self, key: str) -> Optional[str]:
103+
"""
104+
Get Airflow Configuration
105+
106+
:param key: Configuration Option Key
107+
:return: Configuration Option Value
108+
"""
109+
return self._get_secret(self.config_prefix, key)
110+
111+
@staticmethod
112+
def build_path(path_prefix: str, secret_id: str, sep: str = '-') -> str:
113+
"""
114+
Given a path_prefix and secret_id, build a valid secret name for the Azure Key Vault Backend.
115+
Also replaces underscore in the path with dashes to support easy switching between
116+
environment variables, so ``connection_default`` becomes ``connection-default``.
117+
118+
:param path_prefix: The path prefix of the secret to retrieve
119+
:type path_prefix: str
120+
:param secret_id: Name of the secret
121+
:type secret_id: str
122+
:param sep: Separator used to concatenate path_prefix and secret_id
123+
:type sep: str
124+
"""
125+
path = f'{path_prefix}{sep}{secret_id}'
126+
return path.replace('_', sep)
127+
128+
def _get_secret(self, path_prefix: str, secret_id: str) -> Optional[str]:
129+
"""
130+
Get an Azure Key Vault secret value
131+
132+
:param path_prefix: Prefix for the Path to get Secret
133+
:type path_prefix: str
134+
:param secret_id: Secret Key
135+
:type secret_id: str
136+
"""
137+
name = self.build_path(path_prefix, secret_id, self.sep)
138+
try:
139+
secret = self.client.get_secret(name=name)
140+
return secret.value
141+
except ResourceNotFoundError as ex:
142+
self.log.debug('Secret %s not found: %s', name, ex)
143+
return None

‎docs/autoapi_templates/index.rst‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -414,6 +414,7 @@ All secrets backends derive from :class:`~airflow.secrets.BaseSecretsBackend`.
414414
airflow/providers/amazon/aws/secrets/index
415415
airflow/providers/hashicorp/secrets/index
416416
airflow/providers/google/cloud/secrets/index
417+
airflow/providers/microsoft/azure/secrets/index
417418

418419
Task Log Handlers
419420
-----------------

‎docs/security/secrets/secrets-backend/azure-key-vault-secrets-backend.rst‎

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
.. Licensed to the Apache Software Foundation (ASF) under one
2+
or more contributor license agreements. See the NOTICE file
3+
distributed with this work for additional information
4+
regarding copyright ownership. The ASF licenses this file
5+
to you under the Apache License, Version 2.0 (the
6+
"License"); you may not use this file except in compliance
7+
with the License. You may obtain a copy of the License at
8+
9+
.. http://www.apache.org/licenses/LICENSE-2.0
10+
11+
.. Unless required by applicable law or agreed to in writing,
12+
software distributed under the License is distributed on an
13+
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
KIND, either express or implied. See the License for the
15+
specific language governing permissions and limitations
16+
under the License.
17+
18+
19+
Azure Key Vault Backend
20+
^^^^^^^^^^^^^^^^^^^^^^^
21+
22+
To enable the Azure Key Vault as secrets backend, specify
23+
:py:class:`~airflow.providers.microsoft.azure.secrets.azure_key_vault.AzureKeyVaultBackend`
24+
as the ``backend`` in ``[secrets]`` section of ``airflow.cfg``.
25+
26+
Here is a sample configuration:
27+
28+
.. code-block:: ini
29+
30+
[secrets]
31+
backend = airflow.providers.microsoft.azure.secrets.azure_key_vault.AzureKeyVaultBackend
32+
backend_kwargs = {"connections_prefix": "airflow-connections", "variables_prefix": "airflow-variables", "vault_url": "https://example-akv-resource-name.vault.azure.net/"}
33+
34+
For client authentication, the ``DefaultAzureCredential`` from the Azure Python SDK is used as credential provider,
35+
which supports service principal, managed identity and user credentials.
36+
37+
38+
Storing and Retrieving Connections
39+
""""""""""""""""""""""""""""""""""
40+
41+
If you have set ``connections_prefix`` as ``airflow-connections``, then for a connection id of ``smtp_default``,
42+
you would want to store your connection at ``airflow-connections-smtp-default``.
43+
44+
The value of the secret must be the :ref:`connection URI representation <generating_connection_uri>`
45+
of the connection object.
46+
47+
Storing and Retrieving Variables
48+
""""""""""""""""""""""""""""""""
49+
50+
If you have set ``variables_prefix`` as ``airflow-variables``, then for an Variable key of ``hello``,
51+
you would want to store your Variable at ``airflow-variables-hello``.

‎setup.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -187,6 +187,8 @@ def write_version(filename: str = os.path.join(*[my_dir, "airflow", "git_version
187187
'azure-batch>=8.0.0',
188188
'azure-cosmos>=3.0.1,<4',
189189
'azure-datalake-store>=0.0.45',
190+
'azure-identity>=1.3.1',
191+
'azure-keyvault>=4.1.0',
190192
'azure-kusto-data>=0.0.43,<0.1',
191193
'azure-mgmt-containerinstance>=1.5.0,<2.0',
192194
'azure-mgmt-datalake-store>=0.5.0',

‎tests/providers/microsoft/azure/secrets/test_azure_key_vault.py‎

Lines changed: 104 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,104 @@
1+
#
2+
# Licensed to the Apache Software Foundation (ASF) under one
3+
# or more contributor license agreements. See the NOTICE file
4+
# distributed with this work for additional information
5+
# regarding copyright ownership. The ASF licenses this file
6+
# to you under the Apache License, Version 2.0 (the
7+
# "License"); you may not use this file except in compliance
8+
# with the License. You may obtain a copy of the License at
9+
#
10+
# http://www.apache.org/licenses/LICENSE-2.0
11+
#
12+
# Unless required by applicable law or agreed to in writing,
13+
# software distributed under the License is distributed on an
14+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
# KIND, either express or implied. See the License for the
16+
# specific language governing permissions and limitations
17+
# under the License.
18+
19+
from unittest import TestCase, mock
20+
21+
from azure.core.exceptions import ResourceNotFoundError
22+
23+
from airflow.providers.microsoft.azure.secrets.azure_key_vault import AzureKeyVaultBackend
24+
25+
26+
class TestAzureKeyVaultBackend(TestCase):
27+
@mock.patch('airflow.providers.microsoft.azure.secrets.azure_key_vault.AzureKeyVaultBackend.get_conn_uri')
28+
def test_get_connections(self, mock_get_uri):
29+
mock_get_uri.return_value = 'scheme://user:pass@host:100'
30+
conn_list = AzureKeyVaultBackend().get_connections('fake_conn')
31+
conn = conn_list[0]
32+
self.assertEqual(conn.host, 'host')
33+
34+
@mock.patch('airflow.providers.microsoft.azure.secrets.azure_key_vault.DefaultAzureCredential')
35+
@mock.patch('airflow.providers.microsoft.azure.secrets.azure_key_vault.SecretClient')
36+
def test_get_conn_uri(self, mock_secret_client, mock_azure_cred):
37+
mock_cred = mock.Mock()
38+
mock_sec_client = mock.Mock()
39+
mock_azure_cred.return_value = mock_cred
40+
mock_secret_client.return_value = mock_sec_client
41+
42+
mock_sec_client.get_secret.return_value = mock.Mock(
43+
value='postgresql://airflow:airflow@host:5432/airflow'
44+
)
45+
46+
backend = AzureKeyVaultBackend(vault_url="https://example-akv-resource-name.vault.azure.net/")
47+
returned_uri = backend.get_conn_uri(conn_id='hi')
48+
mock_secret_client.assert_called_once_with(
49+
credential=mock_cred, vault_url='https://example-akv-resource-name.vault.azure.net/'
50+
)
51+
self.assertEqual(returned_uri, 'postgresql://airflow:airflow@host:5432/airflow')
52+
53+
@mock.patch('airflow.providers.microsoft.azure.secrets.azure_key_vault.AzureKeyVaultBackend.client')
54+
def test_get_conn_uri_non_existent_key(self, mock_client):
55+
"""
56+
Test that if the key with connection ID is not present,
57+
AzureKeyVaultBackend.get_connections should return None
58+
"""
59+
conn_id = 'test_mysql'
60+
mock_client.get_secret.side_effect = ResourceNotFoundError
61+
backend = AzureKeyVaultBackend(vault_url="https://example-akv-resource-name.vault.azure.net/")
62+
63+
self.assertIsNone(backend.get_conn_uri(conn_id=conn_id))
64+
self.assertEqual([], backend.get_connections(conn_id=conn_id))
65+
66+
@mock.patch('airflow.providers.microsoft.azure.secrets.azure_key_vault.AzureKeyVaultBackend.client')
67+
def test_get_variable(self, mock_client):
68+
mock_client.get_secret.return_value = mock.Mock(value='world')
69+
backend = AzureKeyVaultBackend()
70+
returned_uri = backend.get_variable('hello')
71+
mock_client.get_secret.assert_called_with(name='airflow-variables-hello')
72+
self.assertEqual('world', returned_uri)
73+
74+
@mock.patch('airflow.providers.microsoft.azure.secrets.azure_key_vault.AzureKeyVaultBackend.client')
75+
def test_get_variable_non_existent_key(self, mock_client):
76+
"""
77+
Test that if Variable key is not present,
78+
AzureKeyVaultBackend.get_variables should return None
79+
"""
80+
mock_client.get_secret.side_effect = ResourceNotFoundError
81+
backend = AzureKeyVaultBackend()
82+
self.assertIsNone(backend.get_variable('test_mysql'))
83+
84+
@mock.patch('airflow.providers.microsoft.azure.secrets.azure_key_vault.AzureKeyVaultBackend.client')
85+
def test_get_secret_value_not_found(self, mock_client):
86+
"""
87+
Test that if a non-existent secret returns None
88+
"""
89+
mock_client.get_secret.side_effect = ResourceNotFoundError
90+
backend = AzureKeyVaultBackend()
91+
self.assertIsNone(
92+
backend._get_secret(path_prefix=backend.connections_prefix, secret_id='test_non_existent')
93+
)
94+
95+
@mock.patch('airflow.providers.microsoft.azure.secrets.azure_key_vault.AzureKeyVaultBackend.client')
96+
def test_get_secret_value(self, mock_client):
97+
"""
98+
Test that get_secret returns the secret value
99+
"""
100+
mock_client.get_secret.return_value = mock.Mock(value='super-secret')
101+
backend = AzureKeyVaultBackend()
102+
secret_val = backend._get_secret('af-secrets', 'test_mysql_password')
103+
mock_client.get_secret.assert_called_with(name='af-secrets-test-mysql-password')
104+
self.assertEqual(secret_val, 'super-secret')

0 commit comments

Comments
 (0)