Skip to content

Commit dc0159e

Browse files
kezekeze
authored
Add support in GCP connection for reading key from Secret Manager (#19164)
1 parent 244627e commit dc0159e

File tree

6 files changed

+118
-13
lines changed

6 files changed

+118
-13
lines changed

‎airflow/providers/google/cloud/utils/credentials_provider.py‎

Lines changed: 34 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@
3333
from google.auth.environment_vars import CREDENTIALS, LEGACY_PROJECT, PROJECT
3434

3535
from airflow.exceptions import AirflowException
36+
from airflow.providers.google.cloud._internal_client.secret_manager_client import _SecretManagerClient
3637
from airflow.utils.log.logging_mixin import LoggingMixin
3738
from airflow.utils.process_utils import patch_environ
3839

@@ -202,20 +203,23 @@ def __init__(
202203
self,
203204
key_path: Optional[str] = None,
204205
keyfile_dict: Optional[Dict[str, str]] = None,
206+
key_secret_name: Optional[str] = None,
205207
scopes: Optional[Collection[str]] = None,
206208
delegate_to: Optional[str] = None,
207209
disable_logging: bool = False,
208210
target_principal: Optional[str] = None,
209211
delegates: Optional[Sequence[str]] = None,
210212
) -> None:
211213
super().__init__()
212-
if key_path and keyfile_dict:
214+
key_options = [key_path, key_secret_name, keyfile_dict]
215+
if len([x for x in key_options if x]) > 1:
213216
raise AirflowException(
214-
"The `keyfile_dict` and `key_path` fields are mutually exclusive. "
215-
"Please provide only one value."
217+
"The `keyfile_dict`, `key_path`, and `key_secret_name` fields"
218+
"are all mutually exclusive. Please provide only one value."
216219
)
217220
self.key_path = key_path
218221
self.keyfile_dict = keyfile_dict
222+
self.key_secret_name = key_secret_name
219223
self.scopes = scopes
220224
self.delegate_to = delegate_to
221225
self.disable_logging = disable_logging
@@ -231,6 +235,8 @@ def get_credentials_and_project(self) -> Tuple[google.auth.credentials.Credentia
231235
"""
232236
if self.key_path:
233237
credentials, project_id = self._get_credentials_using_key_path()
238+
elif self.key_secret_name:
239+
credentials, project_id = self._get_credentials_using_key_secret_name()
234240
elif self.keyfile_dict:
235241
credentials, project_id = self._get_credentials_using_keyfile_dict()
236242
else:
@@ -283,6 +289,31 @@ def _get_credentials_using_key_path(self):
283289
project_id = credentials.project_id
284290
return credentials, project_id
285291

292+
def _get_credentials_using_key_secret_name(self):
293+
self._log_debug('Getting connection using JSON key data from GCP secret: %s', self.key_secret_name)
294+
295+
# Use ADC to access GCP Secret Manager.
296+
adc_credentials, adc_project_id = google.auth.default(scopes=self.scopes)
297+
secret_manager_client = _SecretManagerClient(credentials=adc_credentials)
298+
299+
if not secret_manager_client.is_valid_secret_name(self.key_secret_name):
300+
raise AirflowException('Invalid secret name specified for fetching JSON key data.')
301+
302+
secret_value = secret_manager_client.get_secret(self.key_secret_name, adc_project_id)
303+
if secret_value is None:
304+
raise AirflowException(f"Failed getting value of secret {self.key_secret_name}.")
305+
306+
try:
307+
keyfile_dict = json.loads(secret_value)
308+
except json.decoder.JSONDecodeError:
309+
raise AirflowException('Key data read from GCP Secret Manager is not valid JSON.')
310+
311+
credentials = google.oauth2.service_account.Credentials.from_service_account_info(
312+
keyfile_dict, scopes=self.scopes
313+
)
314+
project_id = credentials.project_id
315+
return credentials, project_id
316+
286317
def _get_credentials_using_adc(self):
287318
self._log_info(
288319
'Getting connection using `google.auth.default()` since no key file is defined for hook.'

‎airflow/providers/google/common/hooks/base_google.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -182,6 +182,9 @@ def get_connection_form_widgets() -> Dict[str, Any]:
182182
"extra__google_cloud_platform__scope": StringField(
183183
lazy_gettext('Scopes (comma separated)'), widget=BS3TextFieldWidget()
184184
),
185+
"extra__google_cloud_platform__key_secret_name": StringField(
186+
lazy_gettext('Keyfile Secret Name (in GCP Secret Manager)'), widget=BS3TextFieldWidget()
187+
),
185188
"extra__google_cloud_platform__num_retries": IntegerField(
186189
lazy_gettext('Number of Retries'),
187190
validators=[NumberRange(min=0)],
@@ -225,12 +228,14 @@ def _get_credentials_and_project_id(self) -> Tuple[google.auth.credentials.Crede
225228
keyfile_dict_json = json.loads(keyfile_dict)
226229
except json.decoder.JSONDecodeError:
227230
raise AirflowException('Invalid key JSON.')
231+
key_secret_name: Optional[str] = self._get_field('key_secret_name', None)
228232

229233
target_principal, delegates = _get_target_principal_and_delegates(self.impersonation_chain)
230234

231235
credentials, project_id = get_credentials_and_project_id(
232236
key_path=key_path,
233237
keyfile_dict=keyfile_dict_json,
238+
key_secret_name=key_secret_name,
234239
scopes=self.scopes,
235240
delegate_to=self.delegate_to,
236241
target_principal=target_principal,

‎docs/apache-airflow-providers-google/connections/gcp.rst‎

Lines changed: 24 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -27,17 +27,26 @@ The Google Cloud connection type enables the Google Cloud Integrations.
2727
Authenticating to Google Cloud
2828
------------------------------
2929

30-
There are three ways to connect to Google Cloud using Airflow.
30+
There are two ways to connect to Google Cloud using Airflow.
3131

32-
1. Use `Application Default Credentials
32+
1. Using a `Application Default Credentials
3333
<https://google-auth.readthedocs.io/en/latest/reference/google.auth.html#google.auth.default>`_,
34-
2. Use a `service account
35-
<https://cloud.google.com/docs/authentication/#service_accounts>`_ key
36-
file (JSON format) on disk - ``Keyfile Path``.
37-
3. Use a service account key file (JSON format) from connection configuration - ``Keyfile JSON``.
34+
2. Using a `service account
35+
<https://cloud.google.com/docs/authentication/#service_accounts>`_ by specifying a key file in JSON format.
36+
Key can be specified as a path to the key file (``Keyfile Path``), as a key payload (``Keyfile JSON``)
37+
or as secret in Secret Manager (``Keyfile secret name``). Only one way of defining the key can be used at a time.
38+
If you need to manage multiple keys then you should configure multiple connections.
3839

39-
Only one authorization method can be used at a time. If you need to manage multiple keys then you should
40-
configure multiple connections.
40+
.. warning:: Additional permissions might be needed
41+
42+
Connection which uses key from the Secret Manager requires that `Application Default Credentials
43+
<https://google-auth.readthedocs.io/en/latest/reference/google.auth.html#google.auth.default>`_ (ADC)
44+
have permission to access payloads of secrets.
45+
46+
.. note:: Alternative way of storing connections
47+
48+
Besides storing only key in Secret Manager there is an option for storing entire connection.
49+
For more details take a look at :ref:`Google Secret Manager Backend <google_cloud_secret_manager_backend>`.
4150

4251
Default Connection IDs
4352
----------------------
@@ -89,6 +98,12 @@ Keyfile JSON
8998

9099
Not required if using application default credentials.
91100

101+
Secret name which holds Keyfile JSON
102+
Name of the secret in Secret Manager which contains a `service account
103+
<https://cloud.google.com/docs/authentication/#service_accounts>`_ key.
104+
105+
Not required if using application default credentials.
106+
92107
Scopes (comma separated)
93108
A list of comma-separated `Google Cloud scopes
94109
<https://developers.google.com/identity/protocols/googlescopes>`_ to
@@ -112,6 +127,7 @@ Number of Retries
112127
* ``extra__google_cloud_platform__project`` - Project Id
113128
* ``extra__google_cloud_platform__key_path`` - Keyfile Path
114129
* ``extra__google_cloud_platform__keyfile_dict`` - Keyfile JSON
130+
* ``extra__google_cloud_platform__key_secret_name`` - Secret name which holds Keyfile JSON
115131
* ``extra__google_cloud_platform__scope`` - Scopes
116132
* ``extra__google_cloud_platform__num_retries`` - Number of Retries
117133

‎docs/apache-airflow-providers-google/secrets-backends/google-cloud-secret-manager-backend.rst‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -151,6 +151,10 @@ command as in the example below.
151151
--replication-policy=automatic
152152
Created version [1] of the secret [airflow-variables-first-variable].
153153
154+
.. note:: If only key of the connection should be hidden there is an option to store
155+
only that key in Cloud Secret Manager and not entire connection. For more details take
156+
a look at :ref:`Google Cloud Connection <howto/connection:gcp>`.
157+
154158
Checking configuration
155159
======================
156160

‎tests/providers/google/cloud/utils/test_credentials_provider.py‎

Lines changed: 42 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
import unittest
2222
from io import StringIO
2323
from unittest import mock
24+
from unittest.mock import ANY
2425
from uuid import uuid4
2526

2627
import pytest
@@ -267,12 +268,52 @@ def test_get_credentials_and_project_id_with_service_account_info(self, mock_fro
267268
'connection using JSON Dict'
268269
] == cm.output
269270

271+
@mock.patch("google.auth.default", return_value=("CREDENTIALS", "PROJECT_ID"))
272+
@mock.patch('google.oauth2.service_account.Credentials.from_service_account_info')
273+
@mock.patch("airflow.providers.google.cloud.utils.credentials_provider._SecretManagerClient")
274+
def test_get_credentials_and_project_id_with_key_secret_name(
275+
self, mock_secret_manager_client, mock_from_service_account_info, mock_default
276+
):
277+
mock_secret_manager_client.return_value.is_valid_secret_name.return_value = True
278+
mock_secret_manager_client.return_value.get_secret.return_value = (
279+
"{\"type\":\"service_account\",\"project_id\":\"pid\","
280+
"\"private_key_id\":\"pkid\","
281+
"\"private_key\":\"payload\"}"
282+
)
283+
284+
get_credentials_and_project_id(key_secret_name="secret name")
285+
mock_from_service_account_info.assert_called_once_with(
286+
{
287+
'type': 'service_account',
288+
'project_id': 'pid',
289+
'private_key_id': 'pkid',
290+
'private_key': 'payload',
291+
},
292+
scopes=ANY,
293+
)
294+
295+
@mock.patch("google.auth.default", return_value=("CREDENTIALS", "PROJECT_ID"))
296+
@mock.patch("airflow.providers.google.cloud.utils.credentials_provider._SecretManagerClient")
297+
def test_get_credentials_and_project_id_with_key_secret_name_when_key_is_invalid(
298+
self, mock_secret_manager_client, mock_default
299+
):
300+
mock_secret_manager_client.return_value.is_valid_secret_name.return_value = True
301+
mock_secret_manager_client.return_value.get_secret.return_value = ""
302+
303+
with pytest.raises(
304+
AirflowException,
305+
match=re.escape('Key data read from GCP Secret Manager is not valid JSON.'),
306+
):
307+
get_credentials_and_project_id(key_secret_name="secret name")
308+
270309
def test_get_credentials_and_project_id_with_mutually_exclusive_configuration(
271310
self,
272311
):
273312
with pytest.raises(
274313
AirflowException,
275-
match=re.escape('The `keyfile_dict` and `key_path` fields are mutually exclusive.'),
314+
match=re.escape(
315+
'The `keyfile_dict`, `key_path`, and `key_secret_name` fieldsare all mutually exclusive.'
316+
),
276317
):
277318
get_credentials_and_project_id(key_path='KEY.json', keyfile_dict={'private_key': 'PRIVATE_KEY'})
278319

‎tests/providers/google/common/hooks/test_base_google.py‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -332,6 +332,7 @@ def test_get_credentials_and_project_id_with_default_auth(self, mock_get_creds_a
332332
mock_get_creds_and_proj_id.assert_called_once_with(
333333
key_path=None,
334334
keyfile_dict=None,
335+
key_secret_name=None,
335336
scopes=self.instance.scopes,
336337
delegate_to=None,
337338
target_principal=None,
@@ -348,6 +349,7 @@ def test_get_credentials_and_project_id_with_service_account_file(self, mock_get
348349
mock_get_creds_and_proj_id.assert_called_once_with(
349350
key_path='KEY_PATH.json',
350351
keyfile_dict=None,
352+
key_secret_name=None,
351353
scopes=self.instance.scopes,
352354
delegate_to=None,
353355
target_principal=None,
@@ -375,6 +377,7 @@ def test_get_credentials_and_project_id_with_service_account_info(self, mock_get
375377
mock_get_creds_and_proj_id.assert_called_once_with(
376378
key_path=None,
377379
keyfile_dict=service_account,
380+
key_secret_name=None,
378381
scopes=self.instance.scopes,
379382
delegate_to=None,
380383
target_principal=None,
@@ -392,6 +395,7 @@ def test_get_credentials_and_project_id_with_default_auth_and_delegate(self, moc
392395
mock_get_creds_and_proj_id.assert_called_once_with(
393396
key_path=None,
394397
keyfile_dict=None,
398+
key_secret_name=None,
395399
scopes=self.instance.scopes,
396400
delegate_to="USER",
397401
target_principal=None,
@@ -425,6 +429,7 @@ def test_get_credentials_and_project_id_with_default_auth_and_overridden_project
425429
mock_get_creds_and_proj_id.assert_called_once_with(
426430
key_path=None,
427431
keyfile_dict=None,
432+
key_secret_name=None,
428433
scopes=self.instance.scopes,
429434
delegate_to=None,
430435
target_principal=None,
@@ -442,7 +447,9 @@ def test_get_credentials_and_project_id_with_mutually_exclusive_configuration(
442447
}
443448
with pytest.raises(
444449
AirflowException,
445-
match=re.escape('The `keyfile_dict` and `key_path` fields are mutually exclusive.'),
450+
match=re.escape(
451+
"The `keyfile_dict`, `key_path`, and `key_secret_name` fields" "are all mutually exclusive. "
452+
),
446453
):
447454
self.instance._get_credentials_and_project_id()
448455

@@ -626,6 +633,7 @@ def test_get_credentials_and_project_id_with_impersonation_chain(
626633
mock_get_creds_and_proj_id.assert_called_once_with(
627634
key_path=None,
628635
keyfile_dict=None,
636+
key_secret_name=None,
629637
scopes=self.instance.scopes,
630638
delegate_to=None,
631639
target_principal=target_principal,

0 commit comments

Comments
 (0)