From 4f5fd2bd875965fc28986e66e8f52bd66f936ca3 Mon Sep 17 00:00:00 2001 From: Philippe Gagnon <12717218+pgagnon@users.noreply.github.com> Date: Mon, 6 Feb 2023 16:40:06 -0500 Subject: [PATCH 1/9] Implement file web identity credentials provider for AWS hook AssumeRoleWithWebIdentity --- .../providers/amazon/aws/hooks/base_aws.py | 28 +++++++++++---- .../amazon/aws/hooks/test_base_aws.py | 36 ++++++++++++++++++- 2 files changed, 57 insertions(+), 7 deletions(-) diff --git a/airflow/providers/amazon/aws/hooks/base_aws.py b/airflow/providers/amazon/aws/hooks/base_aws.py index 80a40bb734954..ab0613949baae 100644 --- a/airflow/providers/amazon/aws/hooks/base_aws.py +++ b/airflow/providers/amazon/aws/hooks/base_aws.py @@ -312,12 +312,19 @@ def _get_web_identity_credential_fetcher( base_session = self.basic_session._session or botocore.session.get_session() client_creator = base_session.create_client federation = self.extra_config.get("assume_role_with_web_identity_federation") - if federation == "google": - web_identity_token_loader = self._get_google_identity_token_loader() - else: - raise AirflowException( - f'Unsupported federation: {federation}. Currently "google" only are supported.' - ) + + web_identity_token_loader = ( + { + "file": self._get_file_token_loader, + "google": self._get_google_identity_token_loader, + }.get(federation)() + if type(federation) == str + else None + ) + + if not web_identity_token_loader: + raise AirflowException(f"Unsupported federation: {federation}.") + return botocore.credentials.AssumeRoleWithWebIdentityCredentialFetcher( client_creator=client_creator, web_identity_token_loader=web_identity_token_loader, @@ -325,6 +332,15 @@ def _get_web_identity_credential_fetcher( extra_args=self.conn.assume_role_kwargs, ) + def _get_file_token_loader(self): + from botocore.credentials import AssumeRoleWithWebIdentityProvider, FileWebIdentityTokenLoader + + token_file = self.extra_config.get("assume_role_with_web_identity_token_file") or os.getenv( + AssumeRoleWithWebIdentityProvider._CONFIG_TO_ENV_VAR["web_identity_token_file"] + ) + + return FileWebIdentityTokenLoader(token_file) + def _get_google_identity_token_loader(self): from google.auth.transport import requests as requests_transport diff --git a/tests/providers/amazon/aws/hooks/test_base_aws.py b/tests/providers/amazon/aws/hooks/test_base_aws.py index 8db11fbb380ca..1a7238d349abf 100644 --- a/tests/providers/amazon/aws/hooks/test_base_aws.py +++ b/tests/providers/amazon/aws/hooks/test_base_aws.py @@ -22,6 +22,7 @@ from base64 import b64encode from datetime import datetime, timedelta, timezone from unittest import mock +from unittest.mock import mock_open from uuid import UUID import boto3 @@ -29,6 +30,7 @@ from botocore.config import Config from botocore.credentials import ReadOnlyCredentials from botocore.exceptions import NoCredentialsError +from botocore.utils import FileWebIdentityTokenLoader from moto import mock_dynamodb, mock_emr, mock_iam, mock_sts from moto.core import DEFAULT_ACCOUNT_ID @@ -484,10 +486,42 @@ def import_mock(name, *args): [mock.call.get_default_id_token_credentials(target_audience="aws-federation.airflow.apache.org")] ) + @mock.patch( + "botocore.utils.FileWebIdentityTokenLoader.__init__.__defaults__", new=(mock_open(read_data="TOKEN"),) + ) + @mock.patch.object( + AwsBaseHook, + "get_connection", + return_value=Connection( + conn_id="aws_default", + conn_type="aws", + extra=json.dumps( + { + "role_arn": "arn:aws:iam::123456:role/role_arn", + "assume_role_method": "assume_role_with_web_identity", + "assume_role_with_web_identity_token_file": "/my-token-path", + "assume_role_with_web_identity_federation": "file", + } + ), + ), + ) + @mock.patch("botocore.credentials.AssumeRoleWithWebIdentityCredentialFetcher") + @mock.patch("botocore.session.Session") + def test_get_credentials_from_token_file( + self, mock_session, mock_credentials_fetcher, mock_get_connection + ): + hook = AwsBaseHook(aws_conn_id="aws_default", client_type="airflow_test") + hook.get_session() + + assert isinstance( + mock_credentials_fetcher.call_args.kwargs["web_identity_token_loader"], FileWebIdentityTokenLoader + ) + + assert mock_credentials_fetcher.call_args.kwargs["web_identity_token_loader"]() == "TOKEN" + @mock.patch.object(AwsBaseHook, "get_connection") @mock_sts def test_assume_role_with_saml(self, mock_get_connection): - idp_url = "https://my-idp.local.corp" principal_arn = "principal_arn_1234567890" role_arn = "arn:aws:iam::123456:role/role_arn" From 4015d086626383caca16930726fa8d0d7b1862e5 Mon Sep 17 00:00:00 2001 From: Philippe Gagnon <12717218+pgagnon@users.noreply.github.com> Date: Mon, 20 Feb 2023 10:20:54 -0500 Subject: [PATCH 2/9] Update mocks & mypy failure --- airflow/providers/amazon/aws/hooks/base_aws.py | 16 ++++++---------- .../providers/amazon/aws/hooks/test_base_aws.py | 17 ++++++++++++----- 2 files changed, 18 insertions(+), 15 deletions(-) diff --git a/airflow/providers/amazon/aws/hooks/base_aws.py b/airflow/providers/amazon/aws/hooks/base_aws.py index ab0613949baae..6c5a2b4bd3fc5 100644 --- a/airflow/providers/amazon/aws/hooks/base_aws.py +++ b/airflow/providers/amazon/aws/hooks/base_aws.py @@ -311,16 +311,12 @@ def _get_web_identity_credential_fetcher( ) -> botocore.credentials.AssumeRoleWithWebIdentityCredentialFetcher: base_session = self.basic_session._session or botocore.session.get_session() client_creator = base_session.create_client - federation = self.extra_config.get("assume_role_with_web_identity_federation") - - web_identity_token_loader = ( - { - "file": self._get_file_token_loader, - "google": self._get_google_identity_token_loader, - }.get(federation)() - if type(federation) == str - else None - ) + federation = str(self.extra_config.get("assume_role_with_web_identity_federation")) + + web_identity_token_loader = { + "file": self._get_file_token_loader, + "google": self._get_google_identity_token_loader, + }.get(federation, lambda: None)() if not web_identity_token_loader: raise AirflowException(f"Unsupported federation: {federation}.") diff --git a/tests/providers/amazon/aws/hooks/test_base_aws.py b/tests/providers/amazon/aws/hooks/test_base_aws.py index 1a7238d349abf..d48d527251c83 100644 --- a/tests/providers/amazon/aws/hooks/test_base_aws.py +++ b/tests/providers/amazon/aws/hooks/test_base_aws.py @@ -487,7 +487,8 @@ def import_mock(name, *args): ) @mock.patch( - "botocore.utils.FileWebIdentityTokenLoader.__init__.__defaults__", new=(mock_open(read_data="TOKEN"),) + "airflow.providers.amazon.aws.hooks.base_aws.botocore.utils.FileWebIdentityTokenLoader.__init__.__defaults__", + new=(mock_open(read_data="TOKEN"),), ) @mock.patch.object( AwsBaseHook, @@ -505,18 +506,24 @@ def import_mock(name, *args): ), ), ) - @mock.patch("botocore.credentials.AssumeRoleWithWebIdentityCredentialFetcher") - @mock.patch("botocore.session.Session") + @mock.patch( + "airflow.providers.amazon.aws.hooks.base_aws.botocore.credentials.AssumeRoleWithWebIdentityCredentialFetcher" + ) + @mock.patch("airflow.providers.amazon.aws.hooks.base_aws.botocore.session.Session") def test_get_credentials_from_token_file( self, mock_session, mock_credentials_fetcher, mock_get_connection ): - hook = AwsBaseHook(aws_conn_id="aws_default", client_type="airflow_test") - hook.get_session() + AwsBaseHook(aws_conn_id="aws_default", client_type="airflow_test").get_session() assert isinstance( mock_credentials_fetcher.call_args.kwargs["web_identity_token_loader"], FileWebIdentityTokenLoader ) + assert ( + mock_credentials_fetcher.call_args.kwargs["web_identity_token_loader"]._web_identity_token_path + == "/my-token-path" + ) + assert mock_credentials_fetcher.call_args.kwargs["web_identity_token_loader"]() == "TOKEN" @mock.patch.object(AwsBaseHook, "get_connection") From 300b79b559b3808a758c19f39958c6b4ad6ed611 Mon Sep 17 00:00:00 2001 From: Philippe Gagnon <12717218+pgagnon@users.noreply.github.com> Date: Mon, 20 Feb 2023 15:44:36 -0500 Subject: [PATCH 3/9] Fix tests for python 3.7 compatibility --- tests/providers/amazon/aws/hooks/test_base_aws.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/tests/providers/amazon/aws/hooks/test_base_aws.py b/tests/providers/amazon/aws/hooks/test_base_aws.py index d48d527251c83..fc602616c10a2 100644 --- a/tests/providers/amazon/aws/hooks/test_base_aws.py +++ b/tests/providers/amazon/aws/hooks/test_base_aws.py @@ -515,16 +515,18 @@ def test_get_credentials_from_token_file( ): AwsBaseHook(aws_conn_id="aws_default", client_type="airflow_test").get_session() + mock_credentials_fetcher_args, mock_credentials_fetcher_kwargs = mock_credentials_fetcher.call_args + assert isinstance( - mock_credentials_fetcher.call_args.kwargs["web_identity_token_loader"], FileWebIdentityTokenLoader + mock_credentials_fetcher_kwargs["web_identity_token_loader"], FileWebIdentityTokenLoader ) assert ( - mock_credentials_fetcher.call_args.kwargs["web_identity_token_loader"]._web_identity_token_path + mock_credentials_fetcher_kwargs["web_identity_token_loader"]._web_identity_token_path == "/my-token-path" ) - assert mock_credentials_fetcher.call_args.kwargs["web_identity_token_loader"]() == "TOKEN" + assert mock_credentials_fetcher_kwargs["web_identity_token_loader"]() == "TOKEN" @mock.patch.object(AwsBaseHook, "get_connection") @mock_sts From 1c88afd8af7d772ed809462962bda77f840993d6 Mon Sep 17 00:00:00 2001 From: Philippe Gagnon <12717218+pgagnon@users.noreply.github.com> Date: Mon, 20 Feb 2023 16:03:49 -0500 Subject: [PATCH 4/9] Add docs --- docs/apache-airflow-providers-amazon/connections/aws.rst | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/docs/apache-airflow-providers-amazon/connections/aws.rst b/docs/apache-airflow-providers-amazon/connections/aws.rst index 5c2faa5598d10..90e9bccc95a8d 100644 --- a/docs/apache-airflow-providers-amazon/connections/aws.rst +++ b/docs/apache-airflow-providers-amazon/connections/aws.rst @@ -109,6 +109,15 @@ Extra (optional) if not specified then **assume_role** is used. * ``assume_role_kwargs``: Additional **kwargs** passed to ``assume_role_method``. + The following extra parameters if ``assume_role_method`` is set to ``assume_role_with_web_identity``: + + * ``assume_role_with_web_identity_federation``: The federation type, which is used to determine which token loader to use to retrieve + the access token. Currently ``file`` and ``google`` are supported. + * ``assume_role_with_web_identity_token_file``: The path to a file on the filesystem that contains the access token used to + authenticate with the AWS STS service, for the ``file`` federation type. If not specified, then the value of the + ``AWS_WEB_IDENTITY_TOKEN_FILE`` environment variable will be used. + * ``assume_role_with_web_identity_federation_audience``: The ``aud`` claim of the access token, if using the ``google`` federation type. + The following extra parameters pass to :external:py:meth:`boto3.session.Session.client` or :external:py:meth:`boto3.session.Session.resource`. From 57aab597b73bdf5b94ea8f4f6e3d825e91c5a99d Mon Sep 17 00:00:00 2001 From: Philippe Gagnon <12717218+pgagnon@users.noreply.github.com> Date: Tue, 21 Feb 2023 10:14:56 -0500 Subject: [PATCH 5/9] Remove references to botocore non-public APIs --- .../providers/amazon/aws/hooks/base_aws.py | 4 +-- .../amazon/aws/hooks/test_base_aws.py | 29 +++++++------------ 2 files changed, 13 insertions(+), 20 deletions(-) diff --git a/airflow/providers/amazon/aws/hooks/base_aws.py b/airflow/providers/amazon/aws/hooks/base_aws.py index 6c5a2b4bd3fc5..936e1445c2441 100644 --- a/airflow/providers/amazon/aws/hooks/base_aws.py +++ b/airflow/providers/amazon/aws/hooks/base_aws.py @@ -329,10 +329,10 @@ def _get_web_identity_credential_fetcher( ) def _get_file_token_loader(self): - from botocore.credentials import AssumeRoleWithWebIdentityProvider, FileWebIdentityTokenLoader + from botocore.credentials import FileWebIdentityTokenLoader token_file = self.extra_config.get("assume_role_with_web_identity_token_file") or os.getenv( - AssumeRoleWithWebIdentityProvider._CONFIG_TO_ENV_VAR["web_identity_token_file"] + "AWS_WEB_IDENTITY_TOKEN_FILE" ) return FileWebIdentityTokenLoader(token_file) diff --git a/tests/providers/amazon/aws/hooks/test_base_aws.py b/tests/providers/amazon/aws/hooks/test_base_aws.py index fc602616c10a2..4a46e3c18d933 100644 --- a/tests/providers/amazon/aws/hooks/test_base_aws.py +++ b/tests/providers/amazon/aws/hooks/test_base_aws.py @@ -486,10 +486,6 @@ def import_mock(name, *args): [mock.call.get_default_id_token_credentials(target_audience="aws-federation.airflow.apache.org")] ) - @mock.patch( - "airflow.providers.amazon.aws.hooks.base_aws.botocore.utils.FileWebIdentityTokenLoader.__init__.__defaults__", - new=(mock_open(read_data="TOKEN"),), - ) @mock.patch.object( AwsBaseHook, "get_connection", @@ -513,20 +509,17 @@ def import_mock(name, *args): def test_get_credentials_from_token_file( self, mock_session, mock_credentials_fetcher, mock_get_connection ): - AwsBaseHook(aws_conn_id="aws_default", client_type="airflow_test").get_session() - - mock_credentials_fetcher_args, mock_credentials_fetcher_kwargs = mock_credentials_fetcher.call_args - - assert isinstance( - mock_credentials_fetcher_kwargs["web_identity_token_loader"], FileWebIdentityTokenLoader - ) - - assert ( - mock_credentials_fetcher_kwargs["web_identity_token_loader"]._web_identity_token_path - == "/my-token-path" - ) - - assert mock_credentials_fetcher_kwargs["web_identity_token_loader"]() == "TOKEN" + mock_open_ = mock_open(read_data="TOKEN") + with mock.patch( + "airflow.providers.amazon.aws.hooks.base_aws.botocore.utils.FileWebIdentityTokenLoader.__init__.__defaults__", + new=(mock_open_,), + ) as mock_token_loader_init_defaults: + AwsBaseHook(aws_conn_id="aws_default", client_type="airflow_test").get_session() + + _, mock_creds_fetcher_kwargs = mock_credentials_fetcher.call_args + assert isinstance(mock_creds_fetcher_kwargs["web_identity_token_loader"], FileWebIdentityTokenLoader) + assert mock_creds_fetcher_kwargs["web_identity_token_loader"]() == "TOKEN" + assert mock_open_.call_args[0][0] == "/my-token-path" @mock.patch.object(AwsBaseHook, "get_connection") @mock_sts From 71883c0c3c977bb3004d826ceada05e78c5823fd Mon Sep 17 00:00:00 2001 From: Philippe Gagnon <12717218+pgagnon@users.noreply.github.com> Date: Tue, 21 Feb 2023 10:49:25 -0500 Subject: [PATCH 6/9] Fixup linting --- tests/providers/amazon/aws/hooks/test_base_aws.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/providers/amazon/aws/hooks/test_base_aws.py b/tests/providers/amazon/aws/hooks/test_base_aws.py index 4a46e3c18d933..4e6087ae50824 100644 --- a/tests/providers/amazon/aws/hooks/test_base_aws.py +++ b/tests/providers/amazon/aws/hooks/test_base_aws.py @@ -513,7 +513,7 @@ def test_get_credentials_from_token_file( with mock.patch( "airflow.providers.amazon.aws.hooks.base_aws.botocore.utils.FileWebIdentityTokenLoader.__init__.__defaults__", new=(mock_open_,), - ) as mock_token_loader_init_defaults: + ): AwsBaseHook(aws_conn_id="aws_default", client_type="airflow_test").get_session() _, mock_creds_fetcher_kwargs = mock_credentials_fetcher.call_args From fd968f9e64ebf91c42e5a4da140ecff889f9b652 Mon Sep 17 00:00:00 2001 From: Philippe Gagnon <12717218+pgagnon@users.noreply.github.com> Date: Tue, 21 Feb 2023 11:34:44 -0500 Subject: [PATCH 7/9] Add example to documentation --- .../connections/aws.rst | 25 +++++++++++++------ 1 file changed, 18 insertions(+), 7 deletions(-) diff --git a/docs/apache-airflow-providers-amazon/connections/aws.rst b/docs/apache-airflow-providers-amazon/connections/aws.rst index 90e9bccc95a8d..76509f9e71986 100644 --- a/docs/apache-airflow-providers-amazon/connections/aws.rst +++ b/docs/apache-airflow-providers-amazon/connections/aws.rst @@ -89,7 +89,7 @@ Extra (optional) * ``service_config``: json used to specify configuration/parameters for different AWS services, such as S3 or STS. - The following extra parameters used to create an initial :external:py:class:`boto3.session.Session`: + The following extra parameters are used to create an initial :external:py:class:`boto3.session.Session`: * ``aws_access_key_id``: AWS access key ID used for the initial connection. * ``aws_secret_access_key``: AWS secret access key used for the initial connection @@ -99,7 +99,7 @@ Extra (optional) * ``profile_name``: The name of a profile to use listed in `configuration and credential file settings `__. - The following extra parameters used for `assume role `__: + The following extra parameters are used for `assume role `__: * ``role_arn``: If specified, then assume this role, obtaining a set of temporary security credentials using the ``assume_role_method``. * ``assume_role_method``: AWS STS client method, one of @@ -109,7 +109,7 @@ Extra (optional) if not specified then **assume_role** is used. * ``assume_role_kwargs``: Additional **kwargs** passed to ``assume_role_method``. - The following extra parameters if ``assume_role_method`` is set to ``assume_role_with_web_identity``: + The following extra parameters are available if ``assume_role_method`` is set to ``assume_role_with_web_identity``: * ``assume_role_with_web_identity_federation``: The federation type, which is used to determine which token loader to use to retrieve the access token. Currently ``file`` and ``google`` are supported. @@ -118,7 +118,7 @@ Extra (optional) ``AWS_WEB_IDENTITY_TOKEN_FILE`` environment variable will be used. * ``assume_role_with_web_identity_federation_audience``: The ``aud`` claim of the access token, if using the ``google`` federation type. - The following extra parameters pass to :external:py:meth:`boto3.session.Session.client` + The following extra parameters are passed to :external:py:meth:`boto3.session.Session.client` or :external:py:meth:`boto3.session.Session.resource`. * ``config_kwargs``: Additional **kwargs** used to construct a @@ -126,7 +126,7 @@ Extra (optional) * ``endpoint_url``: Endpoint URL for the connection. * ``verify``: Whether or not to verify SSL certificates. -.. warning:: Extra parameters below are deprecated and will be removed in a future version of this provider. +.. warning:: The extra parameters below are deprecated and will be removed in a future version of this provider. * ``aws_account_id``: Used to construct ``role_arn`` if it was not specified. * ``aws_iam_role``: Used to construct ``role_arn`` if it was not specified. @@ -270,7 +270,18 @@ This assumes all other Connection fields eg **AWS Access Key ID** or **AWS Secre } } -4. Using AssumeRoleWithSAML +4. Using AssumeRoleWithWebIdentity (file-based token) + +.. code-block:: json + + { + "role_arn": "arn:aws:iam::112223334444:role/my_role", + "assume_role_method": "assume_role_with_web_identity", + "assume_role_with_web_identity_federation": "file", + "assume_role_with_web_identity_token_file": "/path/to/access_token" + } + +5. Using AssumeRoleWithSAML .. code-block:: json @@ -299,7 +310,7 @@ This assumes all other Connection fields eg **AWS Access Key ID** or **AWS Secre "assume_role_kwargs": { "something":"something" } } -5. Using **service_config** to specify configuration for services such as S3, STS, and EMR +6. Using **service_config** to specify configuration for services such as S3, STS, and EMR .. code-block:: json From b836e77a618cd9b21d740d4bcbe3120acc921e7f Mon Sep 17 00:00:00 2001 From: Philippe Gagnon <12717218+pgagnon@users.noreply.github.com> Date: Wed, 22 Feb 2023 17:28:49 -0500 Subject: [PATCH 8/9] Fixup doc --- docs/apache-airflow-providers-amazon/connections/aws.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/apache-airflow-providers-amazon/connections/aws.rst b/docs/apache-airflow-providers-amazon/connections/aws.rst index 2cccfff194fc1..85d45c62d6a50 100644 --- a/docs/apache-airflow-providers-amazon/connections/aws.rst +++ b/docs/apache-airflow-providers-amazon/connections/aws.rst @@ -348,7 +348,7 @@ S3 Bucket configurations To use S3 bucket name per connection in :class:`~airflow.providers.amazon.aws.hooks.s3.S3Hook` methods, provide selected options in the connection's extra field. -.. note:: The ``bucket_name`` parameter hook methods will override this connection setting. +.. note:: The ``bucket_name`` parameter in hook methods will override this connection setting. .. code-block:: json From c1d1dafaf62a4cef67072f594fb66a316f7f1252 Mon Sep 17 00:00:00 2001 From: Philippe Gagnon <12717218+pgagnon@users.noreply.github.com> Date: Tue, 28 Feb 2023 07:59:43 -0500 Subject: [PATCH 9/9] Update airflow/providers/amazon/aws/hooks/base_aws.py Add requested changes Co-authored-by: Ash Berlin-Taylor --- airflow/providers/amazon/aws/hooks/base_aws.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/airflow/providers/amazon/aws/hooks/base_aws.py b/airflow/providers/amazon/aws/hooks/base_aws.py index 936e1445c2441..ca60902fc075c 100644 --- a/airflow/providers/amazon/aws/hooks/base_aws.py +++ b/airflow/providers/amazon/aws/hooks/base_aws.py @@ -316,14 +316,14 @@ def _get_web_identity_credential_fetcher( web_identity_token_loader = { "file": self._get_file_token_loader, "google": self._get_google_identity_token_loader, - }.get(federation, lambda: None)() + }.get(federation) if not web_identity_token_loader: raise AirflowException(f"Unsupported federation: {federation}.") return botocore.credentials.AssumeRoleWithWebIdentityCredentialFetcher( client_creator=client_creator, - web_identity_token_loader=web_identity_token_loader, + web_identity_token_loader=web_identity_token_loader(), role_arn=self.role_arn, extra_args=self.conn.assume_role_kwargs, )