You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "SameerMesiah97 (via GitHub)" <gi...@apache.org> on 2026/08/28 18:34:45 UTC
[PR] Add AthenaSparkSensor [airflow]
SameerMesiah97 opened a new pull request, #72229:
URL: https://github.com/apache/airflow/pull/72229
**Description**
This change introduces a new sensor, `AthenaSparkSensor`, to wait for an existing Amazon Athena Spark calculation execution to reach a terminal state.
The sensor delegates to `AthenaHook` Spark calculation helpers to poll the calculation status. It succeeds when the calculation reaches `COMPLETED`, continues poking while the calculation is in an in-progress state, and raises when the calculation reaches a failed, canceled, or unexpected state.
**Rationale**
Amazon Athena Spark calculations are identified by a `calculation_execution_id` after they are submitted. While `AthenaSparkOperator` can submit Spark code and wait for completion, users may also need to monitor calculations started outside the operator, for example by another task or by an external system.
Adding `AthenaSparkSensor` gives DAG authors a provider-native way to wait on an existing Athena Spark calculation without writing custom polling logic around boto3.
**Tests**
Added unit tests verifying that:
* the sensor inherits the expected AWS hook configuration behavior.
* successful calculation states are handled correctly.
* in-progress states such as `CREATING`, `CREATED`, `QUEUED`, and `RUNNING` continue waiting.
* failed and canceled calculation states raise an error.
* the state change reason is included when Athena provides one.
**Example DAG**
Updated the Athena Spark example DAG to demonstrate waiting for an existing Athena Spark calculation execution using `AthenaSparkSensor`.
**Documentation**
Added documentation for `AthenaSparkSensor`, including its expected `calculation_execution_id` input and an example usage snippet.
**Backwards Compatibility**
This is a new sensor and does not modify existing Athena sensor behavior. No breaking changes are introduced.
###### Was generative AI tooling used to co-author this PR?
- [X] Yes (please specify the tool below)
Generated-by: [GPT 5.5] following [the guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions)
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add AthenaSparkSensor [airflow]
Posted by "SameerMesiah97 (via GitHub)" <gi...@apache.org>.
SameerMesiah97 commented on code in PR #72229:
URL: https://github.com/apache/airflow/pull/72229#discussion_r4302913947
##########
providers/amazon/tests/unit/amazon/aws/sensors/test_athena_spark.py:
##########
@@ -0,0 +1,106 @@
+# 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
+
+
+@pytest.fixture
+def mock_check_spark_calculation_status():
+ with mock.patch.object(AthenaHook, "check_spark_calculation_status") as m:
Review Comment:
Done.
##########
providers/amazon/docs/operators/athena/athena_spark.rst:
##########
@@ -53,6 +53,28 @@ 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 calling
+:class:`~airflow.providers.amazon.aws.operators.athena_spark.AthenaSparkOperator`
+or by starting the calculation outside Airflow.
Review Comment:
Done.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add AthenaSparkSensor [airflow]
Posted by "SameerMesiah97 (via GitHub)" <gi...@apache.org>.
SameerMesiah97 commented on code in PR #72229:
URL: https://github.com/apache/airflow/pull/72229#discussion_r4302907764
##########
providers/amazon/tests/unit/amazon/aws/sensors/test_athena_spark.py:
##########
@@ -0,0 +1,106 @@
+# 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
+
+
+@pytest.fixture
+def mock_check_spark_calculation_status():
+ with mock.patch.object(AthenaHook, "check_spark_calculation_status") as m:
+ yield m
+
+
+@pytest.fixture
+def mock_get_spark_calculation_state_change_reason():
+ with mock.patch.object(AthenaHook, "get_spark_calculation_state_change_reason") as m:
+ yield m
+
+
+@mock_aws
+class TestAthenaSparkSensor:
+ def setup_method(self, _):
+ self.default_op_kwargs = dict(
+ 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_template_fields(self):
+ assert AthenaSparkSensor.template_fields == (
+ "calculation_execution_id",
+ "aws_conn_id",
+ "region_name",
+ )
+
+ @pytest.mark.parametrize("state", ["COMPLETED"])
+ def test_poke_success_states(self, state, mock_check_spark_calculation_status):
+ mock_check_spark_calculation_status.side_effect = [state]
+
+ assert self.sensor.poke({}) is True
+ mock_check_spark_calculation_status.assert_called_once_with(calculation_execution_id="abc")
+
+ @pytest.mark.parametrize("state", ["CREATING", "CREATED", "QUEUED", "RUNNING"])
Review Comment:
All three adjustments have been made to the tests.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add AthenaSparkSensor [airflow]
Posted by "SameerMesiah97 (via GitHub)" <gi...@apache.org>.
SameerMesiah97 commented on PR #72229:
URL: https://github.com/apache/airflow/pull/72229#issuecomment-5998826586
@potiuk
do you want this to remain a sensor or shall I substitute it with deferrable mode as @ashb recommended?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add AthenaSparkSensor [airflow]
Posted by "shahar1 (via GitHub)" <gi...@apache.org>.
shahar1 merged PR #72229:
URL: https://github.com/apache/airflow/pull/72229
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add AthenaSparkSensor [airflow]
Posted by "SameerMesiah97 (via GitHub)" <gi...@apache.org>.
SameerMesiah97 commented on code in PR #72229:
URL: https://github.com/apache/airflow/pull/72229#discussion_r4302905141
##########
providers/amazon/src/airflow/providers/amazon/aws/sensors/athena_spark.py:
##########
@@ -0,0 +1,88 @@
+# 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
+
+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. Defaults to ``aws_default``.
+ :param region_name: AWS region name. Defaults to ``None``. (templated)
+ """
+
+ aws_hook_class = AthenaHook
+
+ template_fields: Sequence[str] = (
+ "calculation_execution_id",
+ "aws_conn_id",
+ "region_name",
+ )
+
+ 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:
+ self.log.info(
+ "Athena Spark calculation %s is in state %s.",
+ self.calculation_execution_id,
+ state,
+ )
+ return False
Review Comment:
The logging has been removed from here.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add AthenaSparkSensor [airflow]
Posted by "SameerMesiah97 (via GitHub)" <gi...@apache.org>.
SameerMesiah97 commented on PR #72229:
URL: https://github.com/apache/airflow/pull/72229#issuecomment-5878762784
requesting review for this
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add AthenaSparkSensor [airflow]
Posted by "ashb (via GitHub)" <gi...@apache.org>.
ashb commented on PR #72229:
URL: https://github.com/apache/airflow/pull/72229#issuecomment-5968788472
I would generally say that a deferable mode on the operator is better than a sensor these days.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add AthenaSparkSensor [airflow]
Posted by "ashb (via GitHub)" <gi...@apache.org>.
ashb commented on PR #72229:
URL: https://github.com/apache/airflow/pull/72229#issuecomment-5968793044
Additionally using Task State to make the base operator Durable would be good too. This is a perfect candidate for 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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add AthenaSparkSensor [airflow]
Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk commented on code in PR #72229:
URL: https://github.com/apache/airflow/pull/72229#discussion_r4186189365
##########
providers/amazon/tests/system/amazon/aws/example_athena_spark.py:
##########
@@ -98,6 +108,17 @@ def stop_athena_spark_session(session_id: str) -> None:
)
# [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=run_spark_calculation.output,
Review Comment:
`AthenaSparkOperator.execute()` returns a metadata dict (`{"calculation_execution_id": ..., "state": ..., ...}`), not the ID string, so the sensor receives the whole dict. `get_calculation_execution(CalculationExecutionId=<dict>)` then raises boto3's `ParamValidationError` on the first poke, and the system test fails every run. Meanwhile the calculation started by `start_athena_spark_calculation` on line 111 is chained in but never awaited.
```suggestion
calculation_execution_id=calculation_execution_id,
```
This block is also the published `howto_sensor_athena_spark` snippet, so the docs show the broken wiring too. Relatedly, `athena_spark.rst` suggests getting the ID "for example by calling `AthenaSparkOperator`". That operator already waits for the calculation to finish, so pairing it with the sensor isn't useful; it may be clearer to say the ID comes from a calculation started without waiting.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add AthenaSparkSensor [airflow]
Posted by "shahar1 (via GitHub)" <gi...@apache.org>.
shahar1 commented on code in PR #72229:
URL: https://github.com/apache/airflow/pull/72229#discussion_r4297647131
##########
providers/amazon/src/airflow/providers/amazon/aws/sensors/athena_spark.py:
##########
@@ -0,0 +1,88 @@
+# 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
+
+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. Defaults to ``aws_default``.
+ :param region_name: AWS region name. Defaults to ``None``. (templated)
+ """
+
+ aws_hook_class = AthenaHook
+
+ template_fields: Sequence[str] = (
+ "calculation_execution_id",
+ "aws_conn_id",
+ "region_name",
+ )
+
+ 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
Review Comment:
**minor** — other AWS sensors build `template_fields` with `aws_template_fields(...)` (see `AthenaSensor` in `sensors/athena.py`). This tuple drops `verify`, so it is no longer templated here, unlike every sibling sensor.
The `aws_conn_id` / `region_name` parameters also re-declare what `AwsBaseSensor.__init__` already accepts, and narrow `aws_conn_id` to `str` while the base allows `None` (the unit tests construct the sensor with `aws_conn_id=None`). Letting the base handle them removes both issues:
```suggestion
template_fields: Sequence[str] = aws_template_fields("calculation_execution_id")
def __init__(self, *, calculation_execution_id: str, **kwargs) -> None:
super().__init__(**kwargs)
self.calculation_execution_id = calculation_execution_id
```
This needs `from airflow.providers.amazon.aws.utils.mixins import aws_template_fields`, the docstring's `aws_conn_id` / `region_name` entries can point at the base class like `AthenaSensor` does, and `test_template_fields` needs the new sorted tuple.
##########
providers/amazon/docs/operators/athena/athena_spark.rst:
##########
@@ -53,6 +53,28 @@ 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 calling
+:class:`~airflow.providers.amazon.aws.operators.athena_spark.AthenaSparkOperator`
+or by starting the calculation outside Airflow.
Review Comment:
**minor** — this still suggests getting the ID from `AthenaSparkOperator`. As noted in the earlier review round, that operator already waits for the calculation to finish, so pairing it with the sensor adds nothing. Pointing at a calculation started without waiting (another task, or a system outside Airflow, as the example Dag does) would be clearer.
##########
providers/amazon/src/airflow/providers/amazon/aws/sensors/athena_spark.py:
##########
@@ -0,0 +1,88 @@
+# 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
+
+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. Defaults to ``aws_default``.
+ :param region_name: AWS region name. Defaults to ``None``. (templated)
+ """
+
+ aws_hook_class = AthenaHook
+
+ template_fields: Sequence[str] = (
+ "calculation_execution_id",
+ "aws_conn_id",
+ "region_name",
+ )
+
+ 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:
+ self.log.info(
+ "Athena Spark calculation %s is in state %s.",
+ self.calculation_execution_id,
+ state,
+ )
+ return False
Review Comment:
**nit** — line 63 already logs the state on every poke, so this logs the same thing a second time for every in-progress poke.
```suggestion
if state in AthenaHook.SPARK_INTERMEDIATE_STATES:
return False
```
##########
providers/amazon/tests/unit/amazon/aws/sensors/test_athena_spark.py:
##########
@@ -0,0 +1,106 @@
+# 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
+
+
+@pytest.fixture
+def mock_check_spark_calculation_status():
+ with mock.patch.object(AthenaHook, "check_spark_calculation_status") as m:
Review Comment:
**minor** — these patches create unspec'd mocks. The review checklist says:
> Flag any `mock.Mock()` or `mock.MagicMock()` without `spec` or `autospec`. Unspec'd mocks silently accept any attribute access, hiding real bugs.
`mock.patch.object(AthenaHook, "check_spark_calculation_status", autospec=True)` (same for the reason patch) would catch a renamed or re-signatured hook method. With `autospec` the recorded call includes `self`, so the assertions become `assert_called_once_with(mock.ANY, calculation_execution_id="abc")`.
##########
providers/amazon/tests/unit/amazon/aws/sensors/test_athena_spark.py:
##########
@@ -0,0 +1,106 @@
+# 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
+
+
+@pytest.fixture
+def mock_check_spark_calculation_status():
+ with mock.patch.object(AthenaHook, "check_spark_calculation_status") as m:
+ yield m
+
+
+@pytest.fixture
+def mock_get_spark_calculation_state_change_reason():
+ with mock.patch.object(AthenaHook, "get_spark_calculation_state_change_reason") as m:
+ yield m
+
+
+@mock_aws
+class TestAthenaSparkSensor:
+ def setup_method(self, _):
+ self.default_op_kwargs = dict(
+ 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_template_fields(self):
+ assert AthenaSparkSensor.template_fields == (
+ "calculation_execution_id",
+ "aws_conn_id",
+ "region_name",
+ )
+
+ @pytest.mark.parametrize("state", ["COMPLETED"])
+ def test_poke_success_states(self, state, mock_check_spark_calculation_status):
+ mock_check_spark_calculation_status.side_effect = [state]
+
+ assert self.sensor.poke({}) is True
+ mock_check_spark_calculation_status.assert_called_once_with(calculation_execution_id="abc")
+
+ @pytest.mark.parametrize("state", ["CREATING", "CREATED", "QUEUED", "RUNNING"])
Review Comment:
**minor** — the review checklist asks for every changed behaviour to have a test that fails when the change is reverted:
> Flag any changed or added behaviour without a corresponding test, and flag tests that a reviewer cannot fail by reverting the PR's change.
Three paths in `poke` have no such test today:
- `CANCELING` is in `SPARK_INTERMEDIATE_STATES` but not in this parametrize list, so dropping it from the hook constant fails nothing.
- The final `raise RuntimeError("Unexpected ...")` branch is untested, including `state is None`, which `check_spark_calculation_status` returns on a malformed response.
- The failure test matches only `failed with state ...`. It never asserts that the state change reason reaches the message, or that a missing reason falls back to `Unknown`. The PR description lists the reason case as tested.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org