From fd57277962efa9effdad7fdf032e1ced6e8ced45 Mon Sep 17 00:00:00 2001 From: Mike Ellis Date: Fri, 31 Jan 2025 21:03:42 +0000 Subject: [PATCH 1/3] Adding extra links for EC2 --- .../airflow/providers/amazon/aws/links/ec2.py | 42 +++++++++++++++ .../providers/amazon/aws/operators/ec2.py | 46 ++++++++++++++++ .../airflow/providers/amazon/provider.yaml | 3 +- providers/tests/amazon/aws/links/test_ec2.py | 53 +++++++++++++++++++ 4 files changed, 143 insertions(+), 1 deletion(-) create mode 100644 providers/src/airflow/providers/amazon/aws/links/ec2.py create mode 100644 providers/tests/amazon/aws/links/test_ec2.py diff --git a/providers/src/airflow/providers/amazon/aws/links/ec2.py b/providers/src/airflow/providers/amazon/aws/links/ec2.py new file mode 100644 index 0000000000000..ca9a770a9f531 --- /dev/null +++ b/providers/src/airflow/providers/amazon/aws/links/ec2.py @@ -0,0 +1,42 @@ +# 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. +from __future__ import annotations + +from airflow.providers.amazon.aws.links.base_aws import BASE_AWS_CONSOLE_LINK, BaseAwsLink + + +class EC2InstanceLink(BaseAwsLink): + """Helper class for constructing Amazon EC2 instance links.""" + + name = "Instance" + key = "_instance_id" + format_str = ( + BASE_AWS_CONSOLE_LINK + "/ec2/home?region={region_name}#InstanceDetails:instanceId={instance_id}" + ) + + +class EC2InstanceDashboardLink(BaseAwsLink): + """ + Helper class for constructing Amazon EC2 console links. + + This is useful for displaying the list of EC2 instances, rather + than a single instance. + """ + + name = "EC2 Instances" + key = "_instance_dashboard" + format_str = BASE_AWS_CONSOLE_LINK + "/ec2/home?region={region_name}#Instances:instanceState={state}" diff --git a/providers/src/airflow/providers/amazon/aws/operators/ec2.py b/providers/src/airflow/providers/amazon/aws/operators/ec2.py index 5b25b27fd0555..45042453c8937 100644 --- a/providers/src/airflow/providers/amazon/aws/operators/ec2.py +++ b/providers/src/airflow/providers/amazon/aws/operators/ec2.py @@ -23,6 +23,7 @@ from airflow.exceptions import AirflowException from airflow.models import BaseOperator from airflow.providers.amazon.aws.hooks.ec2 import EC2Hook +from airflow.providers.amazon.aws.links.ec2 import EC2InstanceDashboardLink, EC2InstanceLink if TYPE_CHECKING: from airflow.utils.context import Context @@ -47,6 +48,7 @@ class EC2StartInstanceOperator(BaseOperator): between each instance state checks until operation is completed """ + operator_extra_links = (EC2InstanceLink(),) template_fields: Sequence[str] = ("instance_id", "region_name") ui_color = "#eeaa11" ui_fgcolor = "#ffffff" @@ -71,6 +73,13 @@ def execute(self, context: Context): self.log.info("Starting EC2 instance %s", self.instance_id) instance = ec2_hook.get_instance(instance_id=self.instance_id) instance.start() + EC2InstanceLink.persist( + context=context, + operator=self, + aws_partition=ec2_hook.conn_partition, + instance_id=self.instance_id, + region_name=ec2_hook.conn_region_name, + ) ec2_hook.wait_for_state( instance_id=self.instance_id, target_state="running", @@ -97,6 +106,7 @@ class EC2StopInstanceOperator(BaseOperator): between each instance state checks until operation is completed """ + operator_extra_links = (EC2InstanceLink(),) template_fields: Sequence[str] = ("instance_id", "region_name") ui_color = "#eeaa11" ui_fgcolor = "#ffffff" @@ -120,7 +130,15 @@ def execute(self, context: Context): ec2_hook = EC2Hook(aws_conn_id=self.aws_conn_id, region_name=self.region_name) self.log.info("Stopping EC2 instance %s", self.instance_id) instance = ec2_hook.get_instance(instance_id=self.instance_id) + EC2InstanceLink.persist( + context=context, + operator=self, + aws_partition=ec2_hook.conn_partition, + instance_id=self.instance_id, + region_name=ec2_hook.conn_region_name, + ) instance.stop() + ec2_hook.wait_for_state( instance_id=self.instance_id, target_state="stopped", @@ -154,6 +172,7 @@ class EC2CreateInstanceOperator(BaseOperator): in the `running` state before returning. """ + operator_extra_links = (EC2InstanceDashboardLink(),) template_fields: Sequence[str] = ( "image_id", "max_count", @@ -197,6 +216,14 @@ def execute(self, context: Context): **self.config, )["Instances"] + # Console link is for EC2 dashboard list, not individual instances + EC2InstanceDashboardLink.persist( + context=context, + operator=self, + region_name=ec2_hook.conn_region_name, + aws_partition=ec2_hook.conn_partition, + state="running", + ) instance_ids = self._on_kill_instance_ids = [instance["InstanceId"] for instance in instances] for instance_id in instance_ids: self.log.info("Created EC2 instance %s", instance_id) @@ -311,6 +338,7 @@ class EC2RebootInstanceOperator(BaseOperator): in the `running` state before returning. """ + operator_extra_links = (EC2InstanceDashboardLink(),) template_fields: Sequence[str] = ("instance_ids", "region_name") ui_color = "#eeaa11" ui_fgcolor = "#ffffff" @@ -341,6 +369,14 @@ def execute(self, context: Context): self.log.info("Rebooting EC2 instances %s", ", ".join(self.instance_ids)) ec2_hook.conn.reboot_instances(InstanceIds=self.instance_ids) + # Console link is for EC2 dashboard list, not individual instances + EC2InstanceDashboardLink.persist( + context=context, + operator=self, + region_name=ec2_hook.conn_region_name, + aws_partition=ec2_hook.conn_partition, + state="running", + ) if self.wait_for_completion: ec2_hook.get_waiter("instance_running").wait( InstanceIds=self.instance_ids, @@ -374,6 +410,7 @@ class EC2HibernateInstanceOperator(BaseOperator): in the `stopped` state before returning. """ + operator_extra_links = (EC2InstanceDashboardLink(),) template_fields: Sequence[str] = ("instance_ids", "region_name") ui_color = "#eeaa11" ui_fgcolor = "#ffffff" @@ -404,6 +441,15 @@ def execute(self, context: Context): self.log.info("Hibernating EC2 instances %s", ", ".join(self.instance_ids)) instances = ec2_hook.get_instances(instance_ids=self.instance_ids) + # Console link is for EC2 dashboard list, not individual instances + EC2InstanceDashboardLink.persist( + context=context, + operator=self, + region_name=ec2_hook.conn_region_name, + aws_partition=ec2_hook.conn_partition, + state="stopping", + ) + for instance in instances: hibernation_options = instance.get("HibernationOptions") if not hibernation_options or not hibernation_options["Configured"]: diff --git a/providers/src/airflow/providers/amazon/provider.yaml b/providers/src/airflow/providers/amazon/provider.yaml index 43569a28827ab..5192052800079 100644 --- a/providers/src/airflow/providers/amazon/provider.yaml +++ b/providers/src/airflow/providers/amazon/provider.yaml @@ -891,7 +891,8 @@ extra-links: - airflow.providers.amazon.aws.links.comprehend.ComprehendDocumentClassifierLink - airflow.providers.amazon.aws.links.datasync.DataSyncTaskLink - airflow.providers.amazon.aws.links.datasync.DataSyncTaskExecutionLink - + - airflow.providers.amazon.aws.links.ec2.EC2InstanceLink + - airflow.providers.amazon.aws.links.ec2.EC2InstanceDashboardLink connection-types: - hook-class-name: airflow.providers.amazon.aws.hooks.base_aws.AwsGenericHook diff --git a/providers/tests/amazon/aws/links/test_ec2.py b/providers/tests/amazon/aws/links/test_ec2.py new file mode 100644 index 0000000000000..534a9670338b7 --- /dev/null +++ b/providers/tests/amazon/aws/links/test_ec2.py @@ -0,0 +1,53 @@ +# 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. +from __future__ import annotations + +from airflow.providers.amazon.aws.links.ec2 import EC2InstanceDashboardLink, EC2InstanceLink + +from providers.tests.amazon.aws.links.test_base_aws import BaseAwsLinksTestCase + + +class TestEC2InstanceLink(BaseAwsLinksTestCase): + link_class = EC2InstanceLink + + INSTANCE_ID = "i-xxxxxxxxxxxx" + + def test_extra_link(self): + self.assert_extra_link_url( + expected_url=( + "https://console.aws.amazon.com/ec2/home" + f"?region=eu-west-1#InstanceDetails:instanceId={self.INSTANCE_ID}" + ), + region_name="eu-west-1", + aws_partition="aws", + instance_id=self.INSTANCE_ID, + ) + + +class TestEC2InstanceDashboardLink(BaseAwsLinksTestCase): + link_class = EC2InstanceDashboardLink + + STATE = "running" + BASE_URL = "https://console.aws.amazon.com/ec2/home" + + def test_extra_link(self): + self.assert_extra_link_url( + expected_url=(f"{self.BASE_URL}?region=eu-west-1#Instances:instanceState={self.STATE}"), + region_name="eu-west-1", + aws_partition="aws", + state=self.STATE, + ) From d0058241d3109da19f2547214368100e28a49790 Mon Sep 17 00:00:00 2001 From: Mike Ellis Date: Mon, 3 Feb 2025 19:17:31 +0000 Subject: [PATCH 2/3] Update console link to filter on multiple instance ids --- .../airflow/providers/amazon/aws/links/ec2.py | 7 ++++++- .../providers/amazon/aws/operators/ec2.py | 17 +++++++++++------ providers/tests/amazon/aws/links/test_ec2.py | 13 ++++++++++--- 3 files changed, 27 insertions(+), 10 deletions(-) diff --git a/providers/src/airflow/providers/amazon/aws/links/ec2.py b/providers/src/airflow/providers/amazon/aws/links/ec2.py index ca9a770a9f531..143e873911503 100644 --- a/providers/src/airflow/providers/amazon/aws/links/ec2.py +++ b/providers/src/airflow/providers/amazon/aws/links/ec2.py @@ -39,4 +39,9 @@ class EC2InstanceDashboardLink(BaseAwsLink): name = "EC2 Instances" key = "_instance_dashboard" - format_str = BASE_AWS_CONSOLE_LINK + "/ec2/home?region={region_name}#Instances:instanceState={state}" + format_str = BASE_AWS_CONSOLE_LINK + "/ec2/home?region={region_name}#Instances:instanceId=:{instance_ids}" + # Instances:instanceId=: + + @staticmethod + def format_instance_id_filter(instance_ids: list[str]) -> str: + return ",:".join(instance_ids) diff --git a/providers/src/airflow/providers/amazon/aws/operators/ec2.py b/providers/src/airflow/providers/amazon/aws/operators/ec2.py index 45042453c8937..2851810170b7d 100644 --- a/providers/src/airflow/providers/amazon/aws/operators/ec2.py +++ b/providers/src/airflow/providers/amazon/aws/operators/ec2.py @@ -23,7 +23,11 @@ from airflow.exceptions import AirflowException from airflow.models import BaseOperator from airflow.providers.amazon.aws.hooks.ec2 import EC2Hook -from airflow.providers.amazon.aws.links.ec2 import EC2InstanceDashboardLink, EC2InstanceLink +from airflow.providers.amazon.aws.links.ec2 import ( + EC2InstanceDashboardLink, + EC2InstanceLink, + # format_instance_id_filter, +) if TYPE_CHECKING: from airflow.utils.context import Context @@ -216,15 +220,16 @@ def execute(self, context: Context): **self.config, )["Instances"] - # Console link is for EC2 dashboard list, not individual instances + instance_ids = self._on_kill_instance_ids = [instance["InstanceId"] for instance in instances] + # Console link is for EC2 dashboard list, not individual instances when more than 1 instance + EC2InstanceDashboardLink.persist( context=context, operator=self, region_name=ec2_hook.conn_region_name, aws_partition=ec2_hook.conn_partition, - state="running", + instance_ids=EC2InstanceDashboardLink.format_instance_id_filter(instance_ids), ) - instance_ids = self._on_kill_instance_ids = [instance["InstanceId"] for instance in instances] for instance_id in instance_ids: self.log.info("Created EC2 instance %s", instance_id) @@ -375,7 +380,7 @@ def execute(self, context: Context): operator=self, region_name=ec2_hook.conn_region_name, aws_partition=ec2_hook.conn_partition, - state="running", + instance_ids=EC2InstanceDashboardLink.format_instance_id_filter(self.instance_ids), ) if self.wait_for_completion: ec2_hook.get_waiter("instance_running").wait( @@ -447,7 +452,7 @@ def execute(self, context: Context): operator=self, region_name=ec2_hook.conn_region_name, aws_partition=ec2_hook.conn_partition, - state="stopping", + instance_ids=EC2InstanceDashboardLink.format_instance_id_filter(self.instance_ids), ) for instance in instances: diff --git a/providers/tests/amazon/aws/links/test_ec2.py b/providers/tests/amazon/aws/links/test_ec2.py index 534a9670338b7..98f8e68476e30 100644 --- a/providers/tests/amazon/aws/links/test_ec2.py +++ b/providers/tests/amazon/aws/links/test_ec2.py @@ -41,13 +41,20 @@ def test_extra_link(self): class TestEC2InstanceDashboardLink(BaseAwsLinksTestCase): link_class = EC2InstanceDashboardLink - STATE = "running" + # STATE = "running" BASE_URL = "https://console.aws.amazon.com/ec2/home" + INSTANCE_IDS = ["i-xxxxxxxxxxxx", "i-yyyyyyyyyyyy"] + + def test_instance_id_filter(self): + instance_list = ",:".join(self.INSTANCE_IDS) + result = EC2InstanceDashboardLink.format_instance_id_filter(self.INSTANCE_IDS) + assert result == instance_list def test_extra_link(self): + instance_list = ",:".join(self.INSTANCE_IDS) self.assert_extra_link_url( - expected_url=(f"{self.BASE_URL}?region=eu-west-1#Instances:instanceState={self.STATE}"), + expected_url=(f"{self.BASE_URL}?region=eu-west-1#Instances:instanceId=:{instance_list}"), region_name="eu-west-1", aws_partition="aws", - state=self.STATE, + instance_ids=EC2InstanceDashboardLink.format_instance_id_filter(self.INSTANCE_IDS), ) From e5f9a024afcfc820cb94dfc30027eaf6191fb549 Mon Sep 17 00:00:00 2001 From: Mike Ellis Date: Mon, 3 Feb 2025 20:22:14 +0000 Subject: [PATCH 3/3] Removed test notes --- providers/src/airflow/providers/amazon/aws/links/ec2.py | 1 - providers/src/airflow/providers/amazon/aws/operators/ec2.py | 1 - providers/tests/amazon/aws/links/test_ec2.py | 1 - 3 files changed, 3 deletions(-) diff --git a/providers/src/airflow/providers/amazon/aws/links/ec2.py b/providers/src/airflow/providers/amazon/aws/links/ec2.py index 143e873911503..38a23956cddbb 100644 --- a/providers/src/airflow/providers/amazon/aws/links/ec2.py +++ b/providers/src/airflow/providers/amazon/aws/links/ec2.py @@ -40,7 +40,6 @@ class EC2InstanceDashboardLink(BaseAwsLink): name = "EC2 Instances" key = "_instance_dashboard" format_str = BASE_AWS_CONSOLE_LINK + "/ec2/home?region={region_name}#Instances:instanceId=:{instance_ids}" - # Instances:instanceId=: @staticmethod def format_instance_id_filter(instance_ids: list[str]) -> str: diff --git a/providers/src/airflow/providers/amazon/aws/operators/ec2.py b/providers/src/airflow/providers/amazon/aws/operators/ec2.py index 2851810170b7d..f3d0e9fc2af25 100644 --- a/providers/src/airflow/providers/amazon/aws/operators/ec2.py +++ b/providers/src/airflow/providers/amazon/aws/operators/ec2.py @@ -26,7 +26,6 @@ from airflow.providers.amazon.aws.links.ec2 import ( EC2InstanceDashboardLink, EC2InstanceLink, - # format_instance_id_filter, ) if TYPE_CHECKING: diff --git a/providers/tests/amazon/aws/links/test_ec2.py b/providers/tests/amazon/aws/links/test_ec2.py index 98f8e68476e30..922b12275e5aa 100644 --- a/providers/tests/amazon/aws/links/test_ec2.py +++ b/providers/tests/amazon/aws/links/test_ec2.py @@ -41,7 +41,6 @@ def test_extra_link(self): class TestEC2InstanceDashboardLink(BaseAwsLinksTestCase): link_class = EC2InstanceDashboardLink - # STATE = "running" BASE_URL = "https://console.aws.amazon.com/ec2/home" INSTANCE_IDS = ["i-xxxxxxxxxxxx", "i-yyyyyyyyyyyy"]