diff --git a/airflow/providers/amazon/aws/hooks/glue.py b/airflow/providers/amazon/aws/hooks/glue.py index 7879ba7fe0849..d7f538fb5d755 100644 --- a/airflow/providers/amazon/aws/hooks/glue.py +++ b/airflow/providers/amazon/aws/hooks/glue.py @@ -17,6 +17,7 @@ # under the License. import time +from typing import Dict, List, Optional from airflow.exceptions import AirflowException from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook @@ -26,6 +27,8 @@ class AwsGlueJobHook(AwsBaseHook): """ Interact with AWS Glue - create job, trigger, crawler + :param s3_bucket: S3 bucket where logs and local etl script will be uploaded + :type s3_bucket: Optional[str] :param job_name: unique job name per AWS account :type job_name: Optional[str] :param desc: job description @@ -42,21 +45,19 @@ class AwsGlueJobHook(AwsBaseHook): :type region_name: Optional[str] :param iam_role_name: AWS IAM Role for Glue Job :type iam_role_name: Optional[str] - :param s3_bucket: S3 bucket where logs and local etl script will be uploaded - :type s3_bucket: Optional[str] """ JOB_POLL_INTERVAL = 6 # polls job status after every JOB_POLL_INTERVAL seconds def __init__(self, - s3_bucket=None, - job_name=None, - desc=None, - concurrent_run_limit=1, - script_location=None, - retry_limit=0, - num_of_dpus=10, - region_name=None, - iam_role_name=None, + s3_bucket: Optional[str] = None, + job_name: Optional[str] = None, + desc: Optional[str] = None, + concurrent_run_limit: int = 1, + script_location: Optional[str] = None, + retry_limit: int = 0, + num_of_dpus: int = 10, + region_name: Optional[str] = None, + iam_role_name: Optional[str] = None, *args, **kwargs): self.job_name = job_name self.desc = desc @@ -68,16 +69,17 @@ def __init__(self, self.s3_bucket = s3_bucket self.role_name = iam_role_name self.s3_glue_logs = 'logs/glue-logs/' - super(AwsGlueJobHook, self).__init__(client_type='glue', *args, **kwargs) + kwargs['client_type'] = 'glue' + super().__init__(*args, **kwargs) - def list_jobs(self): + def list_jobs(self) -> List: """ :return: Lists of Jobs """ conn = self.get_conn() return conn.get_jobs() - def get_iam_execution_role(self): + def get_iam_execution_role(self) -> Dict: """ :return: iam role for job execution """ @@ -91,13 +93,14 @@ def get_iam_execution_role(self): self.log.error("Failed to create aws glue job, error: %s", general_error) raise - def initialize_job(self, script_arguments=None): + def initialize_job(self, script_arguments: Optional[List] = None) -> Dict[str, str]: """ Initializes connection with AWS Glue to run job :return: """ glue_client = self.get_conn() + script_arguments = script_arguments or [] try: job_name = self.get_or_create_glue_job() @@ -110,7 +113,7 @@ def initialize_job(self, script_arguments=None): self.log.error("Failed to run aws glue job, error: %s", general_error) raise - def job_completion(self, job_name=None, run_id=None): + def job_completion(self, job_name: str, run_id: str) -> Dict[str, str]: """ :param job_name: unique job name per AWS account :type job_name: str @@ -141,7 +144,7 @@ def job_completion(self, job_name=None, run_id=None): job_name, job_run_state) time.sleep(self.JOB_POLL_INTERVAL) - def get_or_create_glue_job(self): + def get_or_create_glue_job(self) -> str: """ Creates(or just returns) and returns the Job name :return:Name of the Job diff --git a/airflow/providers/amazon/aws/hooks/redshift.py b/airflow/providers/amazon/aws/hooks/redshift.py index 54732941863df..57f59c86ee85d 100644 --- a/airflow/providers/amazon/aws/hooks/redshift.py +++ b/airflow/providers/amazon/aws/hooks/redshift.py @@ -19,6 +19,8 @@ Interact with AWS Redshift, using the boto3 library. """ +from typing import List, Optional + from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook @@ -37,7 +39,7 @@ def __init__(self, *args, **kwargs): super().__init__(client_type='redshift', *args, **kwargs) # TODO: Wrap create_cluster_snapshot - def cluster_status(self, cluster_identifier): + def cluster_status(self, cluster_identifier: str) -> str: """ Return status of a cluster @@ -53,9 +55,9 @@ def cluster_status(self, cluster_identifier): def delete_cluster( # pylint: disable=invalid-name self, - cluster_identifier, - skip_final_cluster_snapshot=True, - final_cluster_snapshot_identifier=''): + cluster_identifier: str, + skip_final_cluster_snapshot: bool = True, + final_cluster_snapshot_identifier: Optional[str] = None): """ Delete a cluster and optionally create a snapshot @@ -66,6 +68,8 @@ def delete_cluster( # pylint: disable=invalid-name :param final_cluster_snapshot_identifier: name of final cluster snapshot :type final_cluster_snapshot_identifier: str """ + final_cluster_snapshot_identifier = final_cluster_snapshot_identifier or '' + response = self.get_conn().delete_cluster( ClusterIdentifier=cluster_identifier, SkipFinalClusterSnapshot=skip_final_cluster_snapshot, @@ -73,7 +77,7 @@ def delete_cluster( # pylint: disable=invalid-name ) return response['Cluster'] if response['Cluster'] else None - def describe_cluster_snapshots(self, cluster_identifier): + def describe_cluster_snapshots(self, cluster_identifier: str) -> Optional[List[str]]: """ Gets a list of snapshots for a cluster @@ -90,7 +94,10 @@ def describe_cluster_snapshots(self, cluster_identifier): snapshots.sort(key=lambda x: x['SnapshotCreateTime'], reverse=True) return snapshots - def restore_from_cluster_snapshot(self, cluster_identifier, snapshot_identifier): + def restore_from_cluster_snapshot( + self, + cluster_identifier: str, + snapshot_identifier: str) -> str: """ Restores a cluster from its snapshot @@ -105,7 +112,10 @@ def restore_from_cluster_snapshot(self, cluster_identifier, snapshot_identifier) ) return response['Cluster'] if response['Cluster'] else None - def create_cluster_snapshot(self, snapshot_identifier, cluster_identifier): + def create_cluster_snapshot( + self, + snapshot_identifier: str, + cluster_identifier: str) -> str: """ Creates a snapshot of a cluster diff --git a/airflow/providers/amazon/aws/hooks/s3.py b/airflow/providers/amazon/aws/hooks/s3.py index 710a4bc123f5d..976e7ebd8575b 100644 --- a/airflow/providers/amazon/aws/hooks/s3.py +++ b/airflow/providers/amazon/aws/hooks/s3.py @@ -106,8 +106,9 @@ class S3Hook(AwsBaseHook): :class:`~airflow.providers.amazon.aws.hooks.base_aws.AwsBaseHook` """ - def __init__(self, *args, **kwargs): - super().__init__(client_type='s3', *args, **kwargs) + def __init__(self, *args, **kwargs) -> None: + kwargs['client_type'] = 's3' + super().__init__(*args, **kwargs) @staticmethod def parse_s3_url(s3url: str) -> Tuple[str, str]: