vandonr-amz commented on code in PR #32563:
URL: https://github.com/apache/airflow/pull/32563#discussion_r1261786859
##########
airflow/providers/amazon/aws/operators/step_function.py:
##########
@@ -74,9 +88,25 @@ def execute(self, context: Context):
raise AirflowException(f"Failed to start State Machine execution
for: {self.state_machine_arn}")
self.log.info("Started State Machine execution for %s: %s",
self.state_machine_arn, execution_arn)
-
+ if self.deferrable:
+ self.defer(
+ trigger=StepFunctionsStartExecutionTrigger(
+ execution_arn=execution_arn,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ aws_conn_id=self.aws_conn_id,
+ region_name=self.region_name,
+ ),
+ method_name="execute_complete",
+ timeout=timedelta(seconds=self.waiter_max_attempts *
self.waiter_delay),
Review Comment:
no extra time buffer here ? I seem to recall we gave it 30 extra seconds in
other operators ?
##########
airflow/providers/amazon/aws/operators/step_function.py:
##########
@@ -74,9 +88,25 @@ def execute(self, context: Context):
raise AirflowException(f"Failed to start State Machine execution
for: {self.state_machine_arn}")
self.log.info("Started State Machine execution for %s: %s",
self.state_machine_arn, execution_arn)
-
+ if self.deferrable:
+ self.defer(
+ trigger=StepFunctionsStartExecutionTrigger(
+ execution_arn=execution_arn,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ aws_conn_id=self.aws_conn_id,
+ region_name=self.region_name,
+ ),
+ method_name="execute_complete",
+ timeout=timedelta(seconds=self.waiter_max_attempts *
self.waiter_delay),
+ )
return execution_arn
+ def execute_complete(self, context: Context, event: dict[str, Any] | None
= None) -> None:
+ if event and event["status"] == "success":
Review Comment:
I don't really know in which case event could be None here, but I think it's
not supposed to happen ? So we shouldn't plan around it, we'd rather have it
blow ?
##########
airflow/providers/amazon/aws/triggers/stepfunction.py:
##########
@@ -0,0 +1,61 @@
+# 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.hooks.base_aws import AwsGenericHook
+from airflow.providers.amazon.aws.hooks.step_function import StepFunctionHook
+from airflow.providers.amazon.aws.triggers.base import AwsBaseWaiterTrigger
+
+
+class StepFunctionsStartExecutionTrigger(AwsBaseWaiterTrigger):
Review Comment:
```suggestion
class StepFunctionsExecutionCompleteTrigger(AwsBaseWaiterTrigger):
```
we tend to name triggers after the state they are waiting on I think, rather
than after the operator creating them ?
##########
airflow/providers/amazon/aws/waiters/stepfunctions.json:
##########
@@ -0,0 +1,36 @@
+{
+ "version": 2,
+ "waiters": {
+ "step_function_succeeded": {
+ "operation": "DescribeExecution",
+ "delay": 30,
+ "maxAttempts": 60,
+ "acceptors": [
+ {
+ "matcher": "path",
+ "argument": "status",
+ "expected": "SUCCEEDED",
+ "state": "success"
+ },
+ {
+ "matcher": "error",
+ "argument": "status",
+ "expected": "RUNNING",
+ "state": "retry"
+ },
Review Comment:
what is this ? It's the first time I see it in a waiter definition 🤨
What does it do that the waiter wouldn't do without those lines ?
##########
airflow/providers/amazon/aws/waiters/stepfunctions.json:
##########
@@ -0,0 +1,36 @@
+{
+ "version": 2,
+ "waiters": {
+ "step_function_succeeded": {
+ "operation": "DescribeExecution",
+ "delay": 30,
+ "maxAttempts": 60,
+ "acceptors": [
+ {
+ "matcher": "path",
+ "argument": "status",
+ "expected": "SUCCEEDED",
+ "state": "success"
+ },
+ {
+ "matcher": "error",
+ "argument": "status",
+ "expected": "RUNNING",
+ "state": "retry"
+ },
Review Comment:
what is this ? It's the first time I see it in a waiter definition 🤨
What does it do that the waiter wouldn't do without those lines ?
##########
airflow/providers/amazon/aws/operators/step_function.py:
##########
@@ -42,6 +45,11 @@ class StepFunctionStartExecutionOperator(BaseOperator):
:param state_machine_input: JSON data input to pass to the State Machine
:param aws_conn_id: aws connection to uses
:param do_xcom_push: if True, execution_arn is pushed to XCom with key
execution_arn.
+ :param waiter_max_attempts: Maximum number of attempts to poll the
execution.
+ :param waiter_delay: Number of seconds between polling the state of the
execution.
+ :param deferrable: If True, the operator will wait asynchronously for the
job to complete.
+ This implies waiting for completion. This mode requires aiobotocore
module to be installed.
+ (default: False)
Review Comment:
```suggestion
```
this default is not exactly true anymore
##########
airflow/providers/amazon/aws/triggers/stepfunction.py:
##########
@@ -0,0 +1,61 @@
+# 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.hooks.base_aws import AwsGenericHook
+from airflow.providers.amazon.aws.hooks.step_function import StepFunctionHook
+from airflow.providers.amazon.aws.triggers.base import AwsBaseWaiterTrigger
+
+
+class StepFunctionsStartExecutionTrigger(AwsBaseWaiterTrigger):
+ """
+ Trigger to poll for the completion of a Step Functions execution.
+
+ :param execution_arn: ARN of the state machine to poll
+ :param waiter_delay: The amount of time in seconds to wait between
attempts.
+ :param waiter_max_attempts: The maximum number of attempts to be made.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ """
+
+ def __init__(
+ self,
+ *,
+ execution_arn: str,
+ waiter_delay: int = 60,
+ waiter_max_attempts: int = 30,
+ aws_conn_id: str | None = None,
+ region_name: str | None = None,
+ ):
+ self.aws_conn_id = aws_conn_id
+ self.region_name = region_name
+
+ super().__init__(
+ serialized_fields={"execution_arn": execution_arn, "region_name":
region_name},
+ waiter_name="step_function_succeeded",
+ waiter_args={"executionArn": execution_arn},
+ failure_message="Step function failed",
+ status_message="Status of step function execution is",
+ status_queries=["status"],
Review Comment:
```suggestion
status_queries=["status", "error", "cause"],
```
I think those would be useful too
https://docs.aws.amazon.com/step-functions/latest/apireference/API_DescribeExecution.html
##########
airflow/providers/amazon/aws/triggers/stepfunction.py:
##########
@@ -0,0 +1,61 @@
+# 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.hooks.base_aws import AwsGenericHook
+from airflow.providers.amazon.aws.hooks.step_function import StepFunctionHook
+from airflow.providers.amazon.aws.triggers.base import AwsBaseWaiterTrigger
+
+
+class StepFunctionsStartExecutionTrigger(AwsBaseWaiterTrigger):
+ """
+ Trigger to poll for the completion of a Step Functions execution.
+
+ :param execution_arn: ARN of the state machine to poll
+ :param waiter_delay: The amount of time in seconds to wait between
attempts.
+ :param waiter_max_attempts: The maximum number of attempts to be made.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ """
+
+ def __init__(
+ self,
+ *,
+ execution_arn: str,
+ waiter_delay: int = 60,
+ waiter_max_attempts: int = 30,
+ aws_conn_id: str | None = None,
+ region_name: str | None = None,
+ ):
+ self.aws_conn_id = aws_conn_id
+ self.region_name = region_name
Review Comment:
you don't need to do that, `super` handles it
--
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]
