Adding support for using client_type API for interacting with EC2 and support for using filters

airflow/providers/amazon/aws/hooks/ec2.py

@@ -17,11 +17,35 @@
 # under the License.
 #
 
+import functools
 import time
 
+from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
 
 
+def only_client_type(func):
+    @functools.wraps(func)
+    def checker(self, *args, **kwargs):
+        if self._api_type == "client_type":
+            return func(self, *args, **kwargs)
+
+        raise AirflowException(
+            """
+            This method is only callable when using client_type API for interacting with EC2.
+            Create the EC2Hook object as follows to use this method
+
+            ec2 = EC2Hook(api_type="client_type")
+
+            Read following for details on client_type and resource_type APIs:
+            1. https://boto3.amazonaws.com/v1/documentation/api/1.9.42/reference/services/ec2.html#client
+            2. https://boto3.amazonaws.com/v1/documentation/api/1.9.42/reference/services/ec2.html#service-resource # noqa
+            """
+        )
+
+    return checker
+
+
 class EC2Hook(AwsBaseHook):
     """
     Interact with AWS EC2 Service.
@@ -33,21 +57,112 @@ class EC2Hook(AwsBaseHook):
         :class:`~airflow.providers.amazon.aws.hooks.base_aws.AwsBaseHook`
     """
 
-    def __init__(self, *args, **kwargs) -> None:
-        kwargs["resource_type"] = "ec2"
+    API_TYPES = frozenset({"resource_type", "client_type"})
+
+    def __init__(self, api_type="resource_type", *args, **kwargs) -> None:
+        if api_type not in self.API_TYPES:
+            raise AirflowException("api_type can only be one of %s", self.API_TYPES)
+
+        kwargs[api_type] = "ec2"
+
+        self._api_type = api_type
+
         super().__init__(*args, **kwargs)
 
-    def get_instance(self, instance_id: str):
+    def get_instance(self, instance_id: str, filters: list = None):
         """
         Get EC2 instance by id and return it.
 
         :param instance_id: id of the AWS EC2 instance
         :type instance_id: str
+        :param filters: List of filters to specify instances to get
+        :type filters: list
         :return: Instance object
         :rtype: ec2.Instance
         """
+        if self._api_type == "client_type":
+            return self.get_instances(filters=filters, instance_ids=[instance_id])
+
         return self.conn.Instance(id=instance_id)
 
+    @only_client_type
+    def stop_instances(self, instance_ids: list) -> dict:
+        """
+        Stop instances with given ids
+
+        :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)
+
+    @only_client_type
+    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)
+
+    @only_client_type
+    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
+        """
+        self.log.info("Terminating instances: %s", instance_ids)
+
+        return self.conn.terminate_instances(InstanceIds=instance_ids)
+
+    @only_client_type
+    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 = filters or []
+        instance_ids = instance_ids or []
+
+        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)
+
+    @only_client_type
+    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"]
+        ]
+
+    @only_client_type
+    def get_instance_ids(self, filters: list = None) -> list:
+        """
+        Get list of instance ids, optionally applying filters to fetch selective instances
+
+        :param filters: List of filters to specify instances to get
+        :return: List of instance ids
+        """
+        return [instance["InstanceId"] for instance in self.get_instances(filters=filters)]
+
     def get_instance_state(self, instance_id: str) -> str:
         """
         Get EC2 instance state by id and return it.
@@ -57,6 +172,9 @@ def get_instance_state(self, instance_id: str) -> str:
         :return: current state of the instance
         :rtype: str
         """
+        if self._api_type == "client_type":
+            return self.get_instances(instance_ids=[instance_id])[0]["State"]["Name"]
+
         return self.get_instance(instance_id=instance_id).state["Name"]
 
     def wait_for_state(self, instance_id: str, target_state: str, check_interval: float) -> None:
@@ -74,7 +192,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)
+
         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
+            )

tests/providers/amazon/aws/hooks/test_ec2.py

