dazza-codes commented on a change in pull request #6811: [RFC][AIRFLOW-6245] Add custom waiters for AWS batch jobs URL: https://github.com/apache/airflow/pull/6811#discussion_r360728327
########## File path: airflow/providers/amazon/aws/hooks/batch_waiters.py ########## @@ -0,0 +1,210 @@ +# -*- coding: utf-8 -*- +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# + +""" +AWS batch service waiters + +.. seealso:: + + - https://boto3.amazonaws.com/v1/documentation/api/latest/guide/clients.html#waiters + - https://github.com/boto/botocore/blob/develop/botocore/waiter.py +""" + +import json +import sys +from copy import deepcopy +from pathlib import Path + +import botocore.client +import botocore.exceptions +import botocore.waiter + +from airflow import AirflowException +from airflow.providers.amazon.aws.hooks.batch_client import AwsBatchClient + + +class AwsBatchWaiters(AwsBatchClient): + """ + A utility to manage waiters for AWS batch services. + + :param waiter_config: a custom waiter configuration for AWS batch services + :type waiter_config: Optional[dict] + + :param aws_conn_id: connection id of AWS credentials / region name. If None, + credential boto3 strategy will be used + (http://boto3.readthedocs.io/en/latest/guide/configuration.html). + :type aws_conn_id: Optional[str] + + :param region_name: region name to use in AWS client. + Override the AWS region in connection (if provided) + :type region_name: Optional[str] + + Examples: + + .. code-block:: python + + from airflow.providers.amazon.aws.operators.batch_waiters import AwsBatchWaiters + + # to inspect default waiters + waiters = AwsBatchWaiters() + config = waiters.default_config # type: dict + waiter_names = waiters.list_waiters() # -> ["JobComplete", "JobExists", "JobRunning"] + + # the default_config is a useful stepping stone to creating custom waiters; + # to customize waiters (modify `config`) before any calls to `waiters.get_waiter()` + custom_config = waiters.default_config # this is a deep copy of the default_config + # modify custom_config['waiters'] as necessary + waiters = AwsBatchWaiters(waiter_config=custom_config) + waiters.waiter_config # check the custom configuration + waiters.list_waiters() # names of custom waiters + + # WARNING: the preferred way to customize waiters is using the approach above, + # because the `waiters.waiter_config` is a mutable dict; e.g. don't do + # this at home: + waiters.waiter_config['version'] # -> 2 + config = waiters.waiter_config + config['version'] = 3 # a change like this will likely break botocore validation + waiters.waiter_config['version'] # -> 3 + + # During the init for AwsBatchWaiters, the waiter_config is used to build a waiter_model; + # and note that this only occurs during the class init, to avoid any accidental mutations + # of waiter_config leaking into the waiter_model. + waiters.waiter_model # -> botocore.waiter.WaiterModel object + + # the waiter_model is combined with the waiters.client to get a specific waiter + # and the details of the config on that waiter can be further modified without any + # accidental impact on the generation of new waiters from the defined waiter_model, e.g. + waiters.get_waiter("JobExists").config.delay # -> 5 + waiter = waiters.get_waiter("JobExists") # -> botocore.waiter.Batch.Waiter.JobExists object + waiter.config.delay = 10 + waiters.get_waiter("JobExists").config.delay # -> 5 - defined by waiter_model + + # to use a specific waiter, update the config and call the `wait()` method for jobId, e.g. + waiter = waiters.get_waiter("JobExists") # -> botocore.waiter.Batch.Waiter.JobExists object + waiter.config.delay = randint(1, 10) # seconds + waiter.config.max_attempts = 10 + waiter.wait(jobs=[jobId]) + + .. seealso:: + + - https://www.2ndwatch.com/blog/use-waiters-boto3-write/ + - https://github.com/boto/botocore/blob/develop/botocore/waiter.py + - https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/ec2.html#waiters + - https://github.com/boto/botocore/tree/develop/botocore/data/ec2/2016-11-15 + - https://github.com/boto/botocore/issues/1915 + """ + + def __init__(self, waiter_config=None, aws_conn_id=None, region_name=None): + + AwsBatchClient.__init__(self, aws_conn_id=aws_conn_id, region_name=region_name) + + self._default_config = None + self._waiter_config = waiter_config or self.default_config + self._waiter_model = botocore.waiter.WaiterModel(self._waiter_config) + + @property + def default_config(self) -> dict: + """An immutable default waiter configuration + :return: a waiter configuration for AWS batch services + :rtype: dict + """ + if self._default_config is None: + config_path = Path(__file__).with_name("batch_waiters.json").absolute() + with open(config_path) as config_file: + self._default_config = json.load(config_file) + # return a deep copy, to avoid accidental mutation + return deepcopy(self._default_config) + + @property + def waiter_config(self): + """ + :return: a mutable waiter configuration for AWS batch services + :rtype: dict + """ + return self._waiter_config + + @property + def waiter_model(self): + """ + :return: a waiter model for AWS batch services + :rtype: botocore.waiter.WaiterModel + """ + return self._waiter_model + + def get_waiter(self, waiter_name): + """ + Get an AWS Batch service waiter + + :param waiter_name: The name of the waiter. The name should match Review comment: OK, this is likely a remnant of my academic education in writing science articles. ---------------------------------------------------------------- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. For queries about this service, please contact Infrastructure at: us...@infra.apache.org With regards, Apache Git Services
