Allow Tags on AWS Batch Job Submission (#13396)

Co-authored-by: Derek Flionis <dflionis@tenable.com>

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

@@ -106,6 +106,7 @@ class AwsBatchProtocol(Protocol):
         arrayProperties: Dict,
         parameters: Dict,
         containerOverrides: Dict,
+        tags: Dict,
     ) -> Dict:
         """
         Submit a batch job
@@ -128,6 +129,9 @@ class AwsBatchProtocol(Protocol):
         :param containerOverrides: the same parameter that boto3 will receive
         :type containerOverrides: Dict
 
+        :param tags: the same parameter that boto3 will receive
+        :type tags: Dict
+
         :return: an API response
         :rtype: Dict
         """

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

@@ -80,6 +80,10 @@ class AwsBatchOperator(BaseOperator):
         Override the region_name in connection (if provided)
     :type region_name: str
 
+    :param tags: collection of tags to apply to the AWS Batch job submission
+        if None, no tags are submitted
+    :type tags: dict
+
     .. note::
         Any custom waiters must return a waiter for these calls:
         .. code-block:: python
@@ -113,6 +117,7 @@ class AwsBatchOperator(BaseOperator):
         status_retries: Optional[int] = None,
         aws_conn_id: Optional[str] = None,
         region_name: Optional[str] = None,
+        tags: Optional[dict] = None,
         **kwargs,
     ):  # pylint: disable=too-many-arguments
 
@@ -125,6 +130,7 @@ class AwsBatchOperator(BaseOperator):
         self.array_properties = array_properties or {}
         self.parameters = parameters or {}
         self.waiters = waiters
+        self.tags = tags or {}
         self.hook = AwsBatchClientHook(
             max_retries=max_retries,
             status_retries=status_retries,
@@ -166,6 +172,7 @@ class AwsBatchOperator(BaseOperator):
                 arrayProperties=self.array_properties,
                 parameters=self.parameters,
                 containerOverrides=self.overrides,
+                tags=self.tags,
             )
             self.job_id = response["jobId"]
 

tests/providers/amazon/aws/operators/test_batch.py

@@ -63,6 +63,7 @@ class TestAwsBatchOperator(unittest.TestCase):
             array_properties=None,
             aws_conn_id='airflow_test',
             region_name="eu-west-1",
+            tags={},
         )
         self.client_mock = self.get_client_type_mock.return_value
         self.assertEqual(self.batch.hook.client, self.client_mock)  # setup client property
@@ -91,6 +92,7 @@ class TestAwsBatchOperator(unittest.TestCase):
         self.assertEqual(self.batch.hook.region_name, "eu-west-1")
         self.assertEqual(self.batch.hook.aws_conn_id, "airflow_test")
         self.assertEqual(self.batch.hook.client, self.client_mock)
+        self.assertEqual(self.batch.tags, {})
 
         self.get_client_type_mock.assert_called_once_with("batch", region_name="eu-west-1")