shahar1 commented on code in PR #75915:
URL: https://github.com/apache/airflow/pull/75915#discussion_r4541611289
##########
providers/amazon/tests/unit/amazon/aws/triggers/test_athena.py:
##########
@@ -40,3 +43,53 @@ def test_serialization(self):
"waiter_max_attempts": 60,
"aws_conn_id": "aws_default",
}
+
+
+class TestAthenaSparkTrigger:
+ @pytest.mark.parametrize(
+ "aws_kwargs",
+ [
+ pytest.param({}, id="defaults"),
+ pytest.param(
+ {
+ "region_name": "eu-west-1",
+ "verify": False,
+ "botocore_config": {"read_timeout": 42},
+ },
+ id="custom-aws-settings",
+ ),
+ ],
+ )
+ def test_serialization(self, aws_kwargs):
+ expected_kwargs = {
+ "calculation_execution_id": "test_calculation_execution_id",
+ "waiter_delay": 30,
+ "waiter_max_attempts": 60,
+ "aws_conn_id": "aws_default",
+ **aws_kwargs,
+ }
+ trigger = AthenaSparkTrigger(**expected_kwargs)
+
+ classpath, kwargs = trigger.serialize()
+
+ assert classpath ==
"airflow.providers.amazon.aws.triggers.athena.AthenaSparkTrigger"
+ assert kwargs == expected_kwargs
+ assert AthenaSparkTrigger(**kwargs).serialize() == (classpath, kwargs)
+
+ def test_waiter_configuration(self):
Review Comment:
**minor**: Testing. From `.github/instructions/code-review.instructions.md`
§ Testing Requirements:
> **Flag any changed or added behaviour without a corresponding test, and
flag tests that a reviewer cannot fail by reverting the PR's change.** The
target is exactly 100% coverage of what the PR changes — no more, no less.
The new `spark_calculation_terminal` waiter treats `FAILED` and `CANCELED`
as `success` acceptors. That lets the trigger return so `execute_complete` can
report the state-change reason through `_handle_terminal_state`. This test only
checks the waiter name, so every test still passes if either acceptor flips to
`failure` or is dropped.
A waiter test in the style of `tests/unit/amazon/aws/waiters/test_glue.py`
would cover it. It should assert that `COMPLETED`, `FAILED` and `CANCELED` all
end the wait, and that a non-terminal state such as `RUNNING` keeps polling.
##########
providers/amazon/src/airflow/providers/amazon/aws/waiters/athena.json:
##########
@@ -81,6 +81,31 @@
"argument": "Status.State"
}
]
+ },
+ "spark_calculation_terminal": {
+ "operation": "GetCalculationExecutionStatus",
Review Comment:
**minor**: The deferrable path polls `GetCalculationExecutionStatus`. The
sync path (the `calculation_complete` waiter) and every other call this
operator makes use `GetCalculationExecution`. These are separate IAM actions.
So a least-privilege role that works with `deferrable=False` gets
`AccessDenied` here: the trigger yields an error event, `execute_complete`
stops the calculation, and the task fails.
Both responses carry `Status.State`. Using the same operation keeps the IAM
permissions the same in both modes:
```suggestion
"operation": "GetCalculationExecution",
```
##########
providers/amazon/src/airflow/providers/amazon/aws/triggers/athena.py:
##########
@@ -66,3 +66,55 @@ def hook(self) -> AwsGenericHook:
verify=self.verify,
config=self.botocore_config,
)
+
+
+class AthenaSparkTrigger(AwsBaseWaiterTrigger):
+ """
+ Wait asynchronously for an Athena Spark calculation to reach a terminal
state.
+
+ :param calculation_execution_id: Calculation execution ID to monitor.
+ :param waiter_delay: Seconds between status checks. Defaults to 30.
+ :param waiter_max_attempts: Maximum number of polling attempts. Defaults
to 120.
+ :param aws_conn_id: Airflow connection used for AWS credentials. Defaults
to ``aws_default``.
+ """
+
+ aws_hook_class = AthenaHook
+
+ def __init__(
+ self,
+ *,
+ calculation_execution_id: str,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 120,
+ aws_conn_id: str | None = "aws_default",
+ region_name: str | None = None,
+ verify: bool | str | None = None,
+ botocore_config: dict | None = None,
+ ) -> None:
+ super().__init__(
+ serialized_fields={
+ "calculation_execution_id": calculation_execution_id,
+ },
+ waiter_name="spark_calculation_terminal",
+ waiter_args={
+ "CalculationExecutionId": calculation_execution_id,
+ },
+ failure_message=(f"Error waiting for Athena Spark calculation
{calculation_execution_id}"),
+ status_message=(f"Athena Spark calculation execution ID:
{calculation_execution_id}"),
+ status_queries=["Status.State", "Status.StateChangeReason"],
+ return_value=calculation_execution_id,
+ waiter_delay=waiter_delay,
+ waiter_max_attempts=waiter_max_attempts,
+ aws_conn_id=aws_conn_id,
+ region_name=region_name,
+ verify=verify,
+ botocore_config=botocore_config,
+ )
+
+ def hook(self) -> AwsGenericHook:
Review Comment:
**nit**: `AwsBaseWaiterTrigger.hook()` already returns
`self.aws_hook_class(**self._hook_parameters)`. Those parameters map to exactly
these four arguments, and `aws_hook_class = AthenaHook` is set above. You can
drop this override.
##########
providers/amazon/src/airflow/providers/amazon/aws/operators/athena_spark.py:
##########
@@ -135,6 +157,42 @@ def execute(self, context: Context) -> dict[str, Any]:
return self._handle_terminal_state(calculation_execution_id,
final_state)
+ def execute_complete(
+ self,
+ context: Context,
+ calculation_execution_id: str,
+ event: dict[str, Any] | None = None,
+ ) -> dict[str, Any]:
+ """Handle the trigger result and return calculation metadata."""
+ self._calculation_execution_id = calculation_execution_id
+
Review Comment:
**nit**: Raising inside `try` only to catch it a few lines later is harder
to follow than a single guard. For example:
```python
if event is None or event.get("status") != "success":
self._stop_calculation(calculation_execution_id)
message = "event is None" if event is None else event.get("message",
"Unknown trigger error.")
raise RuntimeError(f"Error waiting for Athena Spark calculation
{calculation_execution_id}: {message}")
```
This keeps the same messages, so the `test_execute_complete_trigger_error`
cases still match.
##########
providers/amazon/src/airflow/providers/amazon/aws/operators/athena_spark.py:
##########
@@ -49,11 +51,14 @@ class AthenaSparkOperator(AwsBaseOperator[AthenaHook]):
:param waiter_delay: Seconds to wait between status checks. Defaults to 30.
:param waiter_max_attempts: Maximum number of polling attempts before
timing out. Defaults to 120.
To limit total task time, use execution_timeout on the task as well.
+ :param deferrable: Whether to wait asynchronously using the triggerer.
+ Defaults to the ``operators.default_deferrable`` configuration value
(False).
:param log_query: Whether to log submission details. Defaults to True.
:param aws_conn_id: The Airflow connection used for AWS credentials.
Defaults to ``aws_default``.
:param region_name: AWS region. If not set, default boto3 behavior is
used. Defaults to None.
:param verify: Whether to verify SSL certificates. Defaults to None.
:param botocore_config: Optional botocore configuration dict. Defaults to
None.
+
Review Comment:
**nit**: Stray blank line before the closing `"""`.
```suggestion
```
##########
providers/amazon/src/airflow/providers/amazon/aws/triggers/athena.py:
##########
@@ -66,3 +66,55 @@ def hook(self) -> AwsGenericHook:
verify=self.verify,
config=self.botocore_config,
)
+
+
+class AthenaSparkTrigger(AwsBaseWaiterTrigger):
+ """
+ Wait asynchronously for an Athena Spark calculation to reach a terminal
state.
+
+ :param calculation_execution_id: Calculation execution ID to monitor.
+ :param waiter_delay: Seconds between status checks. Defaults to 30.
+ :param waiter_max_attempts: Maximum number of polling attempts. Defaults
to 120.
+ :param aws_conn_id: Airflow connection used for AWS credentials. Defaults
to ``aws_default``.
Review Comment:
**nit**: The docstring stops at `aws_conn_id`. The trigger also accepts and
serializes `region_name`, `verify` and `botocore_config`. Please document them
as the operator docstring does.
--
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]
