[GitHub] [airflow] RosterIn commented on a change in pull request #9011: Changing from Resource to Client for EC2 and adding support for filters

[GitBox <gi...@apache.org>]

RosterIn commented on a change in pull request #9011:
URL: https://github.com/apache/airflow/pull/9011#discussion_r547905265



##########
File path: airflow/providers/amazon/aws/hooks/ec2.py
##########
@@ -33,20 +33,82 @@ class EC2Hook(AwsBaseHook):
         :class:`~airflow.providers.amazon.aws.hooks.base_aws.AwsBaseHook`
     """
 
-    def __init__(self, *args, **kwargs) -> None:
-        kwargs["resource_type"] = "ec2"
-        super().__init__(*args, **kwargs)
+    def __init__(self, *args, **kwargs):
+        super().__init__(client_type="ec2", *args, **kwargs)
 
-    def get_instance(self, instance_id: str):
+    def stop_instances(self, instance_ids: list) -> dict:
         """
-        Get EC2 instance by id and return it.
+        Stop instances with given ids
 
-        :param instance_id: id of the AWS EC2 instance
-        :type instance_id: str
-        :return: Instance object
-        :rtype: ec2.Instance
+        :param instance_ids: List of instance ids to stop
+        :return: Dict with key `StoppingInstances` and value as list of instances being stopped
+        """
+        self.log.info("Stopping instances: %s", instance_ids)
+
+        return self.conn.stop_instances(InstanceIds=instance_ids)
+
+    def start_instances(self, instance_ids: list) -> dict:
+        """
+        Start instances with given ids
+
+        :param instance_ids: List of instance ids to start
+        :return: Dict with key `StartingInstances` and value as list of instances being started
+        """
+        self.log.info("Starting instances: %s", instance_ids)
+
+        return self.conn.start_instances(InstanceIds=instance_ids)
+
+    def terminate_instances(self, instance_ids: list) -> dict:
+        """
+        Terminate instances with given ids
+
+        :param instance_ids: List of instance ids to terminate
+        :return: Dict with key `TerminatingInstances` and value as list of instances being terminated
         """
-        return self.conn.Instance(id=instance_id)
+        self.log.info("Terminating instances: %s", instance_ids)
+
+        return self.conn.terminate_instances(InstanceIds=instance_ids)
+
+    def describe_instances(self, filters: list = None, instance_ids: list = None):
+        """
+        Describe EC2 instances, optionally applying filters and selective instance ids
+
+        :param filters: List of filters to specify instances to describe
+        :param instance_ids: List of instance IDs to describe
+        :return: Response from EC2 describe_instances API
+        """
+        filters = [] if filters is None else filters
+        instance_ids = [] if instance_ids is None else instance_ids
+
+        self.log.info("Filters provided: %s", filters)
+        self.log.info("Instance ids provided: %s", instance_ids)
+
+        return self.conn.describe_instances(Filters=filters, InstanceIds=instance_ids)
+
+    def get_instances(self, filters: list = None, instance_ids: list = None) -> list:
+        """
+        Get list of instance details, optionally applying filters and selective instance ids
+
+        :param instance_ids: List of ids to get instances for
+        :param filters: List of filters to specify instances to get
+        :return: List of instances
+        """
+        description = self.describe_instances(filters=filters, instance_ids=instance_ids)
+
+        return [
+            instance
+            for reservation in description['Reservations']
+            for instance in reservation['Instances']

Review comment:
       ```suggestion
               instance for reservation in description['Reservations'] for instance in reservation['Instances']
   ```

##########
File path: airflow/providers/amazon/aws/hooks/ec2.py
##########
@@ -74,7 +136,11 @@ def wait_for_state(self, instance_id: str, target_state: str, check_interval: fl
         :rtype: None
         """
         instance_state = self.get_instance_state(instance_id=instance_id)
+
+        self.log.info("instance state: %s. Same as target: %s", instance_state, instance_state == target_state)

Review comment:
       ```suggestion
           self.log.info(
               "instance state: %s. Same as target: %s", instance_state, instance_state == target_state
           )
   ```

##########
File path: airflow/providers/amazon/aws/hooks/ec2.py
##########
@@ -74,7 +136,11 @@ def wait_for_state(self, instance_id: str, target_state: str, check_interval: fl
         :rtype: None
         """
         instance_state = self.get_instance_state(instance_id=instance_id)
+
+        self.log.info("instance state: %s. Same as target: %s", instance_state, instance_state == target_state)
+
         while instance_state != target_state:
-            self.log.info("instance state: %s", instance_state)
             time.sleep(check_interval)
             instance_state = self.get_instance_state(instance_id=instance_id)
+
+            self.log.info("instance state: %s. Same as target: %s", instance_state, instance_state == target_state)

Review comment:
       ```suggestion
               self.log.info(
                   "instance state: %s. Same as target: %s", instance_state, instance_state == target_state
               )
   ```




----------------------------------------------------------------
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:
users@infra.apache.org