From d58df468c8d77c5d45e80f2333eb074bb7771a95 Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Wed, 24 Nov 2021 07:45:32 +0800 Subject: [PATCH] Revert "Adjust built-in base_aws methods to avoid Deprecation warnings (#19725)" (#19791) This reverts commit 4be04143a5f7e246127e942bf1d73abcd22ce189. --- airflow/providers/amazon/aws/hooks/base_aws.py | 8 +++----- airflow/providers/amazon/aws/hooks/glue.py | 3 +-- airflow/providers/amazon/aws/hooks/s3.py | 7 ++----- airflow/providers/postgres/hooks/postgres.py | 3 +-- 4 files changed, 7 insertions(+), 14 deletions(-) diff --git a/airflow/providers/amazon/aws/hooks/base_aws.py b/airflow/providers/amazon/aws/hooks/base_aws.py index 90b14733abb39..171437cc740c5 100644 --- a/airflow/providers/amazon/aws/hooks/base_aws.py +++ b/airflow/providers/amazon/aws/hooks/base_aws.py @@ -491,9 +491,9 @@ def conn(self) -> Union[boto3.client, boto3.resource]: :rtype: Union[boto3.client, boto3.resource] """ if self.client_type: - return self.get_client_type(region_name=self.region_name) + return self.get_client_type(self.client_type, region_name=self.region_name) elif self.resource_type: - return self.get_resource_type(region_name=self.region_name) + return self.get_resource_type(self.resource_type, region_name=self.region_name) else: # Rare possibility - subclasses have not specified a client_type or resource_type raise NotImplementedError('Could not get boto3 connection!') @@ -539,9 +539,7 @@ def expand_role(self, role: str) -> str: if "/" in role: return role else: - session, endpoint_url = self._get_credentials() - _client = session.client('iam', endpoint_url=endpoint_url, config=self.config, verify=self.verify) - return _client.get_role(RoleName=role)["Role"]["Arn"] + return self.get_client_type("iam").get_role(RoleName=role)["Role"]["Arn"] @staticmethod def retry(should_retry: Callable[[Exception], bool]): diff --git a/airflow/providers/amazon/aws/hooks/glue.py b/airflow/providers/amazon/aws/hooks/glue.py index 2cf048000aa4c..a5e43278dc35a 100644 --- a/airflow/providers/amazon/aws/hooks/glue.py +++ b/airflow/providers/amazon/aws/hooks/glue.py @@ -85,8 +85,7 @@ def list_jobs(self) -> List: def get_iam_execution_role(self) -> Dict: """:return: iam role for job execution""" - session, endpoint_url = self._get_credentials(self.region_name) - iam_client = session.client('iam', endpoint_url=endpoint_url, config=self.config, verify=self.verify) + iam_client = self.get_client_type('iam', self.region_name) try: glue_execution_role = iam_client.get_role(RoleName=self.role_name) diff --git a/airflow/providers/amazon/aws/hooks/s3.py b/airflow/providers/amazon/aws/hooks/s3.py index 569cbade3e18c..53b2747104380 100644 --- a/airflow/providers/amazon/aws/hooks/s3.py +++ b/airflow/providers/amazon/aws/hooks/s3.py @@ -173,8 +173,7 @@ def get_bucket(self, bucket_name: Optional[str] = None) -> str: :return: the bucket object to the bucket name. :rtype: boto3.S3.Bucket """ - session, endpoint_url = self._get_credentials() - s3_resource = session.client('s3', endpoint_url=endpoint_url, config=self.config, verify=self.verify) + s3_resource = self.get_resource_type('s3') return s3_resource.Bucket(bucket_name) @provide_bucket_name @@ -341,9 +340,7 @@ def get_key(self, key: str, bucket_name: Optional[str] = None) -> S3Transfer: :return: the key object from the bucket :rtype: boto3.s3.Object """ - session, endpoint_url = self._get_credentials() - s3_resource = session.client('s3', endpoint_url=endpoint_url, config=self.config, verify=self.verify) - obj = s3_resource.Object(bucket_name, key) + obj = self.get_resource_type('s3').Object(bucket_name, key) obj.load() return obj diff --git a/airflow/providers/postgres/hooks/postgres.py b/airflow/providers/postgres/hooks/postgres.py index a68046fda03ae..93688d74cd67c 100644 --- a/airflow/providers/postgres/hooks/postgres.py +++ b/airflow/providers/postgres/hooks/postgres.py @@ -184,8 +184,7 @@ def get_iam_token(self, conn: Connection) -> Tuple[str, str, int]: # Pull the custer-identifier from the beginning of the Redshift URL # ex. my-cluster.ccdre4hpd39h.us-east-1.redshift.amazonaws.com returns my-cluster cluster_identifier = conn.extra_dejson.get('cluster-identifier', conn.host.split('.')[0]) - session, endpoint_url = aws_hook._get_credentials() - client = session.client('redshift', endpoint_url=endpoint_url, config=aws_hook.config, verify=aws_hook.verify) + client = aws_hook.get_client_type('redshift') cluster_creds = client.get_cluster_credentials( DbUser=conn.login, DbName=self.schema or conn.schema,