You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by el...@apache.org on 2023/12/04 19:07:12 UTC
(airflow) branch main updated: Fix a bug in get_iam_token for Redshift Serverless (#36001)
This is an automated email from the ASF dual-hosted git repository.
eladkal pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new e0df7441fa Fix a bug in get_iam_token for Redshift Serverless (#36001)
e0df7441fa is described below
commit e0df7441fa607645d0a379c2066ca4ab16f5cb95
Author: Hussein Awala <hu...@awala.fr>
AuthorDate: Mon Dec 4 21:07:03 2023 +0200
Fix a bug in get_iam_token for Redshift Serverless (#36001)
---
airflow/providers/amazon/aws/hooks/redshift_sql.py | 12 +++++++-----
tests/providers/amazon/aws/hooks/test_redshift_sql.py | 10 +++++-----
2 files changed, 12 insertions(+), 10 deletions(-)
diff --git a/airflow/providers/amazon/aws/hooks/redshift_sql.py b/airflow/providers/amazon/aws/hooks/redshift_sql.py
index 87bc6be478..8ce97f26bd 100644
--- a/airflow/providers/amazon/aws/hooks/redshift_sql.py
+++ b/airflow/providers/amazon/aws/hooks/redshift_sql.py
@@ -113,15 +113,17 @@ class RedshiftSQLHook(DbApiHook):
serverless_token_duration_seconds = conn.extra_dejson.get(
"serverless_token_duration_seconds", 3600
)
- redshift_client = AwsBaseHook(
+ redshift_serverless_client = AwsBaseHook(
aws_conn_id=self.aws_conn_id, client_type="redshift-serverless"
).conn
# https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/redshift-serverless/client/get_credentials.html#get-credentials
- cluster_creds = redshift_client.get_cluster_credentials(
- DbName=conn.schema,
+ cluster_creds = redshift_serverless_client.get_credentials(
+ dbName=conn.schema,
workgroupName=serverless_work_group,
durationSeconds=serverless_token_duration_seconds,
)
+ token = cluster_creds["dbPassword"]
+ login = cluster_creds["dbUser"]
else:
# Pull the custer-identifier from the beginning of the Redshift URL
# ex. my-cluster.ccdre4hpd39h.us-east-1.redshift.amazonaws.com returns my-cluster
@@ -139,8 +141,8 @@ class RedshiftSQLHook(DbApiHook):
ClusterIdentifier=cluster_identifier,
AutoCreate=False,
)
- token = cluster_creds["DbPassword"]
- login = cluster_creds["DbUser"]
+ token = cluster_creds["DbPassword"]
+ login = cluster_creds["DbUser"]
return login, token, port
def get_uri(self) -> str:
diff --git a/tests/providers/amazon/aws/hooks/test_redshift_sql.py b/tests/providers/amazon/aws/hooks/test_redshift_sql.py
index cab6053eff..4871522489 100644
--- a/tests/providers/amazon/aws/hooks/test_redshift_sql.py
+++ b/tests/providers/amazon/aws/hooks/test_redshift_sql.py
@@ -139,16 +139,16 @@ class TestRedshiftSQLHookConn:
mock_db_pass = "aws_token"
# Mock AWS Connection
- mock_aws_hook_conn.get_cluster_credentials.return_value = {
- "DbPassword": mock_db_pass,
- "DbUser": mock_db_user,
+ mock_aws_hook_conn.get_credentials.return_value = {
+ "dbPassword": mock_db_pass,
+ "dbUser": mock_db_user,
}
self.db_hook.get_conn()
# Check boto3 'redshift' client method `get_cluster_credentials` call args
- mock_aws_hook_conn.get_cluster_credentials.assert_called_once_with(
- DbName=LOGIN_SCHEMA,
+ mock_aws_hook_conn.get_credentials.assert_called_once_with(
+ dbName=LOGIN_SCHEMA,
workgroupName=mock_work_group,
durationSeconds=3600,
)