Add AWS Batch & AWS CloudWatch Extra Links

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

@@ -410,6 +410,45 @@ def parse_job_description(job_id: str, response: Dict) -> Dict:
 
         return matching_jobs[0]
 
+    def get_job_awslogs_info(self, job_id: str) -> Optional[Dict[str, str]]:
+        """
+        Parse job description to extract AWS CloudWatch information.
+
+        :param job_id: AWS Batch Job ID
+        """
+        job_container_desc = self.get_job_description(job_id=job_id).get("container", {})
+        log_configuration = job_container_desc.get("logConfiguration", {})
+
+        # In case if user select other "logDriver" rather than "awslogs"
+        # than CloudWatch logging should be disabled.
+        # If user not specify anything than expected that "awslogs" will use
+        # with default settings:
+        #   awslogs-group = /aws/batch/job
+        #   awslogs-region = `same as AWS Batch Job region`
+        log_driver = log_configuration.get("logDriver", "awslogs")
+        if log_driver != "awslogs":
+            self.log.warning(
+                "AWS Batch job (%s) uses logDriver (%s). AWS CloudWatch logging disabled.", job_id, log_driver
+            )
+            return None
+
+        awslogs_stream_name = job_container_desc.get("logStreamName")
+        if not awslogs_stream_name:
+            # In case of call this method on very early stage of running AWS Batch
+            # there is possibility than AWS CloudWatch Stream Name not exists yet.
+            # AWS CloudWatch Stream Name also not created in case of misconfiguration.
+            self.log.warning("AWS Batch job (%s) doesn't create AWS CloudWatch Stream.", job_id)
+            return None
+
+        # Try to get user-defined log configuration options
+        log_options = log_configuration.get("options", {})
+
+        return {
+            "awslogs_stream_name": awslogs_stream_name,
+            "awslogs_group": log_options.get("awslogs-group", "/aws/batch/job"),
+            "awslogs_region": log_options.get("awslogs-region", self.conn_region_name),
+        }
+
     @staticmethod
     def add_jitter(
         delay: Union[int, float], width: Union[int, float] = 1, minima: Union[int, float] = 0

airflow/providers/amazon/aws/links/base_aws.py

@@ -18,7 +18,6 @@
 
 from datetime import datetime
 from typing import TYPE_CHECKING, ClassVar, Optional
-from urllib.parse import quote_plus
 
 from airflow.models import BaseOperatorLink, XCom
 
@@ -49,6 +48,18 @@ def get_aws_domain(aws_partition) -> Optional[str]:
 
         return None
 
+    def format_link(self, **kwargs) -> str:
+        """
+        Format AWS Service Link
+
+        Some AWS Service Link should require additional escaping
+        in this case this method should be overridden.
+        """
+        try:
+            return self.format_str.format(**kwargs)
+        except KeyError:
+            return ""
+
     def get_link(
         self,
         operator,
@@ -74,12 +85,8 @@ def get_link(
                 task_id=operator.task_id,
                 execution_date=dttm,
             )
-        if not conf:
-            return ""
 
-        # urlencode special characters, e.g.: CloudWatch links contains `/` character.
-        quoted_conf = {k: quote_plus(v) if isinstance(v, str) else v for k, v in conf.items()}
-        return self.format_str.format(**quoted_conf)
+        return self.format_link(**conf) if conf else ""
 
     @classmethod
     def persist(

airflow/providers/amazon/aws/links/batch.py

@@ -0,0 +1,44 @@
+# 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 airflow.providers.amazon.aws.links.base_aws import BASE_AWS_CONSOLE_LINK, BaseAwsLink
+
+
+class BatchJobDefinitionLink(BaseAwsLink):
+    """Helper class for constructing AWS Batch Job Definition Link"""
+
+    name = "Batch Job Definition"
+    key = "batch_job_definition"
+    format_str = (
+        BASE_AWS_CONSOLE_LINK + "/batch/home?region={region_name}#job-definition/detail/{job_definition_arn}"
+    )
+
+
+class BatchJobDetailsLink(BaseAwsLink):
+    """Helper class for constructing AWS Batch Job Details Link"""
+
+    name = "Batch Job Details"
+    key = "batch_job_details"
+    format_str = BASE_AWS_CONSOLE_LINK + "/batch/home?region={region_name}#jobs/detail/{job_id}"
+
+
+class BatchJobQueueLink(BaseAwsLink):
+    """Helper class for constructing AWS Batch Job Queue Link"""
+
+    name = "Batch Job Queue"
+    key = "batch_job_queue"
+    format_str = BASE_AWS_CONSOLE_LINK + "/batch/home?region={region_name}#queues/detail/{job_queue_arn}"

airflow/providers/amazon/aws/links/logs.py

@@ -0,0 +1,41 @@
+# 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 urllib.parse import quote_plus
+
+from airflow.providers.amazon.aws.links.base_aws import BASE_AWS_CONSOLE_LINK, BaseAwsLink
+
+
+class CloudWatchEventsLink(BaseAwsLink):
+    """Helper class for constructing AWS CloudWatch Events Link"""
+
+    name = "CloudWatch Events"
+    key = "cloudwatch_events"
+    format_str = (
+        BASE_AWS_CONSOLE_LINK
+        + "/cloudwatch/home?region={awslogs_region}#logsV2:log-groups/log-group/{awslogs_group}"
+        + "/log-events/{awslogs_stream_name}"
+    )
+
+    def format_link(self, **kwargs) -> str:
+        for field in ("awslogs_stream_name", "awslogs_group"):
+            if field in kwargs:
+                kwargs[field] = quote_plus(kwargs[field])
+            else:
+                return ""
+
+        return super().format_link(**kwargs)

airflow/providers/amazon/aws/operators/batch.py

@@ -31,6 +31,12 @@
 from airflow.exceptions import AirflowException
 from airflow.models import BaseOperator
 from airflow.providers.amazon.aws.hooks.batch_client import BatchClientHook
+from airflow.providers.amazon.aws.links.batch import (
+    BatchJobDefinitionLink,
+    BatchJobDetailsLink,
+    BatchJobQueueLink,
+)
+from airflow.providers.amazon.aws.links.logs import CloudWatchEventsLink
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context
@@ -95,6 +101,15 @@ class BatchOperator(BaseOperator):
     )
     template_fields_renderers = {"overrides": "json", "parameters": "json"}
 
+    @property
+    def operator_extra_links(self):
+        op_extra_links = [BatchJobDetailsLink(), BatchJobDefinitionLink(), BatchJobQueueLink()]
+        if not self.array_properties:
+            # There is no CloudWatch Link to the parent Batch Job available.
+            op_extra_links.append(CloudWatchEventsLink())
+
+        return tuple(op_extra_links)
+
     def __init__(
         self,
         *,
@@ -167,13 +182,24 @@ def submit_job(self, context: 'Context'):
                 containerOverrides=self.overrides,
                 tags=self.tags,
             )
-            self.job_id = response["jobId"]
-
-            self.log.info("AWS Batch job (%s) started: %s", self.job_id, response)
         except Exception as e:
-            self.log.error("AWS Batch job (%s) failed submission", self.job_id)
+            self.log.error(
+                "AWS Batch job failed submission - job definition: %s - on queue %s",
+                self.job_definition,
+                self.job_queue,
+            )
             raise AirflowException(e)
 
+        self.job_id = response["jobId"]
+        self.log.info("AWS Batch job (%s) started: %s", self.job_id, response)
+        BatchJobDetailsLink.persist(
+            context=context,
+            operator=self,
+            region_name=self.hook.conn_region_name,
+            aws_partition=self.hook.conn_partition,
+            job_id=self.job_id,
+        )
+
     def monitor_job(self, context: 'Context'):
         """
         Monitor an AWS Batch job
@@ -186,11 +212,50 @@ def monitor_job(self, context: 'Context'):
         if not self.job_id:
             raise AirflowException('AWS Batch job - job_id was not found')
 
+        try:
+            job_desc = self.hook.get_job_description(self.job_id)
+            job_definition_arn = job_desc["jobDefinition"]
+            job_queue_arn = job_desc["jobQueue"]
+            self.log.info(
+                "AWS Batch job (%s) Job Definition ARN: %r, Job Queue ARN: %r",
+                self.job_id,
+                job_definition_arn,
+                job_queue_arn,
+            )
+        except KeyError:
+            self.log.warning("AWS Batch job (%s) can't get Job Definition ARN and Job Queue ARN", self.job_id)
+        else:
+            BatchJobDefinitionLink.persist(
+                context=context,
+                operator=self,
+                region_name=self.hook.conn_region_name,
+                aws_partition=self.hook.conn_partition,
+                job_definition_arn=job_definition_arn,
+            )
+            BatchJobQueueLink.persist(
+                context=context,
+                operator=self,
+                region_name=self.hook.conn_region_name,
+                aws_partition=self.hook.conn_partition,
+                job_queue_arn=job_queue_arn,
+            )
+
         if self.waiters:
             self.waiters.wait_for_job(self.job_id)
         else:
             self.hook.wait_for_job(self.job_id)
 
+        awslogs = self.hook.get_job_awslogs_info(self.job_id)
+        if awslogs:
+            self.log.info("AWS Batch job (%s) CloudWatch Events details found: %s", self.job_id, awslogs)
+            CloudWatchEventsLink.persist(
+                context=context,
+                operator=self,
+                region_name=self.hook.conn_region_name,
+                aws_partition=self.hook.conn_partition,
+                **awslogs,
+            )
+
         self.hook.check_job_success(self.job_id)
         self.log.info("AWS Batch job (%s) succeeded", self.job_id)
 

airflow/providers/amazon/provider.yaml

@@ -556,7 +556,11 @@ hook-class-names:  # deprecated - to be removed after providers add dependency o
   - airflow.providers.amazon.aws.hooks.redshift_sql.RedshiftSQLHook
 
 extra-links:
+  - airflow.providers.amazon.aws.links.batch.BatchJobDefinitionLink
+  - airflow.providers.amazon.aws.links.batch.BatchJobDetailsLink
+  - airflow.providers.amazon.aws.links.batch.BatchJobQueueLink
   - airflow.providers.amazon.aws.links.emr.EmrClusterLink
+  - airflow.providers.amazon.aws.links.logs.CloudWatchEventsLink
   - airflow.providers.amazon.aws.operators.emr_create_job_flow.EmrClusterLink
 
 connection-types:

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

@@ -33,6 +33,7 @@
 AWS_SECRET_ACCESS_KEY = "airflow_dummy_secret"
 
 JOB_ID = "8ba9d676-4108-4474-9dca-8bbac1da9b19"
+LOG_STREAM_NAME = "test/stream/d56a66bb98a14c4593defa1548686edf"
 
 
 class TestBatchClient(unittest.TestCase):
@@ -229,6 +230,76 @@ def test_terminate_job(self):
         self.client_mock.terminate_job.assert_called_once_with(jobId=JOB_ID, reason=reason)
         assert response == {}
 
+    def test_job_awslogs_default(self):
+        self.client_mock.describe_jobs.return_value = {
+            "jobs": [
+                {
+                    "jobId": JOB_ID,
+                    "container": {"logStreamName": LOG_STREAM_NAME},
+                }
+            ]
+        }
+        self.client_mock.meta.client.meta.region_name = AWS_REGION
+
+        awslogs = self.batch_client.get_job_awslogs_info(JOB_ID)
+        assert awslogs["awslogs_stream_name"] == LOG_STREAM_NAME
+        assert awslogs["awslogs_group"] == "/aws/batch/job"
+        assert awslogs["awslogs_region"] == AWS_REGION
+
+    def test_job_awslogs_user_defined(self):
+        self.client_mock.describe_jobs.return_value = {
+            "jobs": [
+                {
+                    "jobId": JOB_ID,
+                    "container": {
+                        "logStreamName": LOG_STREAM_NAME,
+                        "logConfiguration": {
+                            "logDriver": "awslogs",
+                            "options": {
+                                "awslogs-group": "/test/batch/job",
+                                "awslogs-region": "ap-southeast-2",
+                            },
+                        },
+                    },
+                }
+            ]
+        }
+        awslogs = self.batch_client.get_job_awslogs_info(JOB_ID)
+        assert awslogs["awslogs_stream_name"] == LOG_STREAM_NAME
+        assert awslogs["awslogs_group"] == "/test/batch/job"
+        assert awslogs["awslogs_region"] == "ap-southeast-2"
+
+    def test_job_no_awslogs_stream(self):
+        self.client_mock.describe_jobs.return_value = {
+            "jobs": [
+                {
+                    "jobId": JOB_ID,
+                    "container": {},
+                }
+            ]
+        }
+        with self.assertLogs(level='WARNING') as capture_logs:
+            assert self.batch_client.get_job_awslogs_info(JOB_ID) is None
+            assert len(capture_logs.records) == 1
+
+    def test_job_splunk_logs(self):
+        self.client_mock.describe_jobs.return_value = {
+            "jobs": [
+                {
+                    "jobId": JOB_ID,
+                    "logStreamName": LOG_STREAM_NAME,
+                    "container": {
+                        "logConfiguration": {
+                            "logDriver": "splunk",
+                        }
+                    },
+                }
+            ]
+        }
+        with self.assertLogs(level='WARNING') as capture_logs:
+            assert self.batch_client.get_job_awslogs_info(JOB_ID) is None
+            assert len(capture_logs.records) == 1
+
 
 class TestBatchClientDelays(unittest.TestCase):
     @mock.patch.dict("os.environ", AWS_DEFAULT_REGION=AWS_REGION)