Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

FIX AWS deferrable operators by using AioCredentials when using assume_role #32733

Merged
merged 5 commits into from
Jul 22, 2023
Merged

FIX AWS deferrable operators by using AioCredentials when using assume_role #32733

Changes from 2 commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
fbf6b15
FIX: deferrable operators now use AioCredentials
rishi-kulkarni Jul 21, 2023
86ea3a6
formatting
rishi-kulkarni Jul 21, 2023
efcecdd
use dict unpacking
rishi-kulkarni Jul 21, 2023
38a4fd2
TEST: add test that checks that credentials._refresh is a coroutine
rishi-kulkarni Jul 21, 2023
64cbc53
TEST: assert that `credentials.get_frozen_credentials` is a coroutine
rishi-kulkarni Jul 21, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
41 changes: 30 additions & 11 deletions airflow/providers/amazon/aws/hooks/base_aws.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -197,22 +197,41 @@ def _create_basic_session(self, session_kwargs: dict[str, Any]) -> boto3.session
def _create_session_with_assume_role(
self, session_kwargs: dict[str, Any], deferrable: bool = False
) -> boto3.session.Session:

if self.conn.assume_role_method == "assume_role_with_web_identity":
# Deferred credentials have no initial credentials
credential_fetcher = self._get_web_identity_credential_fetcher()
credentials = botocore.credentials.DeferredRefreshableCredentials(
method="assume-role-with-web-identity",
refresh_using=credential_fetcher.fetch_credentials,
time_fetcher=lambda: datetime.datetime.now(tz=tzlocal()),
)

if deferrable:
from aiobotocore.credentials import AioDeferredRefreshableCredentials

credentials = AioDeferredRefreshableCredentials(
method="assume-role-with-web-identity",
refresh_using=credential_fetcher.fetch_credentials,
time_fetcher=lambda: datetime.datetime.now(tz=tzlocal()),
rishi-kulkarni marked this conversation as resolved.
Show resolved Hide resolved
)
else:
credentials = botocore.credentials.DeferredRefreshableCredentials(
method="assume-role-with-web-identity",
refresh_using=credential_fetcher.fetch_credentials,
time_fetcher=lambda: datetime.datetime.now(tz=tzlocal()),
)
else:
# Refreshable credentials do have initial credentials
credentials = botocore.credentials.RefreshableCredentials.create_from_metadata(
metadata=self._refresh_credentials(),
refresh_using=self._refresh_credentials,
method="sts-assume-role",
)

if deferrable:
from aiobotocore.credentials import AioRefreshableCredentials

credentials = AioRefreshableCredentials.create_from_metadata(
metadata=self._refresh_credentials(),
refresh_using=self._refresh_credentials,
method="sts-assume-role",
)
else:
credentials = botocore.credentials.RefreshableCredentials.create_from_metadata(
metadata=self._refresh_credentials(),
refresh_using=self._refresh_credentials,
method="sts-assume-role",
)

if deferrable:
from aiobotocore.session import get_session as async_get_session
Expand Down