Skip to content

Commit

Permalink
Fix async conn for none aws_session_token (#30868)
Browse files Browse the repository at this point in the history
  • Loading branch information
pankajastro authored Apr 26, 2023
1 parent b335e47 commit a044a21
Show file tree
Hide file tree
Showing 2 changed files with 22 additions and 3 deletions.
6 changes: 3 additions & 3 deletions airflow/providers/amazon/aws/hooks/base_aws.py
Original file line number Diff line number Diff line change
Expand Up @@ -138,9 +138,9 @@ def _apply_session_kwargs(self, session):
or self.conn.session_kwargs.get("aws_session_token", None)
):
session.set_credentials(
self.conn.session_kwargs["aws_access_key_id"],
self.conn.session_kwargs["aws_secret_access_key"],
self.conn.session_kwargs["aws_session_token"],
access_key=self.conn.session_kwargs.get("aws_access_key_id"),
secret_key=self.conn.session_kwargs.get("aws_secret_access_key"),
token=self.conn.session_kwargs.get("aws_session_token"),
)

if self.conn.session_kwargs.get("region_name", None) is not None:
Expand Down
19 changes: 19 additions & 0 deletions tests/providers/amazon/aws/hooks/test_base_aws.py
Original file line number Diff line number Diff line change
Expand Up @@ -246,6 +246,25 @@ def test_async_create_session_from_credentials(self, region_name, profile_name):

assert session_profile == profile_name

@pytest.mark.asyncio
async def test_async_create_a_session_from_credentials_without_token(self):
mock_conn = Connection(
conn_type=MOCK_CONN_TYPE,
conn_id=MOCK_AWS_CONN_ID,
extra={
"aws_access_key_id": "test_aws_access_key_id",
"aws_secret_access_key": "test_aws_secret_access_key",
"region_name": "eu-central-1",
},
)
mock_conn_config = AwsConnectionWrapper(conn=mock_conn)
sf = BaseSessionFactory(conn=mock_conn_config, config=None)
async_session = sf.create_session(deferrable=True)
cred = await async_session.get_credentials()
assert cred.access_key == "test_aws_access_key_id"
assert cred.secret_key == "test_aws_secret_access_key"
assert cred.token is None

config_for_credentials_test = [
(
"assume-with-initial-creds",
Expand Down

0 comments on commit a044a21

Please sign in to comment.