Skip to content

Commit

Permalink
Import aiobotocore only if deferrable is true
Browse files Browse the repository at this point in the history
We should import aiobotocore only if deferrable is set to true. currently, it is failing if a task uses assume role.
  • Loading branch information
pankajastro committed May 5, 2023
1 parent a6be96d commit 582c37f
Showing 1 changed file with 6 additions and 2 deletions.
8 changes: 6 additions & 2 deletions airflow/providers/amazon/aws/hooks/base_aws.py
Original file line number Diff line number Diff line change
Expand Up @@ -194,7 +194,6 @@ 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:
from aiobotocore.session import get_session as async_get_session

if self.conn.assume_role_method == "assume_role_with_web_identity":
# Deferred credentials have no initial credentials
Expand All @@ -212,7 +211,12 @@ def _create_session_with_assume_role(
method="sts-assume-role",
)

session = async_get_session() if deferrable else botocore.session.get_session()
if deferrable:
from aiobotocore.session import get_session as async_get_session

session = async_get_session()
else:
session = botocore.session.get_session()

session._credentials = credentials
session.set_config_variable("region", self.basic_session.region_name)
Expand Down

0 comments on commit 582c37f

Please sign in to comment.