ramitkataria commented on code in PR #73935:
URL: https://github.com/apache/airflow/pull/73935#discussion_r4407843255
##########
providers/amazon/src/airflow/providers/amazon/aws/sensors/mwaa_serverless.py:
##########
@@ -53,13 +65,17 @@ def __init__(
run_id: str,
success_states: set[str] | None = None,
failure_states: set[str] | None = None,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable",
fallback=False),
+ max_retries: int = 720,
**kwargs,
) -> None:
super().__init__(**kwargs)
self.workflow_arn = workflow_arn
self.run_id = run_id
- self.success_states = success_states or {"SUCCESS"}
- self.failure_states = failure_states or {"FAILED", "TIMEOUT",
"STOPPED"}
+ self.success_states = success_states or set(DEFAULT_SUCCESS_STATES)
Review Comment:
Small regression here: with `success_states=["SUCCESS", "STOPPED"]` (a list)
and no `failure_states`, `DEFAULT_FAILURE_STATES - success_states` raises
`TypeError` when the Dag is parsed. That worked before this PR. Could we wrap
it in `set(...)` like the trigger does (line 69), with a test that passes a
list?
While you're here, could the success/failure overlap `ValueError` move into
`__init__`, like `MwaaDagRunSensor` (`sensors/mwaa.py:106`)? Right now poke
mode quietly treats an overlapping state as a failure, and deferrable mode only
errors at execute time.
##########
providers/amazon/src/airflow/providers/amazon/aws/triggers/mwaa_serverless.py:
##########
@@ -0,0 +1,124 @@
+# 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.
+"""Amazon MWAA Serverless triggers."""
+
+from __future__ import annotations
+
+from collections.abc import Collection
+from typing import Any
+
+from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+from airflow.providers.amazon.aws.triggers.base import AwsBaseWaiterTrigger
+
+WORKFLOW_RUN_STATES = frozenset(
+ {"STARTING", "QUEUED", "RUNNING", "STOPPING", "SUCCESS", "FAILED",
"TIMEOUT", "STOPPED"}
+)
+DEFAULT_SUCCESS_STATES = frozenset({"SUCCESS"})
+DEFAULT_FAILURE_STATES = frozenset({"FAILED", "TIMEOUT", "STOPPED"})
+
+
+def get_failure_states(success_states: set[str], failure_states:
Collection[str] | None) -> set[str]:
+ """Return ``failure_states``, or the defaults minus any state listed in
``success_states``."""
+ if failure_states:
+ return set(failure_states)
+ return set(DEFAULT_FAILURE_STATES - success_states)
+
+
+class MwaaServerlessWorkflowRunCompletedTrigger(AwsBaseWaiterTrigger):
+ """
+ Trigger when an Amazon MWAA Serverless workflow run reaches a terminal
state.
+
+ :param workflow_arn: The ARN of the workflow.
+ :param run_id: The ID of the workflow run to wait for.
+ :param success_states: Collection of run states that mark the run as
successful.
+ Default: ``{"SUCCESS"}``.
+ :param failure_states: Collection of run states that mark the run as
failed.
+ Default: ``{"FAILED", "TIMEOUT", "STOPPED"}``, minus any state listed
in ``success_states``.
+ :param waiter_delay: The amount of time in seconds to wait between
attempts. (default: 60)
+ :param waiter_max_attempts: The maximum number of attempts to be made.
(default: 720)
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ """
+
+ aws_hook_class = AwsBaseHook
+
+ def __init__(
+ self,
+ *,
+ workflow_arn: str,
+ run_id: str,
+ success_states: Collection[str] | None = None,
+ failure_states: Collection[str] | None = None,
+ waiter_delay: int = 60,
+ waiter_max_attempts: int = 720,
+ **kwargs,
+ ) -> None:
+ self.success_states = set(success_states) if success_states else
set(DEFAULT_SUCCESS_STATES)
+ self.failure_states = get_failure_states(self.success_states,
failure_states)
+
+ if self.success_states & self.failure_states:
+ raise ValueError("success_states and failure_states must not have
any values in common")
+
+ in_progress_states = WORKFLOW_RUN_STATES - self.success_states -
self.failure_states
+
+ super().__init__(
+ serialized_fields={
+ "workflow_arn": workflow_arn,
+ "run_id": run_id,
+ "success_states": sorted(self.success_states),
+ "failure_states": sorted(self.failure_states),
+ },
+ waiter_name="workflow_run_complete",
+ waiter_args={"WorkflowArn": workflow_arn, "RunId": run_id},
+ failure_message=f"MWAA Serverless workflow run {run_id} of
{workflow_arn} failed",
+ status_message="State of workflow run",
+ status_queries=["RunDetail.RunState", "RunDetail.ErrorMessage"],
+ return_key="run_id",
+ return_value=run_id,
+ waiter_delay=waiter_delay,
+ waiter_max_attempts=waiter_max_attempts,
+ waiter_config_overrides={
+ "acceptors": _build_waiter_acceptors(
+ success_states=self.success_states,
+ failure_states=self.failure_states,
+ in_progress_states=in_progress_states,
+ )
+ },
+ **kwargs,
+ )
+
+ @property
+ def _hook_parameters(self) -> dict[str, Any]:
+ return {**super()._hook_parameters, "client_type": "mwaa-serverless"}
+
+
+def _build_waiter_acceptors(
Review Comment:
nit: this is nearly identical to `_build_waiter_acceptors` in
`triggers/mwaa.py:195`; only `argument` differs. Could they share one helper
that takes `argument`?
##########
providers/amazon/src/airflow/providers/amazon/aws/sensors/mwaa_serverless.py:
##########
@@ -75,3 +91,28 @@ def poke(self, context: Context) -> bool:
raise RuntimeError(f"Workflow run {self.run_id} failed with state
{state}: {error_msg}")
return state in self.success_states
+
+ def execute(self, context: Context) -> None:
+ if self.deferrable:
+ self.defer(
Review Comment:
Could deferrable mode behave the same as poke mode? A few gaps:
- `defer()` has no `timeout=`, so `timeout` is ignored once the sensor
defers. The example's `timeout=600` turns into 720 × 30 s = 6 h. Most Amazon
sensors pass `timeout=timedelta(seconds=self.timeout)` (e.g.
`sensors/emr.py:165`).
- `execute_complete` (line 118) raises `RuntimeError`, but
`BaseSensorOperator.resume_execution` only applies `soft_fail`/`never_fail` to
`AirflowException`/`TaskDeferralError`, so both flags stop working after
deferral. A small dedicated `AirflowException` subclass would fix this
(`check-no-new-airflow-exceptions` rules out raising `AirflowException`
directly).
- `int(self.poke_interval)` turns a sub-second interval into 0, which makes
the trigger poll in a tight loop.
##########
providers/amazon/src/airflow/providers/amazon/aws/operators/mwaa_serverless.py:
##########
@@ -80,11 +103,49 @@ def execute(self, context: Context) -> str:
"WorkflowArn": self.workflow_arn,
"OverrideParameters": self.override_parameters,
"WorkflowVersion": self.workflow_version,
+ "ClientToken": self.client_token,
}
)
response = self.hook.conn.start_workflow_run(**kwargs)
run_id = response["RunId"]
self.log.info("Started workflow run %s (status: %s)", run_id,
response.get("Status"))
+
+ if not self.wait_for_completion:
+ return run_id
+
+ if self.deferrable:
+ self.log.info("Deferring until workflow run %s completes", run_id)
+ self.defer(
+ trigger=MwaaServerlessWorkflowRunCompletedTrigger(
+ workflow_arn=self.workflow_arn,
+ run_id=run_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ aws_conn_id=self.aws_conn_id,
+ region_name=self.region_name,
+ verify=self.verify,
+ botocore_config=self.botocore_config,
+ ),
+ method_name="execute_complete",
+ )
+ else:
+ self.log.info("Waiting for workflow run %s to complete", run_id)
+ self.hook.get_waiter("workflow_run_complete").wait(
Review Comment:
Now that `wait_for_completion=True` can block for hours, could we add an
`on_kill` that stops the run (`stop_workflow_run`), like the Glue/EMR/ECS
operators do? Otherwise a killed task leaves the workflow running.
Related: passing `timeout=` to `defer()` at line 118 (EMR uses
`waiter_max_attempts * waiter_delay + 60`) would stop a stuck trigger from
leaving the task deferred indefinitely.
--
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.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
