ashb commented on a change in pull request #9011:
URL: https://github.com/apache/airflow/pull/9011#discussion_r441597364
##########
File path: airflow/providers/amazon/aws/hooks/ec2.py
##########
@@ -33,21 +33,90 @@ class EC2Hook(AwsBaseHook):
:class:`~airflow.providers.amazon.aws.hooks.base_aws.AwsBaseHook`
"""
+ # Describe response
+ RESERVATIONS = 'Reservations'
+ INSTANCES = 'Instances'
+ STATE = 'State'
+ NAME = 'Name'
+ INSTANCE_ID = 'InstanceId'
+
def __init__(self,
*args,
**kwargs):
- super().__init__(resource_type="ec2", *args, **kwargs)
+ super().__init__(client_type="ec2", *args, **kwargs)
- def get_instance(self, instance_id: str):
Review comment:
This is only acceptible since the Ec2Hook hasn't yet been included in a
release -- otherwise this would be a breaking change and you would have have to
include a compatibility layer.
##########
File path: airflow/providers/amazon/aws/hooks/ec2.py
##########
@@ -33,21 +33,90 @@ class EC2Hook(AwsBaseHook):
:class:`~airflow.providers.amazon.aws.hooks.base_aws.AwsBaseHook`
"""
+ # Describe response
+ RESERVATIONS = 'Reservations'
+ INSTANCES = 'Instances'
+ STATE = 'State'
+ NAME = 'Name'
+ INSTANCE_ID = 'InstanceId'
Review comment:
Please don't do this. Use the string directly. Having this extracted out
just makes it harder to understand the code at the call site.
##########
File path: airflow/providers/amazon/aws/hooks/ec2.py
##########
@@ -33,21 +33,90 @@ class EC2Hook(AwsBaseHook):
:class:`~airflow.providers.amazon.aws.hooks.base_aws.AwsBaseHook`
"""
+ # Describe response
+ RESERVATIONS = 'Reservations'
+ INSTANCES = 'Instances'
+ STATE = 'State'
+ NAME = 'Name'
+ INSTANCE_ID = 'InstanceId'
+
def __init__(self,
*args,
**kwargs):
- super().__init__(resource_type="ec2", *args, **kwargs)
+ super().__init__(client_type="ec2", *args, **kwargs)
- def get_instance(self, instance_id: str):
+ def stop_instances(self, instance_ids):
"""
- 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):
+ """
+ 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):
+ """
+ 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)
+
+ def describe_instances(self, filters=None, instance_ids=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=None, instance_ids=None):
"""
- return self.conn.Instance(id=instance_id)
+ 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[self.RESERVATIONS] for instance in
reservation[self.INSTANCES]
Review comment:
What about spot instances? I think this would break if the query
returned any spot instances.
##########
File path: airflow/providers/amazon/aws/hooks/ec2.py
##########
@@ -58,7 +127,7 @@ def get_instance_state(self, instance_id: str) -> str:
:return: current state of the instance
:rtype: str
"""
- return self.get_instance(instance_id=instance_id).state["Name"]
+ return
self.get_instances(instance_ids=[instance_id])[0][self.STATE][self.NAME]
Review comment:
If the instanced ID doesn't exist this probably returns a very unhelpful
error (`list index out of range` for example.)
Please fix this, and and a unit test case covering this.
----------------------------------------------------------------
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:
[email protected]
