Skip to content

Commit

Permalink
Revert "Adjust built-in base_aws methods to avoid Deprecation warnings (
Browse files Browse the repository at this point in the history
#19725)" (#19791)

This reverts commit 4be0414.
  • Loading branch information
uranusjr authored Nov 23, 2021
1 parent 1f0bb04 commit d58df46
Show file tree
Hide file tree
Showing 4 changed files with 7 additions and 14 deletions.
8 changes: 3 additions & 5 deletions airflow/providers/amazon/aws/hooks/base_aws.py
Original file line number Diff line number Diff line change
Expand Up @@ -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!')
Expand Down Expand Up @@ -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]):
Expand Down
3 changes: 1 addition & 2 deletions airflow/providers/amazon/aws/hooks/glue.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
7 changes: 2 additions & 5 deletions airflow/providers/amazon/aws/hooks/s3.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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

Expand Down
3 changes: 1 addition & 2 deletions airflow/providers/postgres/hooks/postgres.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down

0 comments on commit d58df46

Please sign in to comment.