This is an automated email from the ASF dual-hosted git repository.
shahar1 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 bfe64b1991e Add AthenaSparkSensor (#72229)
bfe64b1991e is described below
commit bfe64b1991e0990c9d638cae094454ed6ff73be9
Author: SameerMesiah97 <[email protected]>
AuthorDate: Sat Oct 17 15:56:06 2026 +0100
Add AthenaSparkSensor (#72229)
Add a sensor for monitoring existing Athena Spark calculations.
Include unit tests, documentation, and an example system test.
---
.../amazon/docs/operators/athena/athena_spark.rst | 25 ++-
providers/amazon/provider.yaml | 1 +
.../airflow/providers/amazon/aws/hooks/athena.py | 7 +
.../providers/amazon/aws/sensors/athena_spark.py | 84 ++++++++++
.../airflow/providers/amazon/get_provider_info.py | 5 +-
.../system/amazon/aws/example_athena_spark.py | 25 ++-
.../unit/amazon/aws/sensors/test_athena_spark.py | 175 +++++++++++++++++++++
7 files changed, 318 insertions(+), 4 deletions(-)
diff --git a/providers/amazon/docs/operators/athena/athena_spark.rst
b/providers/amazon/docs/operators/athena/athena_spark.rst
index 83fe2427052..260675ba40b 100644
--- a/providers/amazon/docs/operators/athena/athena_spark.rst
+++ b/providers/amazon/docs/operators/athena/athena_spark.rst
@@ -15,8 +15,8 @@
specific language governing permissions and limitations
under the License.
-Athena Spark Operators
-======================
+Athena Spark Operators and Sensors
+==================================
`Amazon Athena <https://aws.amazon.com/athena/>`__ supports Apache Spark
calculations through session-based APIs.
This page documents the provider support for submitting and monitoring those
@@ -53,6 +53,27 @@ this operator, please see the `Sample Dag
<https://github.com/apache/airflow/blo
:start-after: [START howto_operator_athena_spark]
:end-before: [END howto_operator_athena_spark]
+
+Sensors
+-------
+
+Wait for an Athena Spark calculation
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Use
:class:`~airflow.providers.amazon.aws.sensors.athena_spark.AthenaSparkSensor`
+to wait for an existing Athena Spark calculation execution to reach a terminal
+state.
+
+The sensor expects a calculation execution ID that was already created, for
+example by another Airflow task that submits a calculation without waiting,
+or by a system outside Airflow.
+
+.. exampleinclude::
/../../amazon/tests/system/amazon/aws/example_athena_spark.py
+ :language: python
+ :dedent: 4
+ :start-after: [START howto_sensor_athena_spark]
+ :end-before: [END howto_sensor_athena_spark]
+
Reference
---------
diff --git a/providers/amazon/provider.yaml b/providers/amazon/provider.yaml
index 159724fea0f..a2536a08ee1 100644
--- a/providers/amazon/provider.yaml
+++ b/providers/amazon/provider.yaml
@@ -566,6 +566,7 @@ sensors:
- integration-name: Amazon Athena
python-modules:
- airflow.providers.amazon.aws.sensors.athena
+ - airflow.providers.amazon.aws.sensors.athena_spark
- integration-name: Amazon Web Services
python-modules:
- airflow.providers.amazon.aws.sensors.base_aws
diff --git a/providers/amazon/src/airflow/providers/amazon/aws/hooks/athena.py
b/providers/amazon/src/airflow/providers/amazon/aws/hooks/athena.py
index f3bfbf6b6a9..7d8ab21fa82 100644
--- a/providers/amazon/src/airflow/providers/amazon/aws/hooks/athena.py
+++ b/providers/amazon/src/airflow/providers/amazon/aws/hooks/athena.py
@@ -87,6 +87,13 @@ class AthenaHook(AwsBaseHook):
"CANCELLED",
)
+ SPARK_INTERMEDIATE_STATES = (
+ "CREATING",
+ "CREATED",
+ "QUEUED",
+ "RUNNING",
+ "CANCELING",
+ )
SPARK_FAILURE_STATES = (
"FAILED",
"CANCELED",
diff --git
a/providers/amazon/src/airflow/providers/amazon/aws/sensors/athena_spark.py
b/providers/amazon/src/airflow/providers/amazon/aws/sensors/athena_spark.py
new file mode 100644
index 00000000000..81a6d4be99a
--- /dev/null
+++ b/providers/amazon/src/airflow/providers/amazon/aws/sensors/athena_spark.py
@@ -0,0 +1,84 @@
+# 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 collections.abc import Sequence
+from typing import TYPE_CHECKING
+
+from airflow.providers.amazon.aws.hooks.athena import AthenaHook
+from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
+from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
+
+if TYPE_CHECKING:
+ from airflow.sdk import Context
+
+
+class AthenaSparkSensor(AwsBaseSensor[AthenaHook]):
+ """
+ Waits for an Amazon Athena Spark calculation to reach a terminal state.
+
+ :param calculation_execution_id: Athena Spark calculation execution ID.
(templated)
+ :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). (templated)
+ :param region_name: AWS region_name. If not specified then the default
boto3 behaviour is used. (templated)
+ """
+
+ aws_hook_class = AthenaHook
+
+ template_fields: Sequence[str] =
aws_template_fields("calculation_execution_id")
+
+ def __init__(
+ self,
+ *,
+ calculation_execution_id: str,
+ aws_conn_id: str = "aws_default",
+ region_name: str | None = None,
+ **kwargs,
+ ) -> None:
+ super().__init__(aws_conn_id=aws_conn_id, region_name=region_name,
**kwargs)
+ self.calculation_execution_id = calculation_execution_id
+
+ def poke(self, context: Context) -> bool:
+ state = self.hook.check_spark_calculation_status(
+ calculation_execution_id=self.calculation_execution_id,
+ )
+
+ self.log.info("Calculation %s state is: %s",
self.calculation_execution_id, state)
+
+ if state in AthenaHook.SPARK_SUCCESS_STATES:
+ self.log.info("Athena Spark calculation %s completed.",
self.calculation_execution_id)
+ return True
+
+ if state in AthenaHook.SPARK_FAILURE_STATES:
+ reason = self.hook.get_spark_calculation_state_change_reason(
+ calculation_execution_id=self.calculation_execution_id,
+ )
+ raise RuntimeError(
+ f"Athena Spark calculation {self.calculation_execution_id}
failed with state {state}. "
+ f"Reason: {reason or 'Unknown'}"
+ )
+
+ if state in AthenaHook.SPARK_INTERMEDIATE_STATES:
+ return False
+
+ raise RuntimeError(
+ f"Unexpected Athena Spark calculation state for
{self.calculation_execution_id}: {state!r}"
+ )
diff --git a/providers/amazon/src/airflow/providers/amazon/get_provider_info.py
b/providers/amazon/src/airflow/providers/amazon/get_provider_info.py
index 45b899ed65e..0e816d7c0da 100644
--- a/providers/amazon/src/airflow/providers/amazon/get_provider_info.py
+++ b/providers/amazon/src/airflow/providers/amazon/get_provider_info.py
@@ -573,7 +573,10 @@ def get_provider_info():
"sensors": [
{
"integration-name": "Amazon Athena",
- "python-modules":
["airflow.providers.amazon.aws.sensors.athena"],
+ "python-modules": [
+ "airflow.providers.amazon.aws.sensors.athena",
+ "airflow.providers.amazon.aws.sensors.athena_spark",
+ ],
},
{
"integration-name": "Amazon Web Services",
diff --git a/providers/amazon/tests/system/amazon/aws/example_athena_spark.py
b/providers/amazon/tests/system/amazon/aws/example_athena_spark.py
index a1e08050d0a..f479b45347a 100644
--- a/providers/amazon/tests/system/amazon/aws/example_athena_spark.py
+++ b/providers/amazon/tests/system/amazon/aws/example_athena_spark.py
@@ -14,7 +14,6 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-
from __future__ import annotations
from datetime import datetime
@@ -24,6 +23,7 @@ import boto3
from airflow.providers.amazon.aws.hooks.athena import AthenaHook
from airflow.providers.amazon.aws.operators.athena_spark import
AthenaSparkOperator
from airflow.providers.amazon.aws.operators.s3 import S3CreateBucketOperator,
S3DeleteBucketOperator
+from airflow.providers.amazon.aws.sensors.athena_spark import AthenaSparkSensor
from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS
@@ -96,6 +96,16 @@ def stop_athena_spark_session(session_id: str) -> None:
client.terminate_session(SessionId=session_id)
+@task
+def start_athena_spark_calculation(session_id: str) -> str:
+ client = boto3.client("athena")
+ response = client.start_calculation_execution(
+ SessionId=session_id,
+ CodeBlock="print('hello from the athena spark sensor test')",
+ )
+ return response["CalculationExecutionId"]
+
+
with DAG(
dag_id=DAG_ID,
schedule="@once",
@@ -125,6 +135,17 @@ with DAG(
)
# [END howto_operator_athena_spark]
+ calculation_execution_id = start_athena_spark_calculation(idle_session_id)
+
+ # [START howto_sensor_athena_spark]
+ await_spark_calculation = AthenaSparkSensor(
+ task_id="await_spark_calculation",
+ calculation_execution_id=calculation_execution_id,
+ poke_interval=30,
+ timeout=3600,
+ )
+ # [END howto_sensor_athena_spark]
+
stop_session = stop_athena_spark_session(session_id)
delete_bucket = S3DeleteBucketOperator(
@@ -143,6 +164,8 @@ with DAG(
idle_session_id,
# TEST BODY
run_spark_calculation,
+ calculation_execution_id,
+ await_spark_calculation,
# TEST TEARDOWN
stop_session,
delete_work_group(work_group),
diff --git
a/providers/amazon/tests/unit/amazon/aws/sensors/test_athena_spark.py
b/providers/amazon/tests/unit/amazon/aws/sensors/test_athena_spark.py
new file mode 100644
index 00000000000..a428255980d
--- /dev/null
+++ b/providers/amazon/tests/unit/amazon/aws/sensors/test_athena_spark.py
@@ -0,0 +1,175 @@
+# 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 unittest import mock
+
+import pytest
+from moto import mock_aws
+
+from airflow.providers.amazon.aws.hooks.athena import AthenaHook
+from airflow.providers.amazon.aws.sensors.athena_spark import AthenaSparkSensor
+
+
[email protected]
+def mock_check_spark_calculation_status():
+ with mock.patch.object(
+ AthenaHook,
+ "check_spark_calculation_status",
+ autospec=True,
+ ) as m:
+ yield m
+
+
[email protected]
+def mock_get_spark_calculation_state_change_reason():
+ with mock.patch.object(
+ AthenaHook,
+ "get_spark_calculation_state_change_reason",
+ autospec=True,
+ ) as m:
+ yield m
+
+
+@mock_aws
+class TestAthenaSparkSensor:
+ def setup_method(self, _):
+ self.default_op_kwargs = {
+ "task_id": "test_athena_spark_sensor",
+ "calculation_execution_id": "abc",
+ }
+ self.sensor = AthenaSparkSensor(**self.default_op_kwargs,
aws_conn_id=None)
+
+ def test_base_aws_op_attributes(self):
+ op = AthenaSparkSensor(**self.default_op_kwargs)
+ assert op.hook.aws_conn_id == "aws_default"
+ assert op.hook._region_name is None
+ assert op.hook._verify is None
+ assert op.hook._config is None
+ assert op.hook.log_query is True
+
+ op = AthenaSparkSensor(
+ **self.default_op_kwargs,
+ aws_conn_id="aws-test-custom-conn",
+ region_name="eu-west-1",
+ verify=False,
+ botocore_config={"read_timeout": 42},
+ )
+ assert op.hook.aws_conn_id == "aws-test-custom-conn"
+ assert op.hook._region_name == "eu-west-1"
+ assert op.hook._verify is False
+ assert op.hook._config is not None
+ assert op.hook._config.read_timeout == 42
+
+ def test_poke_success_state(self, mock_check_spark_calculation_status):
+ mock_check_spark_calculation_status.return_value = "COMPLETED"
+
+ assert self.sensor.poke({}) is True
+
+ mock_check_spark_calculation_status.assert_called_once_with(
+ mock.ANY,
+ calculation_execution_id="abc",
+ )
+
+ @pytest.mark.parametrize(
+ "state",
+ ["CREATING", "CREATED", "QUEUED", "RUNNING", "CANCELING"],
+ )
+ def test_poke_intermediate_states(
+ self,
+ state,
+ mock_check_spark_calculation_status,
+ ):
+ mock_check_spark_calculation_status.return_value = state
+
+ assert self.sensor.poke({}) is False
+
+ mock_check_spark_calculation_status.assert_called_once_with(
+ mock.ANY,
+ calculation_execution_id="abc",
+ )
+
+ @pytest.mark.parametrize(
+ ("state", "reason", "expected_reason"),
+ [
+ pytest.param(
+ "FAILED",
+ "Calculation failed",
+ "Calculation failed",
+ id="failed-with-reason",
+ ),
+ pytest.param(
+ "CANCELED",
+ None,
+ "Unknown",
+ id="canceled-without-reason",
+ ),
+ ],
+ )
+ def test_poke_failure_states(
+ self,
+ state,
+ reason,
+ expected_reason,
+ mock_check_spark_calculation_status,
+ mock_get_spark_calculation_state_change_reason,
+ ):
+ mock_check_spark_calculation_status.return_value = state
+ mock_get_spark_calculation_state_change_reason.return_value = reason
+
+ with pytest.raises(
+ RuntimeError,
+ match=rf"failed with state {state}\. Reason: {expected_reason}",
+ ):
+ self.sensor.poke({})
+
+ mock_check_spark_calculation_status.assert_called_once_with(
+ mock.ANY,
+ calculation_execution_id="abc",
+ )
+ mock_get_spark_calculation_state_change_reason.assert_called_once_with(
+ mock.ANY,
+ calculation_execution_id="abc",
+ )
+
+ @pytest.mark.parametrize(
+ "state",
+ [
+ pytest.param(None, id="missing-state"),
+ pytest.param("UNKNOWN", id="unexpected-state"),
+ ],
+ )
+ def test_poke_unexpected_states(
+ self,
+ state,
+ mock_check_spark_calculation_status,
+ mock_get_spark_calculation_state_change_reason,
+ ):
+ mock_check_spark_calculation_status.return_value = state
+
+ with pytest.raises(
+ RuntimeError,
+ match="Unexpected Athena Spark calculation state",
+ ):
+ self.sensor.poke({})
+
+ mock_check_spark_calculation_status.assert_called_once_with(
+ mock.ANY,
+ calculation_execution_id="abc",
+ )
+ mock_get_spark_calculation_state_change_reason.assert_not_called()
