Skip to content

Commit

Permalink
FIX: deferrable operators now use AioCredentials
Browse files Browse the repository at this point in the history
Fixes airflow.providers.amazong.aws.hooks.base_aws.BaseSessionFactory feeds synchronous credentials to aiobotocore when using `assume_role` apache#32732
  • Loading branch information
rishi-kulkarni committed Jul 21, 2023
1 parent 15d42b4 commit fbf6b15
Showing 1 changed file with 26 additions and 6 deletions.
32 changes: 26 additions & 6 deletions airflow/providers/amazon/aws/hooks/base_aws.py
Original file line number Diff line number Diff line change
Expand Up @@ -201,18 +201,38 @@ def _create_session_with_assume_role(
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(

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()),
)
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

0 comments on commit fbf6b15

Please sign in to comment.