Adjust built-in base_aws methods to avoid Deprecation warnings (#19725)

dimon222 · uranusjr · web-flow · commit 4be04143a5f7 · 2021-11-24T03:41:56.000+08:00
Co-authored-by: Tzu-ping Chung &lt;uranusjr@gmail.com&gt;
diff --git 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(self.client_type, region_name=self.region_name)
+            return self.get_client_type(region_name=self.region_name)
         elif self.resource_type:
-            return self.get_resource_type(self.resource_type, region_name=self.region_name)
+            return self.get_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,7 +539,9 @@ def expand_role(self, role: str) -> str:
         if "/" in role:
             return role
         else:
-            return self.get_client_type("iam").get_role(RoleName=role)["Role"]["Arn"]
+            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"]
 
     @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
@@ -85,7 +85,8 @@ def list_jobs(self) -> List:
 
     def get_iam_execution_role(self) -> Dict:
         """:return: iam role for job execution"""
-        iam_client = self.get_client_type('iam', self.region_name)
+        session, endpoint_url = self._get_credentials(self.region_name)
+        iam_client = session.client('iam', endpoint_url=endpoint_url, config=self.config, verify=self.verify)
 
         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
@@ -173,7 +173,8 @@ def get_bucket(self, bucket_name: Optional[str] = None) -> str:
         :return: the bucket object to the bucket name.
         :rtype: boto3.S3.Bucket
         """
-        s3_resource = self.get_resource_type('s3')
+        session, endpoint_url = self._get_credentials()
+        s3_resource = session.client('s3', endpoint_url=endpoint_url, config=self.config, verify=self.verify)
         return s3_resource.Bucket(bucket_name)
 
     @provide_bucket_name
@@ -340,7 +341,9 @@ def get_key(self, key: str, bucket_name: Optional[str] = None) -> S3Transfer:
         :return: the key object from the bucket
         :rtype: boto3.s3.Object
         """
-        obj = self.get_resource_type('s3').Object(bucket_name, key)
+        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.load()
         return obj
 
diff --git a/airflow/providers/postgres/hooks/postgres.py b/airflow/providers/postgres/hooks/postgres.py
@@ -184,7 +184,8 @@ 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])
-            client = aws_hook.get_client_type('redshift')
+            session, endpoint_url = aws_hook._get_credentials()
+            client = session.client('redshift', endpoint_url=endpoint_url, config=aws_hook.config, verify=aws_hook.verify)
             cluster_creds = client.get_cluster_credentials(
                 DbUser=conn.login,
                 DbName=self.schema or conn.schema,
