From 18383709caa423bf2d250708770ef0c10969eda7 Mon Sep 17 00:00:00 2001 From: Andisha Date: Thu, 7 May 2026 18:57:39 -0400 Subject: [PATCH 1/3] Amazon: Add Athena Spark operator and sensor support --- airflow-core/newsfragments/00000.feature.rst | 1 + .../docs/operators/athena/athena_spark.rst | 73 ++++++ providers/amazon/provider.yaml | 6 + .../providers/amazon/aws/hooks/athena.py | 132 ++++++++++ .../amazon/aws/operators/athena_spark.py | 229 ++++++++++++++++++ .../amazon/aws/sensors/athena_spark.py | 64 +++++ .../unit/amazon/aws/hooks/test_athena.py | 161 ++++++++++++ .../amazon/aws/operators/test_athena_spark.py | 147 +++++++++++ .../amazon/aws/sensors/test_athena_spark.py | 60 +++++ 9 files changed, 873 insertions(+) create mode 100644 airflow-core/newsfragments/00000.feature.rst create mode 100644 providers/amazon/docs/operators/athena/athena_spark.rst create mode 100644 providers/amazon/src/airflow/providers/amazon/aws/operators/athena_spark.py create mode 100644 providers/amazon/src/airflow/providers/amazon/aws/sensors/athena_spark.py create mode 100644 providers/amazon/tests/unit/amazon/aws/operators/test_athena_spark.py create mode 100644 providers/amazon/tests/unit/amazon/aws/sensors/test_athena_spark.py diff --git a/airflow-core/newsfragments/00000.feature.rst b/airflow-core/newsfragments/00000.feature.rst new file mode 100644 index 0000000000000..0059e9312d8cc --- /dev/null +++ b/airflow-core/newsfragments/00000.feature.rst @@ -0,0 +1 @@ +Add Athena Spark operator and sensor support to the Amazon provider. diff --git a/providers/amazon/docs/operators/athena/athena_spark.rst b/providers/amazon/docs/operators/athena/athena_spark.rst new file mode 100644 index 0000000000000..a9abbf244a60e --- /dev/null +++ b/providers/amazon/docs/operators/athena/athena_spark.rst @@ -0,0 +1,73 @@ +.. 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. + +Athena Spark Operators +====================== + +Amazon Athena supports Apache Spark calculations through session-based APIs. +This page documents the provider support for submitting and monitoring those +calculations from Airflow. + +AthenaSparkOperator +------------------- + +Use :class:`~airflow.providers.amazon.aws.operators.athena_spark.AthenaSparkOperator` +to submit Spark code to an existing Athena session and wait until the calculation +reaches a terminal state. + +.. code-block:: python + + from airflow import DAG + from airflow.providers.amazon.aws.operators.athena_spark import AthenaSparkOperator + from datetime import datetime + + with DAG( + dag_id="example_athena_spark_operator", + start_date=datetime(2024, 1, 1), + schedule=None, + catchup=False, + ) as dag: + AthenaSparkOperator( + task_id="run_spark_code", + session_id="my-athena-session-id", + code_block="print('hello from athena spark')", + poll_interval=10, + max_polling_attempts=60, + ) + +AthenaSparkSensor +----------------- + +Use :class:`~airflow.providers.amazon.aws.sensors.athena_spark.AthenaSparkSensor` +to wait for an existing calculation execution ID. + +.. code-block:: python + + from airflow import DAG + from airflow.providers.amazon.aws.sensors.athena_spark import AthenaSparkSensor + from datetime import datetime + + with DAG( + dag_id="example_athena_spark_sensor", + start_date=datetime(2024, 1, 1), + schedule=None, + catchup=False, + ) as dag: + AthenaSparkSensor( + task_id="wait_for_spark_calculation", + calculation_execution_id="calc-exec-123", + ) diff --git a/providers/amazon/provider.yaml b/providers/amazon/provider.yaml index df73ba8c7cadd..1a621675201fc 100644 --- a/providers/amazon/provider.yaml +++ b/providers/amazon/provider.yaml @@ -131,7 +131,9 @@ integrations: how-to-guide: - /docs/apache-airflow-providers-amazon/operators/athena/athena_boto.rst - /docs/apache-airflow-providers-amazon/operators/athena/athena_sql.rst + - /docs/apache-airflow-providers-amazon/operators/athena/athena_spark.rst tags: [aws] + - integration-name: Amazon Bedrock external-doc-url: https://aws.amazon.com/bedrock/ logo: /docs/integration-logos/Amazon-Bedrock_light-bg@4x.png @@ -413,6 +415,8 @@ operators: - integration-name: Amazon Athena python-modules: - airflow.providers.amazon.aws.operators.athena + - airflow.providers.amazon.aws.operators.athena_spark + - integration-name: Amazon Web Services python-modules: - airflow.providers.amazon.aws.operators.base_aws @@ -527,6 +531,8 @@ 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 5c9621cf69848..6afe5aa44d12e 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/hooks/athena.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/hooks/athena.py @@ -344,3 +344,135 @@ def stop_query(self, query_execution_id: str) -> dict: """ self.log.info("Stopping Query with executionId - %s", query_execution_id) return self.get_conn().stop_query_execution(QueryExecutionId=query_execution_id) + + # --- Athena Spark (Calculations) API --- + + SPARK_INTERMEDIATE_STATES = ( + "CREATING", + "CREATED", + "QUEUED", + "RUNNING", + "CANCELING", + ) + SPARK_FAILURE_STATES = ( + "FAILED", + "CANCELED", + ) + SPARK_SUCCESS_STATES = ("COMPLETED",) + SPARK_TERMINAL_STATES = ( + "COMPLETED", + "FAILED", + "CANCELED", + ) + + def start_calculation( + self, + *, + session_id: str, + code_block: str, + description: str | None = None, + calculation_configuration: dict[str, Any] | None = None, + client_request_token: str | None = None, + ) -> str: + """ + Start an Athena Spark calculation execution. + + .. seealso:: + - :external+boto3:py:meth:`Athena.Client.start_calculation_execution` + + :param session_id: The Athena session ID. + :param code_block: Spark code to execute (typically notebook-like code). + :param description: Optional description of the calculation. + :param calculation_configuration: Contains configuration information for the calculation. + :param client_request_token: Optional idempotency token. + :return: CalculationExecutionId + """ + params: dict[str, Any] = { + "SessionId": session_id, + "CodeBlock": code_block, + } + if description: + params["Description"] = description + + if calculation_configuration: + params["CalculationConfiguration"] = calculation_configuration + + if client_request_token: + params["ClientRequestToken"] = client_request_token + + if self.log_query: + self.log.info("Starting CalculationExecution with params:\n%s", query_params_to_string(params)) + response = self.get_conn().start_calculation_execution(**params) + calc_execution_id = response["CalculationExecutionId"] + self.log.info("Calculation execution id: %s", calc_execution_id) + return calc_execution_id + + def get_calculation_info(self, calculation_execution_id: str, use_cache: bool = False) -> dict[str, Any]: + """ + Get information about a single execution of a calculation. + + .. seealso:: + - :external+boto3:py:meth:`Athena.Client.get_calculation_execution` + + :param calculation_execution_id: CalculationExecutionId returned by start_calculation + :param use_cache: If True, use execution information cache + """ + cache_key = f"calc:{calculation_execution_id}" + if use_cache and cache_key in self.__query_results: + return self.__query_results[cache_key] + + response = self.get_conn().get_calculation_execution(CalculationExecutionId=calculation_execution_id) + + if use_cache: + self.__query_results[cache_key] = response + return response + + def check_calculation_status(self, calculation_execution_id: str, use_cache: bool = False) -> str | None: + """ + Fetch the state of a submitted calculation execution. + + .. seealso:: + - :external+boto3:py:meth:`Athena.Client.get_calculation_execution` + + :param calculation_execution_id: CalculationExecutionId returned by start_calculation + :return: One of valid calculation states, or *None* if the response is malformed. + """ + response = self.get_calculation_info( + calculation_execution_id=calculation_execution_id, use_cache=use_cache + ) + + try: + return response["Status"]["State"] + except KeyError: + self.log.error("Could not parse status for calculation %s", calculation_execution_id) + return None + + def get_calculation_state_change_reason( + self, calculation_execution_id: str, use_cache: bool = False + ) -> str | None: + """ + Fetch the reason for a calculation state change (e.g. error message). + + :param calculation_execution_id: CalculationExecutionId returned by start_calculation + :param use_cache: If True, use execution information cache + :return: State change reason string, or None. + """ + response = self.get_calculation_info( + calculation_execution_id=calculation_execution_id, use_cache=use_cache + ) + try: + return response["CalculationExecution"]["Status"].get("StateChangeReason") + except (KeyError, TypeError): + return None + + def stop_calculation(self, calculation_execution_id: str) -> dict[str, Any]: + """ + Cancel the submitted calculation execution. + + .. seealso:: + - :external+boto3:py:meth:`Athena.Client.stop_calculation_execution` + + :param calculation_execution_id: CalculationExecutionId returned by start_calculation + """ + self.log.info("Stopping CalculationExecution with id - %s", calculation_execution_id) + return self.get_conn().stop_calculation_execution(CalculationExecutionId=calculation_execution_id) diff --git a/providers/amazon/src/airflow/providers/amazon/aws/operators/athena_spark.py b/providers/amazon/src/airflow/providers/amazon/aws/operators/athena_spark.py new file mode 100644 index 0000000000000..8b98a816384a0 --- /dev/null +++ b/providers/amazon/src/airflow/providers/amazon/aws/operators/athena_spark.py @@ -0,0 +1,229 @@ +# +# 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. +""" +AthenaSparkOperator for running Apache Spark calculations in Amazon Athena. +""" + +from __future__ import annotations + +import time +from collections.abc import Sequence +from typing import TYPE_CHECKING, Any + +from airflow.providers.amazon.aws.hooks.athena import AthenaHook +from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator +from airflow.providers.amazon.aws.utils.mixins import aws_template_fields +from airflow.providers.common.compat.sdk import AirflowException + +if TYPE_CHECKING: + from airflow.sdk import Context + + +class AthenaSparkOperator(AwsBaseOperator[AthenaHook]): + """ + Run an Apache Spark calculation in an Amazon Athena session. + + Submits a calculation (e.g. PySpark code) via the Athena API, polls until + the calculation reaches a terminal state (COMPLETED, FAILED, or CANCELED), + and returns execution metadata. + + .. seealso:: + - :class:`airflow.providers.amazon.aws.hooks.athena.AthenaHook` + - `Athena for Apache Spark + `__ + + :param session_id: The Athena session ID in which to run the calculation. (templated) + :param code_block: The calculation code (e.g. PySpark) to execute. (templated) + :param description: Optional description of the calculation. + :param client_request_token: Optional idempotency token for the submission. + :param poll_interval: Seconds to wait between status checks. Default 30. + :param max_polling_attempts: Maximum number of polling attempts before timing out. + To limit total task time, use execution_timeout on the task as well. + :param log_query: Whether to log submission details. Default True. + :param aws_conn_id: The Airflow connection used for AWS credentials. + :param region_name: AWS region. If not set, default boto3 behaviour is used. + :param verify: Whether to verify SSL certificates. + :param botocore_config: Optional botocore configuration dict. + """ + + aws_hook_class = AthenaHook + ui_color = "#44b5e2" + template_fields: Sequence[str] = aws_template_fields("session_id", "code_block", "description") + template_ext: Sequence[str] = (".py",) + template_fields_renderers = {"code_block": "python"} + + def __init__( + self, + *, + session_id: str, + code_block: str, + description: str | None = None, + client_request_token: str | None = None, + poll_interval: int = 30, + max_polling_attempts: int = 120, + log_query: bool = True, + aws_conn_id: str | None = "aws_default", + region_name: str | None = None, + verify: bool | str | None = None, + botocore_config: dict | None = None, + **kwargs: Any, + ) -> None: + super().__init__( + aws_conn_id=aws_conn_id, + region_name=region_name, + verify=verify, + botocore_config=botocore_config, + **kwargs, + ) + self.session_id = session_id + self.code_block = code_block + self.description = description + self.client_request_token = client_request_token + self.poll_interval = poll_interval + self.max_polling_attempts = max_polling_attempts + self.log_query = log_query + self._calculation_execution_id: str | None = None + + @property + def _hook_parameters(self) -> dict[str, Any]: + return {**super()._hook_parameters, "log_query": self.log_query} + + def execute(self, context: Context) -> dict[str, Any]: + """Submit the Spark calculation, poll until terminal state, then return metadata.""" + del context + self.log.info("Starting Athena Spark calculation in session %s", self.session_id) + + calculation_execution_id = self.hook.start_calculation( + session_id=self.session_id, + code_block=self.code_block, + description=self.description, + client_request_token=self.client_request_token, + ) + self._calculation_execution_id = calculation_execution_id + initial_state = self.hook.check_calculation_status(calculation_execution_id) + self.log.info( + "Calculation submitted. CalculationExecutionId: %s, initial state: %s", + calculation_execution_id, + initial_state, + ) + + if initial_state and initial_state in AthenaHook.SPARK_TERMINAL_STATES: + return self._handle_terminal_state(calculation_execution_id, initial_state) + + final_state = self._poll_until_terminal(calculation_execution_id) + return self._handle_terminal_state(calculation_execution_id, final_state) + + def _poll_until_terminal(self, calculation_execution_id: str) -> str: + """Poll calculation status until a terminal state or timeout.""" + for attempt in range(1, self.max_polling_attempts + 1): + if attempt > 1: + time.sleep(self.poll_interval) + state = self.hook.check_calculation_status(calculation_execution_id) + + if state is None: + raise AirflowException( + f"Malformed or missing status for calculation {calculation_execution_id}. " + "Cannot continue polling." + ) + + self.log.info( + "CalculationExecutionId: %s, current state: %s (attempt %d/%d)", + calculation_execution_id, + state, + attempt, + self.max_polling_attempts, + ) + + if state in AthenaHook.SPARK_TERMINAL_STATES: + return state + + raise AirflowException( + f"Polling timed out after {self.max_polling_attempts} attempts for calculation " + f"{calculation_execution_id}. Use execution_timeout or increase max_polling_attempts." + ) + + def _handle_terminal_state(self, calculation_execution_id: str, state: str) -> dict[str, Any]: + """Resolve terminal state: raise on failure/cancel, build and return metadata.""" + reason = self.hook.get_calculation_state_change_reason(calculation_execution_id) + execution_info = self.hook.get_calculation_info(calculation_execution_id) + status = ( + execution_info.get("Status") + or (execution_info.get("CalculationExecution") or {}).get("Status") + or {} + ) + submission_time = status.get("SubmissionDateTime") + completion_time = status.get("CompletionDateTime") + workgroup = ( + execution_info.get("WorkGroup") + or (execution_info.get("CalculationExecution") or {}).get("WorkGroup") + or (execution_info.get("CalculationExecution") or {}).get("Workgroup") + ) + output_location = ( + execution_info.get("OutputLocation") + or (execution_info.get("CalculationExecution") or {}).get("OutputLocation") + or (execution_info.get("ResultConfiguration") or {}).get("OutputLocation") + ) + + result = { + "calculation_execution_id": calculation_execution_id, + "state": state, + "state_change_reason": reason, + "submission_time": str(submission_time) if submission_time else None, + "completion_time": str(completion_time) if completion_time else None, + "session_id": self.session_id, + "workgroup": workgroup, + "output_location": output_location, + } + + if state in AthenaHook.SPARK_FAILURE_STATES: + self.log.error( + "Calculation failed. CalculationExecutionId: %s, state: %s, reason: %s", + calculation_execution_id, + state, + reason, + ) + raise AirflowException( + f"Athena Spark calculation ended in {state}. " + f"CalculationExecutionId: {calculation_execution_id}. " + f"Reason: {reason or 'No reason provided.'}" + ) + + if state != "COMPLETED": + raise AirflowException( + f"Unexpected terminal state: {state} for calculation {calculation_execution_id}. " + "Expected COMPLETED, FAILED, or CANCELED." + ) + + self.log.info( + "Calculation completed successfully. CalculationExecutionId: %s", + calculation_execution_id, + ) + return result + + def on_kill(self) -> None: + """Request cancellation of the calculation when the task is killed.""" + if self._calculation_execution_id: + self.log.info("Received kill signal; stopping calculation %s", self._calculation_execution_id) + try: + self.hook.stop_calculation(self._calculation_execution_id) + except Exception as e: + self.log.warning( + "Failed to stop calculation %s: %s", + self._calculation_execution_id, + e, + ) 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 0000000000000..0d950cf0ce7bc --- /dev/null +++ b/providers/amazon/src/airflow/providers/amazon/aws/sensors/athena_spark.py @@ -0,0 +1,64 @@ +# +# 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, Any + +from airflow.exceptions import AirflowException +from airflow.providers.amazon.aws.hooks.athena import AthenaHook +from airflow.sensors.base import BaseSensorOperator + +if TYPE_CHECKING: + from airflow.utils.context import Context + + +class AthenaSparkSensor(BaseSensorOperator): + """ + Poll the status of an AWS Athena Spark calculation until it reaches a terminal state. + + :param calculation_execution_id: The ID of the calculation to monitor. (templated) + :param aws_conn_id: The Airflow connection used for AWS credentials. + """ + + template_fields: Sequence[str] = ("calculation_execution_id",) + ui_color = "#44e2b5" + + def __init__( + self, + *, + calculation_execution_id: str, + aws_conn_id: str = "aws_default", + **kwargs: Any, + ) -> None: + super().__init__(**kwargs) + self.calculation_execution_id = calculation_execution_id + self.aws_conn_id = aws_conn_id + + def poke(self, context: Context) -> bool: + """Check the current status of the Spark calculation.""" + del context + hook = AthenaHook(aws_conn_id=self.aws_conn_id) + state = hook.check_calculation_status(self.calculation_execution_id) + + self.log.info("Calculation %s state is: %s", self.calculation_execution_id, state) + + if state in hook.SPARK_FAILURE_STATES: + raise AirflowException(f"Calculation {self.calculation_execution_id} failed with state: {state}") + + return state == "COMPLETED" diff --git a/providers/amazon/tests/unit/amazon/aws/hooks/test_athena.py b/providers/amazon/tests/unit/amazon/aws/hooks/test_athena.py index e743e831873e5..5ecf4ba3e5f95 100644 --- a/providers/amazon/tests/unit/amazon/aws/hooks/test_athena.py +++ b/providers/amazon/tests/unit/amazon/aws/hooks/test_athena.py @@ -14,11 +14,21 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +""" +Unit tests for AthenaHook with boto3 client mocked. + +Test strategy: +- Mock the boto3 client via AthenaHook.get_conn() so no real AWS calls are made. +- Cover success, failure (exceptions), and bad/edge-case input for each hook method. +- Use botocore.exceptions.ClientError for API failure scenarios. +""" + from __future__ import annotations from unittest import mock import pytest +from botocore.exceptions import ClientError from moto import mock_aws from airflow.providers.amazon.aws.hooks.athena import ( @@ -40,6 +50,10 @@ "query_execution_id": "eac427d0-1c6d-4dfb-96aa-2835d3ac6595", "next_token_id": "eac427d0-1c6d-4dfb-96aa-2835d3ac6595", "max_items": 1000, + "code_block": "print('hello spark')", + "calculation_execution_id": "calc-123456", + "session_id": "session-123456", + "description": "spark-calc", } mock_query_context = {"Database": MOCK_DATA["database"]} @@ -57,6 +71,10 @@ "Status": {"StateChangeReason": "Terminated by user."}, } } +MOCK_CALCULATION_EXECUTION = {"CalculationExecutionId": MOCK_DATA["calculation_execution_id"]} + +MOCK_RUNNING_CALC_EXECUTION = {"Status": {"State": "RUNNING"}} +MOCK_SUCCEEDED_CALC_EXECUTION = {"Status": {"State": "COMPLETED"}} @mock_aws @@ -103,6 +121,21 @@ def test_hook_run_query_with_token(self, mock_conn): mock_conn.return_value.start_query_execution.assert_called_with(**expected_call_params) assert result == MOCK_DATA["query_execution_id"] + @mock.patch.object(AthenaHook, "get_conn") + def test_hook_run_query_boto3_failure(self, mock_conn): + """Failure case: boto3 start_query_execution raises ClientError.""" + mock_conn.return_value.start_query_execution.side_effect = ClientError( + error_response={"Error": {"Code": "InvalidRequestException", "Message": "Invalid query"}}, + operation_name="start_query_execution", + ) + with pytest.raises(ClientError) as exc_info: + self.athena.run_query( + query=MOCK_DATA["query"], + query_context=mock_query_context, + result_configuration=mock_result_configuration, + ) + assert exc_info.value.response["Error"]["Code"] == "InvalidRequestException" + @mock.patch.object(AthenaHook, "log") @mock.patch.object(AthenaHook, "get_conn") def test_hook_run_query_log_query(self, mock_conn, log): @@ -139,6 +172,78 @@ def test_run_query_hook_lineage(self, mock_conn, mock_send_lineage): assert call_kw["sql"] == MOCK_DATA["query"] assert call_kw["job_id"] == MOCK_DATA["query_execution_id"] + # new test cases + @mock.patch.object(AthenaHook, "get_conn") + def test_hook_start_calculation_default_params(self, mock_conn): + mock_conn.return_value.start_calculation_execution.return_value = MOCK_CALCULATION_EXECUTION + + result = self.athena.start_calculation( + session_id=MOCK_DATA["session_id"], + code_block=MOCK_DATA["code_block"], + ) + + expected_call_params = { + "SessionId": MOCK_DATA["session_id"], + "CodeBlock": MOCK_DATA["code_block"], + } + mock_conn.return_value.start_calculation_execution.assert_called_with(**expected_call_params) + assert result == MOCK_DATA["calculation_execution_id"] + + @mock.patch.object(AthenaHook, "get_conn") + def test_hook_start_calculation_with_optional_params(self, mock_conn): + mock_conn.return_value.start_calculation_execution.return_value = MOCK_CALCULATION_EXECUTION + + calculation_configuration = {"CodeBlock": MOCK_DATA["code_block"]} + + result = self.athena.start_calculation( + session_id=MOCK_DATA["session_id"], + code_block=MOCK_DATA["code_block"], + description=MOCK_DATA["description"], + calculation_configuration=calculation_configuration, + client_request_token=MOCK_DATA["client_request_token"], + ) + + expected_call_params = { + "SessionId": MOCK_DATA["session_id"], + "CodeBlock": MOCK_DATA["code_block"], + "Description": MOCK_DATA["description"], + "CalculationConfiguration": calculation_configuration, + "ClientRequestToken": MOCK_DATA["client_request_token"], + } + mock_conn.return_value.start_calculation_execution.assert_called_with(**expected_call_params) + assert result == MOCK_DATA["calculation_execution_id"] + + @mock.patch.object(AthenaHook, "get_conn") + def test_hook_get_calculation_info(self, mock_conn): + mock_conn.return_value.get_calculation_execution.return_value = MOCK_SUCCEEDED_CALC_EXECUTION + + result = self.athena.get_calculation_info( + calculation_execution_id=MOCK_DATA["calculation_execution_id"] + ) + + mock_conn.return_value.get_calculation_execution.assert_called_once_with( + CalculationExecutionId=MOCK_DATA["calculation_execution_id"] + ) + assert result == MOCK_SUCCEEDED_CALC_EXECUTION + + @mock.patch.object(AthenaHook, "get_conn") + def test_check_calculation_status_normal(self, mock_conn): + mock_conn.return_value.get_calculation_execution.return_value = MOCK_RUNNING_CALC_EXECUTION + + state = self.athena.check_calculation_status( + calculation_execution_id=MOCK_DATA["calculation_execution_id"] + ) + + assert state == "RUNNING" + + @mock.patch.object(AthenaHook, "get_conn") + def test_hook_stop_calculation(self, mock_conn): + self.athena.stop_calculation(calculation_execution_id=MOCK_DATA["calculation_execution_id"]) + + mock_conn.return_value.stop_calculation_execution.assert_called_once_with( + CalculationExecutionId=MOCK_DATA["calculation_execution_id"] + ) + @mock.patch.object(AthenaHook, "get_conn") def test_hook_get_query_results_with_non_succeeded_query(self, mock_conn): mock_conn.return_value.get_query_execution.return_value = MOCK_RUNNING_QUERY_EXECUTION @@ -165,6 +270,24 @@ def test_hook_get_query_results_with_next_token(self, mock_conn): } mock_conn.return_value.get_query_results.assert_called_with(**expected_call_params) + @mock.patch.object(AthenaHook, "get_conn") + def test_hook_get_query_results_invalid_state_returns_none(self, mock_conn): + """Edge case: get_query_execution returns malformed response; check_query_status returns None.""" + mock_conn.return_value.get_query_execution.return_value = {"QueryExecution": {}} + result = self.athena.get_query_results(query_execution_id=MOCK_DATA["query_execution_id"]) + assert result is None + mock_conn.return_value.get_query_results.assert_not_called() + + @mock.patch.object(AthenaHook, "get_conn") + def test_hook_get_query_info_boto3_failure(self, mock_conn): + """Failure case: boto3 get_query_execution raises ClientError.""" + mock_conn.return_value.get_query_execution.side_effect = ClientError( + error_response={"Error": {"Code": "InvalidRequestException"}}, + operation_name="get_query_execution", + ) + with pytest.raises(ClientError): + self.athena.get_query_info(query_execution_id=MOCK_DATA["query_execution_id"]) + @mock.patch.object(AthenaHook, "get_conn") def test_hook_get_paginator_with_non_succeeded_query(self, mock_conn): mock_conn.return_value.get_query_execution.return_value = MOCK_RUNNING_QUERY_EXECUTION @@ -181,6 +304,14 @@ def test_hook_get_paginator_with_default_params(self, mock_conn): } mock_conn.return_value.get_paginator.return_value.paginate.assert_called_with(**expected_call_params) + @mock.patch.object(AthenaHook, "get_conn") + def test_hook_get_paginator_invalid_state_returns_none(self, mock_conn): + """Edge case: malformed response leads to None state; paginator not created.""" + mock_conn.return_value.get_query_execution.return_value = {"QueryExecution": {}} + result = self.athena.get_query_results_paginator(query_execution_id=MOCK_DATA["query_execution_id"]) + assert result is None + mock_conn.return_value.get_paginator.assert_not_called() + @mock.patch.object(AthenaHook, "get_conn") def test_hook_get_paginator_with_pagination_config(self, mock_conn): mock_conn.return_value.get_query_execution.return_value = MOCK_SUCCEEDED_QUERY_EXECUTION @@ -233,6 +364,29 @@ def test_hook_get_output_location(self, mock_conn): result = self.athena.get_output_location(query_execution_id=MOCK_DATA["query_execution_id"]) assert result == "s3://test_bucket/test.csv" + @mock.patch.object(AthenaHook, "get_conn") + def test_hook_stop_query_success(self, mock_conn): + """Success case: stop_query_execution returns normally.""" + mock_conn.return_value.stop_query_execution.return_value = {} + result = self.athena.stop_query(query_execution_id=MOCK_DATA["query_execution_id"]) + mock_conn.return_value.stop_query_execution.assert_called_once_with( + QueryExecutionId=MOCK_DATA["query_execution_id"] + ) + assert result == {} + + @mock.patch.object(AthenaHook, "get_conn") + def test_hook_stop_query_boto3_failure(self, mock_conn): + """Failure case: boto3 stop_query_execution raises ClientError.""" + mock_conn.return_value.stop_query_execution.side_effect = ClientError( + error_response={ + "Error": {"Code": "InvalidRequestException", "Message": "Query already finished"} + }, + operation_name="stop_query_execution", + ) + with pytest.raises(ClientError) as exc_info: + self.athena.stop_query(query_execution_id=MOCK_DATA["query_execution_id"]) + assert exc_info.value.response["Error"]["Code"] == "InvalidRequestException" + @pytest.mark.parametrize( "query_execution_id", [pytest.param("", id="empty-string"), pytest.param(None, id="none")] ) @@ -268,6 +422,13 @@ def test_check_query_status_exception(self, mock_get_query_info): state = self.athena.check_query_status(query_execution_id=MOCK_DATA["query_execution_id"]) assert not state + @mock.patch.object(AthenaHook, "get_query_info") + def test_get_state_change_reason_missing_key_returns_none(self, mock_get_query_info): + """Edge case: response has no StateChangeReason; hook returns None and logs.""" + mock_get_query_info.return_value = {"QueryExecution": {"Status": {}}} + result = self.athena.get_state_change_reason(query_execution_id=MOCK_DATA["query_execution_id"]) + assert result is None + @mock.patch.object(AthenaHook, "get_conn") def test_hook_get_query_info_caching(self, mock_conn): mock_conn.return_value.get_query_execution.return_value = MOCK_QUERY_EXECUTION_OUTPUT diff --git a/providers/amazon/tests/unit/amazon/aws/operators/test_athena_spark.py b/providers/amazon/tests/unit/amazon/aws/operators/test_athena_spark.py new file mode 100644 index 0000000000000..f868b80f5d052 --- /dev/null +++ b/providers/amazon/tests/unit/amazon/aws/operators/test_athena_spark.py @@ -0,0 +1,147 @@ +# +# 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 airflow.providers.amazon.aws.hooks.athena import AthenaHook +from airflow.providers.amazon.aws.operators.athena_spark import AthenaSparkOperator +from airflow.providers.common.compat.sdk import AirflowException + +CALC_ID = "calc-exec-123" +SESSION_ID = "session-456" +CODE_BLOCK = "1 + 1" + + +@pytest.fixture +def operator(): + return AthenaSparkOperator( + task_id="test_athena_spark", + session_id=SESSION_ID, + code_block=CODE_BLOCK, + poll_interval=0, + max_polling_attempts=5, + ) + + +@pytest.fixture +def context(): + return {"ti": None} + + +def _exec_info(state: str, submission_time=None, completion_time=None): + return { + "Status": { + "State": state, + "SubmissionDateTime": submission_time, + "CompletionDateTime": completion_time, + } + } + + +class TestAthenaSparkOperator: + def test_init(self, operator): + assert operator.session_id == SESSION_ID + assert operator.code_block == CODE_BLOCK + assert operator.poll_interval == 0 + assert operator.max_polling_attempts == 5 + assert operator._calculation_execution_id is None + + def test_template_fields(self): + assert "session_id" in AthenaSparkOperator.template_fields + assert "code_block" in AthenaSparkOperator.template_fields + assert "description" in AthenaSparkOperator.template_fields + + @mock.patch.object(AthenaHook, "get_calculation_info", return_value={}) + @mock.patch.object(AthenaHook, "get_calculation_state_change_reason", return_value=None) + @mock.patch.object(AthenaHook, "check_calculation_status", return_value="COMPLETED") + @mock.patch.object(AthenaHook, "start_calculation", return_value=CALC_ID) + def test_execute_success_immediate_completed( + self, mock_start, mock_check, mock_reason, mock_info, operator, context + ): + mock_info.return_value = _exec_info("COMPLETED") + result = operator.execute(context) + mock_start.assert_called_once_with( + session_id=SESSION_ID, + code_block=CODE_BLOCK, + description=None, + client_request_token=None, + ) + assert result["calculation_execution_id"] == CALC_ID + assert result["state"] == "COMPLETED" + assert result["session_id"] == SESSION_ID + + @mock.patch.object(AthenaHook, "get_calculation_info", return_value={}) + @mock.patch.object(AthenaHook, "get_calculation_state_change_reason", return_value="Job failed") + @mock.patch.object(AthenaHook, "check_calculation_status", return_value="FAILED") + @mock.patch.object(AthenaHook, "start_calculation", return_value=CALC_ID) + def test_execute_failure_raises(self, mock_start, mock_check, mock_reason, mock_info, operator, context): + mock_info.return_value = _exec_info("FAILED") + with pytest.raises(AirflowException, match="FAILED"): + operator.execute(context) + mock_reason.assert_called() + + @mock.patch.object(AthenaHook, "get_calculation_info", return_value={}) + @mock.patch.object(AthenaHook, "get_calculation_state_change_reason", return_value="Canceled") + @mock.patch.object(AthenaHook, "check_calculation_status", return_value="CANCELED") + @mock.patch.object(AthenaHook, "start_calculation", return_value=CALC_ID) + def test_execute_cancelled_raises( + self, mock_start, mock_check, mock_reason, mock_info, operator, context + ): + mock_info.return_value = _exec_info("CANCELED") + with pytest.raises(AirflowException, match="CANCELED"): + operator.execute(context) + + @mock.patch.object(AthenaHook, "get_calculation_info", return_value={}) + @mock.patch.object(AthenaHook, "get_calculation_state_change_reason", return_value=None) + @mock.patch.object(AthenaHook, "check_calculation_status", side_effect=["RUNNING", "COMPLETED"]) + @mock.patch.object(AthenaHook, "start_calculation", return_value=CALC_ID) + def test_execute_poll_then_success( + self, mock_start, mock_check, mock_reason, mock_info, operator, context + ): + mock_info.return_value = _exec_info("COMPLETED") + result = operator.execute(context) + assert mock_check.call_count == 2 + assert result["state"] == "COMPLETED" + + @mock.patch.object(AthenaHook, "check_calculation_status", return_value="RUNNING") + @mock.patch.object(AthenaHook, "start_calculation", return_value=CALC_ID) + def test_execute_poll_timeout(self, mock_start, mock_check, operator, context): + operator.max_polling_attempts = 2 + with pytest.raises(AirflowException, match="timed out"): + operator.execute(context) + assert mock_check.call_count >= 2 + + @mock.patch.object(AthenaHook, "check_calculation_status", return_value=None) + @mock.patch.object(AthenaHook, "start_calculation", return_value=CALC_ID) + def test_execute_malformed_status_raises(self, mock_start, mock_check, operator, context): + with pytest.raises(AirflowException, match="Malformed or missing status"): + operator.execute(context) + + @mock.patch.object(AthenaHook, "stop_calculation") + def test_on_kill_calls_stop_calculation(self, mock_stop, operator): + operator._calculation_execution_id = CALC_ID + operator.on_kill() + mock_stop.assert_called_once_with(CALC_ID) + + @mock.patch.object(AthenaHook, "stop_calculation") + def test_on_kill_no_op_when_no_calc_id(self, mock_stop, operator): + operator.on_kill() + mock_stop.assert_not_called() 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 0000000000000..f01fb554464ff --- /dev/null +++ b/providers/amazon/tests/unit/amazon/aws/sensors/test_athena_spark.py @@ -0,0 +1,60 @@ +# +# 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 airflow.exceptions import AirflowException +from airflow.providers.amazon.aws.hooks.athena import AthenaHook +from airflow.providers.amazon.aws.sensors.athena_spark import AthenaSparkSensor + +CALC_ID = "calc-exec-123" + + +@pytest.fixture +def sensor() -> AthenaSparkSensor: + return AthenaSparkSensor(task_id="test_athena_spark_sensor", calculation_execution_id=CALC_ID) + + +class TestAthenaSparkSensor: + def test_init(self, sensor: AthenaSparkSensor): + assert sensor.calculation_execution_id == CALC_ID + assert sensor.aws_conn_id == "aws_default" + + def test_template_fields(self): + assert AthenaSparkSensor.template_fields == ("calculation_execution_id",) + + @mock.patch.object(AthenaHook, "check_calculation_status", return_value="COMPLETED") + def test_poke_completed_returns_true(self, mock_check: mock.Mock, sensor: AthenaSparkSensor): + result = sensor.poke({}) + assert result is True + mock_check.assert_called_once_with(CALC_ID) + + @mock.patch.object(AthenaHook, "check_calculation_status", return_value="FAILED") + def test_poke_failed_raises(self, mock_check: mock.Mock, sensor: AthenaSparkSensor): + with pytest.raises(AirflowException, match="failed with state: FAILED"): + sensor.poke({}) + mock_check.assert_called_once_with(CALC_ID) + + @mock.patch.object(AthenaHook, "check_calculation_status", return_value="RUNNING") + def test_poke_running_returns_false(self, mock_check: mock.Mock, sensor: AthenaSparkSensor): + result = sensor.poke({}) + assert result is False + mock_check.assert_called_once_with(CALC_ID) From c38a3d666fd51dc378149c7bc32cdeffdd7cb958 Mon Sep 17 00:00:00 2001 From: Andisha Date: Thu, 7 May 2026 19:07:42 -0400 Subject: [PATCH 2/3] Rename Athena Spark newsfragment for PR number --- .../newsfragments/{00000.feature.rst => 66576.feature.rst} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename airflow-core/newsfragments/{00000.feature.rst => 66576.feature.rst} (100%) diff --git a/airflow-core/newsfragments/00000.feature.rst b/airflow-core/newsfragments/66576.feature.rst similarity index 100% rename from airflow-core/newsfragments/00000.feature.rst rename to airflow-core/newsfragments/66576.feature.rst From c34fb593be636dc6e1b4bb806d96548abbfd68d0 Mon Sep 17 00:00:00 2001 From: Danish Safdariyan Date: Fri, 15 May 2026 14:27:19 -0400 Subject: [PATCH 3/3] Remove newsfragment; provider changes do not require one. Co-authored-by: Cursor --- airflow-core/newsfragments/66576.feature.rst | 1 - 1 file changed, 1 deletion(-) delete mode 100644 airflow-core/newsfragments/66576.feature.rst diff --git a/airflow-core/newsfragments/66576.feature.rst b/airflow-core/newsfragments/66576.feature.rst deleted file mode 100644 index 0059e9312d8cc..0000000000000 --- a/airflow-core/newsfragments/66576.feature.rst +++ /dev/null @@ -1 +0,0 @@ -Add Athena Spark operator and sensor support to the Amazon provider.