@@ -21,6 +21,7 @@
 
 from moto import mock_ec2
 
+from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.ec2 import EC2Hook
 
 
@@ -39,6 +40,12 @@ def test_get_conn_returns_boto3_resource(self):
         instances = list(ec2_hook.conn.instances.all())
         assert instances is not None
 
+    @mock_ec2
+    def test_client_type_get_conn_returns_boto3_resource(self):
+        ec2_hook = EC2Hook(api_type="client_type")
+        instances = list(ec2_hook.get_instances())
+        assert instances is not None
+
     @mock_ec2
     def test_get_instance(self):
         ec2_hook = EC2Hook()
@@ -58,9 +65,188 @@ def test_get_instance_state(self):
             MaxCount=1,
             MinCount=1,
         )
+
         created_instance_id = created_instances[0].instance_id
         all_instances = list(ec2_hook.conn.instances.all())
         created_instance_state = all_instances[0].state["Name"]
         # test get_instance_state method
         existing_instance_state = ec2_hook.get_instance_state(instance_id=created_instance_id)
         assert created_instance_state == existing_instance_state
+
+    @mock_ec2
+    def test_client_type_get_instance_state(self):
+        ec2_hook = EC2Hook(api_type="client_type")
+        created_instances = ec2_hook.conn.run_instances(
+            MaxCount=1,
+            MinCount=1,
+        )
+
+        created_instance_id = created_instances['Instances'][0]['InstanceId']
+        all_instances = ec2_hook.get_instances()
+        created_instance_state = all_instances[0]['State']['Name']
+
+        existing_instance_state = ec2_hook.get_instance_state(instance_id=created_instance_id)
+        assert created_instance_state == existing_instance_state
+
+    @mock_ec2
+    def test_client_type_start_instances(self):
+        ec2_hook = EC2Hook(api_type="client_type")
+        created_instances = ec2_hook.conn.run_instances(
+            MaxCount=1,
+            MinCount=1,
+        )
+
+        created_instance_id = created_instances['Instances'][0]['InstanceId']
+        response = ec2_hook.start_instances(instance_ids=[created_instance_id])
+
+        assert response["StartingInstances"][0]["InstanceId"] == created_instance_id
+        assert ec2_hook.get_instance_state(created_instance_id) == "running"
+
+    @mock_ec2
+    def test_client_type_stop_instances(self):
+        ec2_hook = EC2Hook(api_type="client_type")
+        created_instances = ec2_hook.conn.run_instances(
+            MaxCount=1,
+            MinCount=1,
+        )
+
+        created_instance_id = created_instances['Instances'][0]['InstanceId']
+        response = ec2_hook.stop_instances(instance_ids=[created_instance_id])
+
+        assert response["StoppingInstances"][0]["InstanceId"] == created_instance_id
+        assert ec2_hook.get_instance_state(created_instance_id) == "stopped"
+
+    @mock_ec2
+    def test_client_type_terminate_instances(self):
+        ec2_hook = EC2Hook(api_type="client_type")
+        created_instances = ec2_hook.conn.run_instances(
+            MaxCount=1,
+            MinCount=1,
+        )
+
+        created_instance_id = created_instances['Instances'][0]['InstanceId']
+        response = ec2_hook.terminate_instances(instance_ids=[created_instance_id])
+
+        assert response["TerminatingInstances"][0]["InstanceId"] == created_instance_id
+        assert ec2_hook.get_instance_state(created_instance_id) == "terminated"
+
+    @mock_ec2
+    def test_client_type_describe_instances(self):
+        ec2_hook = EC2Hook(api_type="client_type")
+        created_instances = ec2_hook.conn.run_instances(
+            MaxCount=1,
+            MinCount=1,
+        )
+
+        created_instance_id = created_instances['Instances'][0]['InstanceId']
+
+        # Without filter
+        response = ec2_hook.describe_instances(instance_ids=[created_instance_id])
+
+        assert response["Reservations"][0]["Instances"][0]["InstanceId"] == created_instance_id
+        assert response["Reservations"][0]["Instances"][0]["State"]["Name"] == "running"
+
+        # With valid filter
+        response = ec2_hook.describe_instances(
+            filters=[{"Name": "instance-id", "Values": [created_instance_id]}]
+        )
+
+        assert len(response["Reservations"]) == 1
+        assert response["Reservations"][0]["Instances"][0]["InstanceId"] == created_instance_id
+        assert response["Reservations"][0]["Instances"][0]["State"]["Name"] == "running"
+
+        # With invalid filter
+        response = ec2_hook.describe_instances(
+            filters=[{"Name": "instance-id", "Values": ["invalid_instance_id"]}]
+        )
+
+        assert len(response["Reservations"]) == 0
+
+    @mock_ec2
+    def test_client_type_get_instances(self):
+        ec2_hook = EC2Hook(api_type="client_type")
+        created_instances = ec2_hook.conn.run_instances(
+            MaxCount=2,
+            MinCount=2,
+        )
+
+        created_instance_id_1 = created_instances['Instances'][0]['InstanceId']
+        created_instance_id_2 = created_instances['Instances'][1]['InstanceId']
+
+        # Without filter
+        response = ec2_hook.get_instances(instance_ids=[created_instance_id_1, created_instance_id_2])
+
+        assert response[0]["InstanceId"] == created_instance_id_1
+        assert response[1]["InstanceId"] == created_instance_id_2
+
+        # With valid filter
+        response = ec2_hook.get_instances(
+            filters=[{"Name": "instance-id", "Values": [created_instance_id_1, created_instance_id_2]}]
+        )
+
+        assert len(response) == 2
+        assert response[0]["InstanceId"] == created_instance_id_1
+        assert response[1]["InstanceId"] == created_instance_id_2
+
+        # With filter and instance ids
+        response = ec2_hook.get_instances(
+            filters=[{"Name": "instance-id", "Values": [created_instance_id_1]}],
+            instance_ids=[created_instance_id_1, created_instance_id_2],
+        )
+
+        assert len(response) == 1
+        assert response[0]["InstanceId"] == created_instance_id_1
+
+        # With invalid filter
+        response = ec2_hook.get_instances(
+            filters=[{"Name": "instance-id", "Values": ["invalid_instance_id"]}]
+        )
+
+        assert len(response) == 0
+
+    @mock_ec2
+    def test_client_type_get_instance_ids(self):
+        ec2_hook = EC2Hook(api_type="client_type")
+        created_instances = ec2_hook.conn.run_instances(
+            MaxCount=2,
+            MinCount=2,
+        )
+
+        created_instance_id_1 = created_instances['Instances'][0]['InstanceId']
+        created_instance_id_2 = created_instances['Instances'][1]['InstanceId']
+
+        # Without filter
+        response = ec2_hook.get_instance_ids()
+
+        assert len(response) == 2
+        assert response[0] == created_instance_id_1
+        assert response[1] == created_instance_id_2
+
+        # With valid filter
+        response = ec2_hook.get_instance_ids(filters=[{"Name": "instance-type", "Values": ["m1.small"]}])
+
+        assert len(response) == 2
+        assert response[0] == created_instance_id_1
+        assert response[1] == created_instance_id_2
+
+        # With invalid filter
+        response = ec2_hook.get_instance_ids(
+            filters=[{"Name": "instance-type", "Values": ["invalid_instance_type"]}]
+        )
+
+        assert len(response) == 0
+
+    @mock_ec2
+    def test_decorator_only_client_type(self):
+        ec2_hook = EC2Hook()
+
+        # Try calling a method which is only supported by client_type API
+        with self.assertRaises(AirflowException):
+            ec2_hook.get_instances()
+
+        # Explicitly provide resource_type as api_type
+        ec2_hook = EC2Hook(api_type="resource_type")
+
+        # Try calling a method which is only supported by client_type API
+        with self.assertRaises(AirflowException):
+            ec2_hook.describe_instances()