You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by po...@apache.org on 2025/04/13 13:31:21 UTC
(airflow) branch main updated: AWS Batch Operators/Sensors inherit AWS Base classes (#49172)
This is an automated email from the ASF dual-hosted git repository.
potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 1781dc9d5fc AWS Batch Operators/Sensors inherit AWS Base classes (#49172)
1781dc9d5fc is described below
commit 1781dc9d5fcbb93bce7018a4050ddd36ec1a65d5
Author: ellisms <11...@users.noreply.github.com>
AuthorDate: Sun Apr 13 09:31:14 2025 -0400
AWS Batch Operators/Sensors inherit AWS Base classes (#49172)
* AWS Batch operators inherit base AWS class
* Update AWS Batch sensor to inherit AWS base class
---------
Co-authored-by: mse139 <ms...@users.noereply.github.com>
---
providers/amazon/docs/operators/batch.rst | 5 ++
.../providers/amazon/aws/hooks/batch_client.py | 11 +++
.../providers/amazon/aws/operators/batch.py | 75 +++++++++----------
.../airflow/providers/amazon/aws/sensors/batch.py | 85 ++++++++--------------
.../tests/unit/amazon/aws/operators/test_batch.py | 40 +++++-----
5 files changed, 104 insertions(+), 112 deletions(-)
diff --git a/providers/amazon/docs/operators/batch.rst b/providers/amazon/docs/operators/batch.rst
index 30e3999610a..b992e26e4b0 100644
--- a/providers/amazon/docs/operators/batch.rst
+++ b/providers/amazon/docs/operators/batch.rst
@@ -29,6 +29,11 @@ Prerequisite Tasks
.. include:: ../_partials/prerequisite_tasks.rst
+Generic Parameters
+------------------
+
+.. include:: ../_partials/generic_parameters.rst
+
Operators
---------
diff --git a/providers/amazon/src/airflow/providers/amazon/aws/hooks/batch_client.py b/providers/amazon/src/airflow/providers/amazon/aws/hooks/batch_client.py
index 4d85898d917..0dcedfce084 100644
--- a/providers/amazon/src/airflow/providers/amazon/aws/hooks/batch_client.py
+++ b/providers/amazon/src/airflow/providers/amazon/aws/hooks/batch_client.py
@@ -142,6 +142,17 @@ class BatchProtocol(Protocol):
"""
...
+ def create_compute_environment(self, **kwargs) -> dict:
+ """
+ Create an AWS Batch compute environment.
+
+ :param kwargs: Arguments for boto3 create_compute_environment
+
+ .. seealso::
+ - https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/batch/client/create_compute_environment.html
+ """
+ ...
+
# Note that the use of invalid-name parameters should be restricted to the boto3 mappings only;
# all the Airflow wrappers of boto3 clients should not adopt invalid-names to match boto3.
diff --git a/providers/amazon/src/airflow/providers/amazon/aws/operators/batch.py b/providers/amazon/src/airflow/providers/amazon/aws/operators/batch.py
index a41d1ec98f4..49669f77f62 100644
--- a/providers/amazon/src/airflow/providers/amazon/aws/operators/batch.py
+++ b/providers/amazon/src/airflow/providers/amazon/aws/operators/batch.py
@@ -28,12 +28,10 @@ from __future__ import annotations
from collections.abc import Sequence
from datetime import timedelta
-from functools import cached_property
from typing import TYPE_CHECKING, Any
from airflow.configuration import conf
from airflow.exceptions import AirflowException
-from airflow.models import BaseOperator
from airflow.models.mappedoperator import MappedOperator
from airflow.providers.amazon.aws.hooks.batch_client import BatchClientHook
from airflow.providers.amazon.aws.links.batch import (
@@ -42,18 +40,20 @@ from airflow.providers.amazon.aws.links.batch import (
BatchJobQueueLink,
)
from airflow.providers.amazon.aws.links.logs import CloudWatchEventsLink
+from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator
from airflow.providers.amazon.aws.triggers.batch import (
BatchCreateComputeEnvironmentTrigger,
BatchJobTrigger,
)
from airflow.providers.amazon.aws.utils import trim_none_values, validate_execute_complete_event
+from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
from airflow.providers.amazon.aws.utils.task_log_fetcher import AwsTaskLogFetcher
if TYPE_CHECKING:
from airflow.utils.context import Context
-class BatchOperator(BaseOperator):
+class BatchOperator(AwsBaseOperator[BatchClientHook]):
"""
Execute a job on AWS Batch.
@@ -83,10 +83,14 @@ class BatchOperator(BaseOperator):
polling is only used when waiters is None
:param status_retries: number of HTTP retries to get job status, 10;
polling is only used when waiters is None
- :param aws_conn_id: connection id of AWS credentials / region name. If None,
- credential boto3 strategy will be used.
- :param region_name: region name to use in AWS Hook.
- Override the region_name in connection (if provided)
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+ :param verify: Whether or not to verify SSL certificates. See:
+ https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html
:param tags: collection of tags to apply to the AWS Batch job submission
if None, no tags are submitted
:param deferrable: Run operator in the deferrable mode.
@@ -106,9 +110,11 @@ class BatchOperator(BaseOperator):
waiter = waiters.get_waiter("JobComplete")
"""
+ aws_hook_class = BatchClientHook
+
ui_color = "#c3dae0"
arn: str | None = None
- template_fields: Sequence[str] = (
+ template_fields: Sequence[str] = aws_template_fields(
"job_id",
"job_name",
"job_definition",
@@ -177,8 +183,6 @@ class BatchOperator(BaseOperator):
waiters: Any | None = None,
max_retries: int = 4200,
status_retries: int | None = None,
- aws_conn_id: str | None = None,
- region_name: str | None = None,
tags: dict | None = None,
wait_for_completion: bool = True,
deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
@@ -188,7 +192,7 @@ class BatchOperator(BaseOperator):
submit_job_timeout: int | None = None,
**kwargs,
) -> None:
- BaseOperator.__init__(self, **kwargs)
+ super().__init__(**kwargs)
self.job_id = job_id
self.job_name = job_name
self.job_definition = job_definition
@@ -215,17 +219,14 @@ class BatchOperator(BaseOperator):
# params for hook
self.max_retries = max_retries
self.status_retries = status_retries
- self.aws_conn_id = aws_conn_id
- self.region_name = region_name
-
- @cached_property
- def hook(self) -> BatchClientHook:
- return BatchClientHook(
- max_retries=self.max_retries,
- status_retries=self.status_retries,
- aws_conn_id=self.aws_conn_id,
- region_name=self.region_name,
- )
+
+ @property
+ def _hook_parameters(self):
+ return {
+ **super()._hook_parameters,
+ "max_retries": self.max_retries,
+ "status_retries": self.status_retries,
+ }
def execute(self, context: Context) -> str | None:
"""
@@ -434,7 +435,7 @@ class BatchOperator(BaseOperator):
)
-class BatchCreateComputeEnvironmentOperator(BaseOperator):
+class BatchCreateComputeEnvironmentOperator(AwsBaseOperator[BatchClientHook]):
"""
Create an AWS Batch compute environment.
@@ -460,15 +461,21 @@ class BatchCreateComputeEnvironmentOperator(BaseOperator):
Only useful when deferrable is True.
:param max_retries: How many times to poll for the environment status.
Only useful when deferrable is True.
- :param aws_conn_id: Connection ID of AWS credentials / region name. If None,
- credential boto3 strategy will be used.
- :param region_name: Region name to use in AWS Hook. Overrides the
- ``region_name`` in connection if provided.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+ :param verify: Whether or not to verify SSL certificates. See:
+ https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html
:param deferrable: If True, the operator will wait asynchronously for the environment to be created.
This mode requires aiobotocore module to be installed. (default: False)
"""
- template_fields: Sequence[str] = (
+ aws_hook_class = BatchClientHook
+
+ template_fields: Sequence[str] = aws_template_fields(
"compute_environment_name",
"compute_resources",
"service_role",
@@ -486,8 +493,6 @@ class BatchCreateComputeEnvironmentOperator(BaseOperator):
tags: dict | None = None,
poll_interval: int = 30,
max_retries: int | None = None,
- aws_conn_id: str | None = None,
- region_name: str | None = None,
deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
**kwargs,
):
@@ -502,18 +507,8 @@ class BatchCreateComputeEnvironmentOperator(BaseOperator):
self.tags = tags or {}
self.poll_interval = poll_interval
self.max_retries = max_retries or 120
- self.aws_conn_id = aws_conn_id
- self.region_name = region_name
self.deferrable = deferrable
- @cached_property
- def hook(self):
- """Create and return a BatchClientHook."""
- return BatchClientHook(
- aws_conn_id=self.aws_conn_id,
- region_name=self.region_name,
- )
-
def execute(self, context: Context):
"""Create an AWS batch compute environment."""
kwargs: dict[str, Any] = {
diff --git a/providers/amazon/src/airflow/providers/amazon/aws/sensors/batch.py b/providers/amazon/src/airflow/providers/amazon/aws/sensors/batch.py
index 2fed9d34500..b368edcd093 100644
--- a/providers/amazon/src/airflow/providers/amazon/aws/sensors/batch.py
+++ b/providers/amazon/src/airflow/providers/amazon/aws/sensors/batch.py
@@ -18,20 +18,20 @@ from __future__ import annotations
from collections.abc import Sequence
from datetime import timedelta
-from functools import cached_property
from typing import TYPE_CHECKING, Any
from airflow.configuration import conf
from airflow.exceptions import AirflowException
from airflow.providers.amazon.aws.hooks.batch_client import BatchClientHook
+from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
from airflow.providers.amazon.aws.triggers.batch import BatchJobTrigger
-from airflow.sensors.base import BaseSensorOperator
+from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
if TYPE_CHECKING:
from airflow.utils.context import Context
-class BatchSensor(BaseSensorOperator):
+class BatchSensor(AwsBaseSensor[BatchClientHook]):
"""
Poll the state of the Batch Job until it reaches a terminal state; fails if the job fails.
@@ -40,19 +40,24 @@ class BatchSensor(BaseSensorOperator):
:ref:`howto/sensor:BatchSensor`
:param job_id: Batch job_id to check the state for
- :param aws_conn_id: aws connection to use, defaults to 'aws_default'
- If this is None or empty then the default boto3 behaviour is used. If
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
running Airflow in a distributed manner and aws_conn_id is None or
empty, then default boto3 configuration would be used (and must be
maintained on each worker node).
- :param region_name: aws region name associated with the client
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+ :param verify: Whether or not to verify SSL certificates. See:
+ https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html
:param deferrable: Run sensor in the deferrable mode.
:param poke_interval: polling period in seconds to check for the status of the job.
:param max_retries: Number of times to poll for job state before
returning the current state.
"""
- template_fields: Sequence[str] = ("job_id",)
+ aws_hook_class = BatchClientHook
+ template_fields: Sequence[str] = aws_template_fields(
+ "job_id",
+ )
template_ext: Sequence[str] = ()
ui_color = "#66c3ff"
@@ -60,8 +65,6 @@ class BatchSensor(BaseSensorOperator):
self,
*,
job_id: str,
- aws_conn_id: str | None = "aws_default",
- region_name: str | None = None,
deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
poke_interval: float = 30,
max_retries: int = 4200,
@@ -69,8 +72,6 @@ class BatchSensor(BaseSensorOperator):
):
super().__init__(**kwargs)
self.job_id = job_id
- self.aws_conn_id = aws_conn_id
- self.region_name = region_name
self.deferrable = deferrable
self.poke_interval = poke_interval
self.max_retries = max_retries
@@ -119,15 +120,8 @@ class BatchSensor(BaseSensorOperator):
job_id = event["job_id"]
self.log.info("Batch Job %s complete", job_id)
- @cached_property
- def hook(self) -> BatchClientHook:
- return BatchClientHook(
- aws_conn_id=self.aws_conn_id,
- region_name=self.region_name,
- )
-
-class BatchComputeEnvironmentSensor(BaseSensorOperator):
+class BatchComputeEnvironmentSensor(AwsBaseSensor[BatchClientHook]):
"""
Poll the state of the Batch environment until it reaches a terminal state; fails if the environment fails.
@@ -137,38 +131,31 @@ class BatchComputeEnvironmentSensor(BaseSensorOperator):
:param compute_environment: Batch compute environment name
- :param aws_conn_id: aws connection to use, defaults to 'aws_default'
- If this is None or empty then the default boto3 behaviour is used. If
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
running Airflow in a distributed manner and aws_conn_id is None or
empty, then default boto3 configuration would be used (and must be
maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+ :param verify: Whether or not to verify SSL certificates. See:
+ https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html
- :param region_name: aws region name associated with the client
"""
- template_fields: Sequence[str] = ("compute_environment",)
+ aws_hook_class = BatchClientHook
+ template_fields: Sequence[str] = aws_template_fields(
+ "compute_environment",
+ )
template_ext: Sequence[str] = ()
ui_color = "#66c3ff"
def __init__(
self,
compute_environment: str,
- aws_conn_id: str | None = "aws_default",
- region_name: str | None = None,
**kwargs,
):
super().__init__(**kwargs)
self.compute_environment = compute_environment
- self.aws_conn_id = aws_conn_id
- self.region_name = region_name
-
- @cached_property
- def hook(self) -> BatchClientHook:
- """Create and return a BatchClientHook."""
- return BatchClientHook(
- aws_conn_id=self.aws_conn_id,
- region_name=self.region_name,
- )
def poke(self, context: Context) -> bool:
response = self.hook.client.describe_compute_environments( # type: ignore[union-attr]
@@ -191,7 +178,7 @@ class BatchComputeEnvironmentSensor(BaseSensorOperator):
)
-class BatchJobQueueSensor(BaseSensorOperator):
+class BatchJobQueueSensor(AwsBaseSensor[BatchClientHook]):
"""
Poll the state of the Batch job queue until it reaches a terminal state; fails if the queue fails.
@@ -204,16 +191,20 @@ class BatchJobQueueSensor(BaseSensorOperator):
:param treat_non_existing_as_deleted: If True, a non-existing Batch job queue is considered as a deleted
queue and as such a valid case.
- :param aws_conn_id: aws connection to use, defaults to 'aws_default'
- If this is None or empty then the default boto3 behaviour is used. If
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
running Airflow in a distributed manner and aws_conn_id is None or
empty, then default boto3 configuration would be used (and must be
maintained on each worker node).
-
- :param region_name: aws region name associated with the client
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+ :param verify: Whether or not to verify SSL certificates. See:
+ https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html
"""
- template_fields: Sequence[str] = ("job_queue",)
+ aws_hook_class = BatchClientHook
+ template_fields: Sequence[str] = aws_template_fields(
+ "job_queue",
+ )
template_ext: Sequence[str] = ()
ui_color = "#66c3ff"
@@ -221,23 +212,11 @@ class BatchJobQueueSensor(BaseSensorOperator):
self,
job_queue: str,
treat_non_existing_as_deleted: bool = False,
- aws_conn_id: str | None = "aws_default",
- region_name: str | None = None,
**kwargs,
):
super().__init__(**kwargs)
self.job_queue = job_queue
self.treat_non_existing_as_deleted = treat_non_existing_as_deleted
- self.aws_conn_id = aws_conn_id
- self.region_name = region_name
-
- @cached_property
- def hook(self) -> BatchClientHook:
- """Create and return a BatchClientHook."""
- return BatchClientHook(
- aws_conn_id=self.aws_conn_id,
- region_name=self.region_name,
- )
def poke(self, context: Context) -> bool:
response = self.hook.client.describe_job_queues( # type: ignore[union-attr]
diff --git a/providers/amazon/tests/unit/amazon/aws/operators/test_batch.py b/providers/amazon/tests/unit/amazon/aws/operators/test_batch.py
index c1b1d847b7d..84e41c5c6d9 100644
--- a/providers/amazon/tests/unit/amazon/aws/operators/test_batch.py
+++ b/providers/amazon/tests/unit/amazon/aws/operators/test_batch.py
@@ -33,6 +33,8 @@ from airflow.providers.amazon.aws.triggers.batch import (
BatchJobTrigger,
)
+from unit.amazon.aws.utils.test_template_fields import validate_template_fields
+
AWS_REGION = "eu-west-1"
AWS_ACCESS_KEY_ID = "airflow_dummy_key"
AWS_SECRET_ACCESS_KEY = "airflow_dummy_secret"
@@ -139,31 +141,13 @@ class TestBatchOperator:
assert batch_job.share_identifier is None
assert batch_job.scheduling_priority_override is None
assert batch_job.hook.region_name is None
- assert batch_job.hook.aws_conn_id is None
assert issubclass(type(batch_job.hook.client), botocore.client.BaseClient)
assert batch_job.tags == {}
assert batch_job.wait_for_completion is True
assert batch_job.submit_job_timeout is None
def test_template_fields_overrides(self):
- assert self.batch.template_fields == (
- "job_id",
- "job_name",
- "job_definition",
- "job_queue",
- "container_overrides",
- "array_properties",
- "ecs_properties_override",
- "eks_properties_override",
- "node_overrides",
- "parameters",
- "retry_strategy",
- "waiters",
- "tags",
- "wait_for_completion",
- "awslogs_enabled",
- "awslogs_fetch_interval",
- )
+ validate_template_fields(self.batch)
@mock.patch.object(BatchClientHook, "get_job_description")
@mock.patch.object(BatchClientHook, "wait_for_job")
@@ -563,6 +547,24 @@ class TestBatchOperator:
class TestBatchCreateComputeEnvironmentOperator:
warn_message = "The `status_retries` parameter is unused and should be removed"
+ def test_template_fields(self):
+ environment_name = "environment_name"
+ compute_resources = {}
+ service_role = "test_role"
+ environment_type = "environment_type"
+ environment_state = "state"
+
+ operator = BatchCreateComputeEnvironmentOperator(
+ task_id="task",
+ compute_environment_name=environment_name,
+ environment_type=environment_type,
+ state=environment_state,
+ compute_resources=compute_resources,
+ service_role=service_role,
+ )
+
+ validate_template_fields(operator)
+
@mock.patch.object(BatchClientHook, "client")
def test_execute(self, mock_conn):
environment_name = "environment_name"