SameerMesiah97 commented on code in PR #73522:
URL: https://github.com/apache/airflow/pull/73522#discussion_r4087808360
##########
providers/influxdb/src/airflow/providers/influxdb/utils.py:
##########
@@ -28,3 +28,20 @@
def _convert_dataframe_to_records(dataframe: pd.DataFrame) -> list[dict[str,
Any]]:
"""Convert a query result DataFrame into a JSON-serializable list of
dictionaries."""
return json.loads(dataframe.to_json(orient="records", date_format="iso"))
+
+
+def _first_cell_is_truthy(dataframe: pd.DataFrame) -> bool:
+ """Return whether the first cell meets the sensor condition."""
+ import pandas as pd
Review Comment:
Is there a reason pandas is imported locally? Python caches it after the
first call, so this is not a major runtime cost, but the helper requires pandas
on every execution path and is invoked on each poll. A module-level import
would make that dependency clearer
##########
providers/influxdb/src/airflow/providers/influxdb/sensors/influxdb3.py:
##########
@@ -0,0 +1,102 @@
+# 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.
+"""Sensor waiting for a SQL query to return a truthy first cell in InfluxDB
3.x."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from datetime import timedelta
+from typing import TYPE_CHECKING, Any
+
+from airflow.providers.common.compat.sdk import AirflowFailException,
BaseSensorOperator, conf
+from airflow.providers.influxdb.hooks.influxdb3 import InfluxDB3Hook
+from airflow.providers.influxdb.triggers.influxdb3 import
InfluxDB3SensorTrigger
+from airflow.providers.influxdb.utils import _first_cell_is_truthy
+
+if TYPE_CHECKING:
+ from airflow.sdk.definitions.context import Context
+
+
+class InfluxDB3Sensor(BaseSensorOperator):
+ """
+ Wait until an InfluxDB 3.x SQL query returns a truthy first cell.
+
+ :param sql: The SQL query to poll.
+ :param influxdb3_conn_id: Reference to :ref:`InfluxDB 3 connection id
<howto/connection:influxdb3>`.
+ :param fail_on_empty: Fail instead of waiting when the query returns no
rows.
Review Comment:
I would included the default values for the `influxdb3_conn_id` and
`fail_on_empty` parameters.
##########
providers/influxdb/tests/unit/influxdb/triggers/test_influxdb3.py:
##########
@@ -79,3 +79,111 @@ async def test_run_propagates_cancellation(self,
mock_hook_class):
with pytest.raises(asyncio.CancelledError):
await anext(trigger.run())
+
+
+class TestInfluxDB3SensorTrigger:
+ def test_serialization(self):
+ trigger = InfluxDB3SensorTrigger(
+ sql=SQL,
+ influxdb3_conn_id=CONN_ID,
+ poll_interval=30,
+ fail_on_empty=True,
+ )
+
+ assert trigger.serialize() == (
+
"airflow.providers.influxdb.triggers.influxdb3.InfluxDB3SensorTrigger",
+ {
+ "sql": SQL,
+ "influxdb3_conn_id": CONN_ID,
+ "poll_interval": 30,
+ "fail_on_empty": True,
+ },
+ )
Review Comment:
Could we make the style of this test consistent with its counterpart for
`TestInfluxDB3QueryTrigger`?
##########
providers/influxdb/tests/unit/influxdb/triggers/test_influxdb3.py:
##########
@@ -79,3 +79,111 @@ async def test_run_propagates_cancellation(self,
mock_hook_class):
with pytest.raises(asyncio.CancelledError):
await anext(trigger.run())
+
+
+class TestInfluxDB3SensorTrigger:
+ def test_serialization(self):
+ trigger = InfluxDB3SensorTrigger(
+ sql=SQL,
+ influxdb3_conn_id=CONN_ID,
+ poll_interval=30,
+ fail_on_empty=True,
+ )
+
+ assert trigger.serialize() == (
+
"airflow.providers.influxdb.triggers.influxdb3.InfluxDB3SensorTrigger",
+ {
+ "sql": SQL,
+ "influxdb3_conn_id": CONN_ID,
+ "poll_interval": 30,
+ "fail_on_empty": True,
+ },
+ )
+
+ @pytest.mark.asyncio
+ @mock.patch("airflow.providers.influxdb.triggers.influxdb3.asyncio.sleep",
new_callable=mock.AsyncMock)
+ @mock.patch("airflow.providers.influxdb.triggers.influxdb3.InfluxDB3Hook",
autospec=True)
+ async def test_run_succeeds_without_sleep_when_condition_is_met(self,
mock_hook_class, mock_sleep):
+ mock_hook = mock_hook_class.return_value
+ mock_hook.query_async =
mock.AsyncMock(return_value=pd.DataFrame({"literal": [1]}))
+
+ events = [event async for event in
InfluxDB3SensorTrigger(sql=SQL).run()]
+
+ mock_hook_class.assert_called_once_with(conn_id="influxdb3_default")
+ mock_hook.query_async.assert_awaited_once_with(SQL)
+ mock_sleep.assert_not_awaited()
+ assert events == [TriggerEvent({"status": "success"})]
+
+ @pytest.mark.asyncio
+ @mock.patch("airflow.providers.influxdb.triggers.influxdb3.asyncio.sleep",
new_callable=mock.AsyncMock)
+ @mock.patch("airflow.providers.influxdb.triggers.influxdb3.InfluxDB3Hook",
autospec=True)
+ async def test_run_polls_until_condition_is_met(self, mock_hook_class,
mock_sleep):
+ mock_hook = mock_hook_class.return_value
+ mock_hook.query_async = mock.AsyncMock(
+ side_effect=[
+ pd.DataFrame({"literal": []}),
+ pd.DataFrame({"count": [0]}),
+ pd.DataFrame({"literal": [1]}),
+ ]
+ )
+
+ events = [event async for event in InfluxDB3SensorTrigger(sql=SQL,
poll_interval=30).run()]
+
+ assert mock_hook.query_async.await_count == 3
+ assert mock_sleep.await_args_list == [mock.call(30), mock.call(30)]
+ assert events == [TriggerEvent({"status": "success"})]
+
+ @pytest.mark.asyncio
+ @mock.patch("airflow.providers.influxdb.triggers.influxdb3.asyncio.sleep",
new_callable=mock.AsyncMock)
+ @mock.patch("airflow.providers.influxdb.triggers.influxdb3.InfluxDB3Hook",
autospec=True)
+ async def test_run_fail_on_empty(self, mock_hook_class, mock_sleep):
+ mock_hook = mock_hook_class.return_value
+ mock_hook.query_async =
mock.AsyncMock(return_value=pd.DataFrame({"literal": []}))
+
+ events = [event async for event in InfluxDB3SensorTrigger(sql=SQL,
fail_on_empty=True).run()]
+
+ mock_hook.query_async.assert_awaited_once_with(SQL)
+ mock_sleep.assert_not_awaited()
+ assert events == [
+ TriggerEvent(
+ {"status": "error", "message": "No rows returned, raising as
per fail_on_empty flag"}
+ )
+ ]
+
+ @pytest.mark.asyncio
+ @mock.patch("airflow.providers.influxdb.triggers.influxdb3.InfluxDB3Hook",
autospec=True)
+ async def test_run_failure(self, mock_hook_class):
+ mock_hook = mock_hook_class.return_value
+ mock_hook.query_async = mock.AsyncMock(side_effect=ValueError("boom"))
+
+ events = [event async for event in
InfluxDB3SensorTrigger(sql=SQL).run()]
+
+ mock_hook_class.assert_called_once_with(conn_id="influxdb3_default")
+ mock_hook.query_async.assert_awaited_once_with(SQL)
+ assert events == [TriggerEvent({"status": "error", "message": "boom"})]
+
+ @pytest.mark.asyncio
+ @mock.patch("airflow.providers.influxdb.triggers.influxdb3.InfluxDB3Hook",
autospec=True)
+ async def test_run_propagates_cancellation_during_query(self,
mock_hook_class):
+ mock_hook = mock_hook_class.return_value
+ mock_hook.query_async =
mock.AsyncMock(side_effect=asyncio.CancelledError())
+
+ with pytest.raises(asyncio.CancelledError):
+ await anext(InfluxDB3SensorTrigger(sql=SQL).run())
+
+ @pytest.mark.asyncio
+ @mock.patch(
+ "airflow.providers.influxdb.triggers.influxdb3.asyncio.sleep",
+ new_callable=mock.AsyncMock,
+ side_effect=asyncio.CancelledError(),
+ )
+ @mock.patch("airflow.providers.influxdb.triggers.influxdb3.InfluxDB3Hook",
autospec=True)
+ async def test_run_propagates_cancellation_during_sleep(self,
mock_hook_class, mock_sleep):
+ mock_hook = mock_hook_class.return_value
+ mock_hook.query_async =
mock.AsyncMock(return_value=pd.DataFrame({"literal": []}))
+
+ with pytest.raises(asyncio.CancelledError):
+ await anext(InfluxDB3SensorTrigger(sql=SQL,
poll_interval=30).run())
+
+ mock_hook.query_async.assert_awaited_once_with(SQL)
+ mock_sleep.assert_awaited_once_with(30)
Review Comment:
I would add a trigger test with `side_effect=[empty_dataframe,
ValueError("boom")]`, asserting it sleeps once and then emits the error event.
It confirms a later polling failure is handled identically to an initial
failure.
##########
providers/influxdb/tests/unit/influxdb/sensors/test_influxdb3.py:
##########
@@ -0,0 +1,132 @@
+# 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 datetime import timedelta
+from unittest import mock
+
+import pandas as pd
+import pytest
+
+from airflow.providers.common.compat.sdk import AirflowFailException,
AirflowSensorTimeout, TaskDeferred
+from airflow.providers.influxdb.sensors.influxdb3 import InfluxDB3Sensor
+from airflow.providers.influxdb.triggers.influxdb3 import
InfluxDB3SensorTrigger
+
+SQL = """SELECT 1 FROM "events" WHERE time > now() - INTERVAL '1 hour' LIMIT
1"""
+CONN_ID = "test_influxdb3_conn"
+HOOK_PATH = "airflow.providers.influxdb.sensors.influxdb3.InfluxDB3Hook"
+
+
+class TestInfluxDB3Sensor:
+ def test_init(self):
+ sensor = InfluxDB3Sensor(task_id="wait", sql=SQL)
+
+ assert sensor.sql == SQL
+ assert sensor.influxdb3_conn_id == "influxdb3_default"
+ assert sensor.fail_on_empty is False
+ assert sensor.deferrable is False
+ assert sensor.template_fields == ("sql", "influxdb3_conn_id")
+ assert sensor.template_ext == (".sql",)
+
+ @pytest.mark.parametrize(
+ ("dataframe", "expected"),
+ [
+ pytest.param(pd.DataFrame({"literal": [1]}), True,
id="numeric-one"),
+ pytest.param(pd.DataFrame({"literal": ["ready"]}), True,
id="non-empty-string"),
+ pytest.param(pd.DataFrame({"literal": []}), False, id="no-rows"),
+ pytest.param(pd.DataFrame({"count": [0]}), False,
id="numeric-zero"),
+ pytest.param(pd.DataFrame({"count": ["0"]}), False,
id="string-zero"),
+ pytest.param(pd.DataFrame({"value": [None]}), False, id="none"),
+ ],
+ )
+ @mock.patch(HOOK_PATH, autospec=True)
+ def test_poke(self, mock_hook_class, dataframe, expected):
+ mock_hook_class.return_value.query.return_value = dataframe
+ sensor = InfluxDB3Sensor(task_id="wait", sql=SQL,
influxdb3_conn_id=CONN_ID)
+
+ assert sensor.poke(context={}) is expected
+ mock_hook_class.assert_called_once_with(conn_id=CONN_ID)
+ mock_hook_class.return_value.query.assert_called_once_with(SQL)
+
+ @mock.patch(HOOK_PATH, autospec=True)
+ def test_poke_fail_on_empty(self, mock_hook_class):
+ mock_hook_class.return_value.query.return_value =
pd.DataFrame({"literal": []})
+ sensor = InfluxDB3Sensor(task_id="wait", sql=SQL, fail_on_empty=True)
+
+ with pytest.raises(AirflowFailException, match="fail_on_empty"):
+ sensor.poke(context={})
Review Comment:
I would see if you could cover the scenario where the sensor receives an
empty dataframe with `fail_on_empty=True` but before any deferral. It verifies
the initial worker-side `poke()` has the same semantics as the trigger.
##########
providers/influxdb/src/airflow/providers/influxdb/sensors/influxdb3.py:
##########
@@ -0,0 +1,102 @@
+# 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.
+"""Sensor waiting for a SQL query to return a truthy first cell in InfluxDB
3.x."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from datetime import timedelta
+from typing import TYPE_CHECKING, Any
+
+from airflow.providers.common.compat.sdk import AirflowFailException,
BaseSensorOperator, conf
+from airflow.providers.influxdb.hooks.influxdb3 import InfluxDB3Hook
+from airflow.providers.influxdb.triggers.influxdb3 import
InfluxDB3SensorTrigger
+from airflow.providers.influxdb.utils import _first_cell_is_truthy
+
+if TYPE_CHECKING:
+ from airflow.sdk.definitions.context import Context
+
+
+class InfluxDB3Sensor(BaseSensorOperator):
+ """
+ Wait until an InfluxDB 3.x SQL query returns a truthy first cell.
+
+ :param sql: The SQL query to poll.
+ :param influxdb3_conn_id: Reference to :ref:`InfluxDB 3 connection id
<howto/connection:influxdb3>`.
+ :param fail_on_empty: Fail instead of waiting when the query returns no
rows.
+ :param deferrable: Run polling in the triggerer. Defaults to the
+ ``operators.default_deferrable`` configuration (``False`` if unset).
+ """
+
+ template_fields: Sequence[str] = ("sql", "influxdb3_conn_id")
+ template_ext: Sequence[str] = (".sql",)
+
+ def __init__(
+ self,
+ *,
+ sql: str,
+ influxdb3_conn_id: str = "influxdb3_default",
+ fail_on_empty: bool = False,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable",
fallback=False),
+ **kwargs,
+ ) -> None:
+ super().__init__(**kwargs)
+ self.sql = sql
+ self.influxdb3_conn_id = influxdb3_conn_id
+ self.fail_on_empty = fail_on_empty
+ self.deferrable = deferrable
+
+ def poke(self, context: Context) -> bool:
+ """Return whether the query result meets the sensor condition."""
+ self.log.info("Poking with SQL query: %s", self.sql)
+ dataframe =
InfluxDB3Hook(conn_id=self.influxdb3_conn_id).query(self.sql)
+ if dataframe.empty and self.fail_on_empty:
+ raise AirflowFailException("No rows returned, raising as per
fail_on_empty flag")
+ return _first_cell_is_truthy(dataframe)
+
+ def execute(self, context: Context) -> None:
+ if not self.deferrable:
+ super().execute(context)
+ return
+
+ if self.poke(context):
+ return
+
+ timeout = timedelta(seconds=self.timeout) if isinstance(self.timeout,
(int, float)) else self.timeout
Review Comment:
This could just be:
`timeout=timedelta(seconds=self.timeout)`
There is no need for any normalizatio here as the parent class
`BaseSensorOperator` does this for you via `_coerce_timeout`
--
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]
