AWS Hook - allow IDP HTTP retry (#12639)

[baolsen]

When doing AssumeRoleWithSAML there is an HTTP request made to an IDP endpoint to obtain the SAML Assertion. This PR allows the IDP request to have a configurable Retry, which makes this authentication mechanism a bit more robust.

I'm not sure how to write tests for this.... suggestions welcome.

Tested with my own Airflow environment which uses this auth method, and it is working as intended.
It's working with and without the new configuration.

closes: 12639

[mik-laj]

I'm not sure how to write tests for this.... suggestions welcome.

You can use mock.mock to check if all methods are properly called.
See example:

airflow/tests/providers/amazon/aws/hooks/test_base_aws.py

Lines 179 to 253 in a697c58

	def test_get_credentials_from_gcp_credentials(self):
	mock_connection = Connection(
	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_federation": 'google',
	"assume_role_with_web_identity_federation_audience": 'aws-federation.airflow.apache.org',
	}
	)
	)
	# Store original __import__
	orig_import = __import__
	mock_id_token_credentials = mock.Mock()
	def import_mock(name, *args):
	if name == 'airflow.providers.google.common.utils.id_token_credentials':
	return mock_id_token_credentials
	return orig_import(name, *args)
	with mock.patch('builtins.__import__', side_effect=import_mock), mock.patch.dict(
	'os.environ', AIRFLOW_CONN_AWS_DEFAULT=mock_connection.get_uri()
	), mock.patch('airflow.providers.amazon.aws.hooks.base_aws.boto3') as mock_boto3, mock.patch(
	'airflow.providers.amazon.aws.hooks.base_aws.botocore'
	) as mock_botocore, mock.patch(
	'airflow.providers.amazon.aws.hooks.base_aws.botocore.session'
	) as mock_session:
	hook = AwsBaseHook(aws_conn_id='aws_default', client_type='airflow_test')
	credentials_from_hook = hook.get_credentials()
	mock_get_credentials = mock_boto3.session.Session.return_value.get_credentials
	self.assertEqual(
	mock_get_credentials.return_value.get_frozen_credentials.return_value,
	credentials_from_hook,
	)
	mock_boto3.assert_has_calls(
	[
	mock.call.session.Session(
	aws_access_key_id=None,
	aws_secret_access_key=None,
	aws_session_token=None,
	region_name=None,
	),
	mock.call.session.Session()._session.__bool__(),
	mock.call.session.Session(
	botocore_session=mock_session.Session.return_value,
	region_name=mock_boto3.session.Session.return_value.region_name,
	),
	mock.call.session.Session().get_credentials(),
	mock.call.session.Session().get_credentials().get_frozen_credentials(),
	]
	)
	mock_fetcher = mock_botocore.credentials.AssumeRoleWithWebIdentityCredentialFetcher
	mock_botocore.assert_has_calls(
	[
	mock.call.credentials.AssumeRoleWithWebIdentityCredentialFetcher(
	client_creator=mock_boto3.session.Session.return_value._session.create_client,
	extra_args={},
	role_arn='arn:aws:iam::123456:role/role_arn',
	web_identity_token_loader=mock.ANY,
	),
	mock.call.credentials.DeferredRefreshableCredentials(
	method='assume-role-with-web-identity',
	refresh_using=mock_fetcher.return_value.fetch_credentials,
	time_fetcher=mock.ANY,
	),
	]
	)
	mock_session.assert_has_calls([mock.call.Session()])
	mock_id_token_credentials.assert_has_calls(
	[mock.call.get_default_id_token_credentials(target_audience='aws-federation.airflow.apache.org')]
	)

[baolsen]

Didn't think of that, thanks for the suggestion @mik-laj . Will add something when I can :)
We are going to take the code through our environments meanwhile, will post results here once done.

[github-actions]

The Workflow run is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks,^Build docs$,^Spell check docs$,^Backport packages$,^Provider packages,^Checks: Helm tests$,^Test OpenAPI*.

[baolsen]

Hey @mik-laj , I think this one is ready for review.

[github-actions]

The PR is likely OK to be merged with just subset of tests for default Python and Database versions without running the full matrix of tests, because it does not modify the core of Airflow. If the committers decide that the full tests matrix is needed, they will add the label 'full tests needed'. Then you should rebase to the latest master or amend the last commit of the PR, and push it with --force-with-lease.

[baolsen]

Hey @mik-laj this one is good to merge

[mik-laj]

The title of this PR uses the word "Retry", but I don't see this word appearing in this code. Does this code test the same functionality?

.

[github-actions]

This pull request has been automatically marked as stale because it has not had recent activity. It will be closed in 5 days if no further activity occurs. Thank you for your contributions.

[baolsen]

Hey @mik-laj , trust all is well that side.
Could we re-open this PR please? I'll need to do some rebasing etc of course.
I finally have some capacity to work on Airflow again :) and would like to get some of my issues sorted.

[mik-laj]

I reopened this PR.

[baolsen]

Thank you @mik-laj.
I see it is merging to apache:master, I can't edit it to target apache:main (can only change PR description at this point).
Do you think I should rather close this and open a new one with the correct target?

[mik-laj]

Yes. It is good idea. You can do it.

[baolsen]

Will do, thank you