SameerMesiah97 commented on code in PR #75915:
URL: https://github.com/apache/airflow/pull/75915#discussion_r4569123073
##########
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:
waiter tests have been added following the pattern for the glue waiter.
##########
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:
Removed.
--
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]
