SameerMesiah97 commented on code in PR #71976:
URL: https://github.com/apache/airflow/pull/71976#discussion_r4018654113
##########
providers/influxdb/src/airflow/providers/influxdb/hooks/influxdb3.py:
##########
@@ -205,6 +213,32 @@ def query(self, query: str) -> pd.DataFrame:
return result
+ async def query_async(self, query: str) -> pd.DataFrame:
+ """
+ Run a SQL query from the triggerer and return results as a pandas
DataFrame.
+
+ ``InfluxDBClient3.query_async`` runs the blocking Arrow Flight calls
in the event
+ loop's default executor. It is a plain coroutine that resolves once
the whole result
+ stream has been read -- InfluxDB 3 has no submit-then-poll query API,
so there is
+ nothing to poll in between. Connection retrieval uses the hook's async
connection
+ accessor before invoking the client coroutine.
Review Comment:
This docstring is too large. I would just keep the first line and leave the
remaining context for a comment.
##########
providers/influxdb/tests/unit/influxdb/operators/test_influxdb3.py:
##########
@@ -57,3 +60,64 @@ def test_execute(self, mock_hook_class):
assert isinstance(result[0], dict)
assert "col1" in result[0]
assert "col2" in result[0]
+
+
@mock.patch("airflow.providers.influxdb.operators.influxdb3.InfluxDB3Hook",
autospec=True)
+ def test_execute_deferrable_defers(self, mock_hook_class):
+ """In deferrable mode the operator defers instead of running the query
on the worker."""
+ operator = InfluxDB3Operator(
+ task_id="test_task_deferrable",
+ sql='SELECT "duration" FROM "pyexample"',
+ influxdb3_conn_id="influxdb3_default",
+ deferrable=True,
Review Comment:
I would set a non-default `execution_timeout` argument here and verify that
it is propagated to the trigger.
##########
providers/influxdb/tests/unit/influxdb/hooks/test_influxdb3.py:
##########
@@ -52,6 +52,20 @@ def test_get_conn(self, influx_db_client_3):
assert self.influxdb3_hook.get_client is not None
+ @pytest.mark.asyncio
+ @mock.patch("airflow.providers.influxdb.hooks.influxdb3.InfluxDBClient3")
Review Comment:
I would add the following patch here:
`@mock.patch("airflow.providers.influxdb.hooks.influxdb3.get_async_connection")`
Adjust the rest of the test accordingly. Currently, it seems like you are
patching the hook internals when `get_async_connection` is being called in
`aget_conn`
##########
providers/influxdb/tests/unit/influxdb/triggers/test_influxdb3.py:
##########
@@ -0,0 +1,66 @@
+# 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.influxdb.triggers.influxdb3 import InfluxDB3QueryTrigger
+from airflow.triggers.base import TriggerEvent
+
+SQL = 'SELECT "duration" FROM "pyexample"'
+
+
+class TestInfluxDB3QueryTrigger:
+ def test_serialization(self):
+ """Trigger round-trips its constructor arguments."""
+ trigger = InfluxDB3QueryTrigger(sql=SQL,
influxdb3_conn_id="influxdb3_default")
Review Comment:
Use a non-default ID to make the test stronger.
##########
providers/influxdb/src/airflow/providers/influxdb/operators/influxdb3.py:
##########
@@ -40,6 +41,9 @@ class InfluxDB3Operator(BaseOperator):
:param sql: The SQL query to be executed
:param influxdb3_conn_id: Reference to :ref:`InfluxDB 3 connection id
<howto/connection:influxdb3>`.
+ :param deferrable: Run the query from the triggerer so the worker slot is
released while the
+ query runs. This is most useful for long-running queries that return
small-to-moderate
+ result sets because the full result still flows back through XCom.
Review Comment:
The default value for deferrable should be mentioned in the docstring.
##########
providers/influxdb/tests/unit/influxdb/triggers/test_influxdb3.py:
##########
@@ -0,0 +1,66 @@
+# 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.influxdb.triggers.influxdb3 import InfluxDB3QueryTrigger
+from airflow.triggers.base import TriggerEvent
+
+SQL = 'SELECT "duration" FROM "pyexample"'
+
+
+class TestInfluxDB3QueryTrigger:
+ def test_serialization(self):
+ """Trigger round-trips its constructor arguments."""
Review Comment:
Is 'round-trip' accurate here? wouldn't something like 'Serializes its
constructor arguments' be more correct?
##########
providers/influxdb/tests/unit/influxdb/test_utils.py:
##########
@@ -0,0 +1,37 @@
+# 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
+
+import pytest
+
+from airflow.providers.influxdb.utils import _convert_dataframe_to_records
+
+
+def test_convert_dataframe_to_records_serializes_rows_and_timestamps():
+ pd = pytest.importorskip("pandas")
Review Comment:
Same here. If pandas is a required dependency, importorskip should be
removed.
##########
providers/influxdb/tests/unit/influxdb/operators/test_influxdb3.py:
##########
@@ -57,3 +60,64 @@ def test_execute(self, mock_hook_class):
assert isinstance(result[0], dict)
assert "col1" in result[0]
assert "col2" in result[0]
+
+
@mock.patch("airflow.providers.influxdb.operators.influxdb3.InfluxDB3Hook",
autospec=True)
+ def test_execute_deferrable_defers(self, mock_hook_class):
+ """In deferrable mode the operator defers instead of running the query
on the worker."""
+ operator = InfluxDB3Operator(
+ task_id="test_task_deferrable",
+ sql='SELECT "duration" FROM "pyexample"',
+ influxdb3_conn_id="influxdb3_default",
+ deferrable=True,
+ )
+
+ with pytest.raises(TaskDeferred) as exc:
+ operator.execute(context={})
+
+ assert isinstance(exc.value.trigger, InfluxDB3QueryTrigger)
+ assert exc.value.trigger.sql == 'SELECT "duration" FROM "pyexample"'
+ assert exc.value.trigger.influxdb3_conn_id == "influxdb3_default"
+ assert exc.value.method_name == "execute_complete"
+ mock_hook_class.assert_not_called()
+
+ def test_execute_complete_success(self):
+ """execute_complete returns the records carried by the trigger
event."""
+ operator = InfluxDB3Operator(
+ task_id="test_task_complete",
+ sql='SELECT "duration" FROM "pyexample"',
+ deferrable=True,
+ )
+ records = [{"col1": 1, "col2": 3}]
+
+ result = operator.execute_complete(context={}, event={"status":
"success", "records": records})
+
+ assert result == records
+
+ def test_execute_complete_error(self):
+ """execute_complete surfaces a failed query as a runtime error."""
+ operator = InfluxDB3Operator(
+ task_id="test_task_complete_error",
+ sql='SELECT "duration" FROM "pyexample"',
+ deferrable=True,
+ )
+
+ with pytest.raises(RuntimeError, match="boom"):
+ operator.execute_complete(context={}, event={"status": "error",
"message": "boom"})
Review Comment:
I would cover the fallback error message too. Please see the below:
```
@pytest.mark.parametrize(
("event", "match"),
[
pytest.param(
{"status": "error", "message": "boom"},
"boom",
id="error-with-message",
),
pytest.param(
{"status": "error"},
"InfluxDB 3 query failed",
id="error-without-message",
),
],
)
def test_execute_complete_error(self, event, match):
operator = InfluxDB3Operator(
task_id="test_task_complete_error",
sql='SELECT "duration" FROM "pyexample"',
deferrable=True,
)
with pytest.raises(RuntimeError, match=match):
operator.execute_complete(context={}, event=event)
```
##########
providers/influxdb/tests/unit/influxdb/operators/test_influxdb3.py:
##########
@@ -57,3 +60,64 @@ def test_execute(self, mock_hook_class):
assert isinstance(result[0], dict)
assert "col1" in result[0]
assert "col2" in result[0]
+
+
@mock.patch("airflow.providers.influxdb.operators.influxdb3.InfluxDB3Hook",
autospec=True)
+ def test_execute_deferrable_defers(self, mock_hook_class):
+ """In deferrable mode the operator defers instead of running the query
on the worker."""
+ operator = InfluxDB3Operator(
+ task_id="test_task_deferrable",
+ sql='SELECT "duration" FROM "pyexample"',
+ influxdb3_conn_id="influxdb3_default",
+ deferrable=True,
+ )
+
+ with pytest.raises(TaskDeferred) as exc:
+ operator.execute(context={})
+
+ assert isinstance(exc.value.trigger, InfluxDB3QueryTrigger)
+ assert exc.value.trigger.sql == 'SELECT "duration" FROM "pyexample"'
+ assert exc.value.trigger.influxdb3_conn_id == "influxdb3_default"
+ assert exc.value.method_name == "execute_complete"
+ mock_hook_class.assert_not_called()
+
+ def test_execute_complete_success(self):
+ """execute_complete returns the records carried by the trigger
event."""
+ operator = InfluxDB3Operator(
+ task_id="test_task_complete",
+ sql='SELECT "duration" FROM "pyexample"',
+ deferrable=True,
+ )
+ records = [{"col1": 1, "col2": 3}]
+
+ result = operator.execute_complete(context={}, event={"status":
"success", "records": records})
+
+ assert result == records
+
+ def test_execute_complete_error(self):
+ """execute_complete surfaces a failed query as a runtime error."""
+ operator = InfluxDB3Operator(
+ task_id="test_task_complete_error",
+ sql='SELECT "duration" FROM "pyexample"',
+ deferrable=True,
+ )
+
+ with pytest.raises(RuntimeError, match="boom"):
+ operator.execute_complete(context={}, event={"status": "error",
"message": "boom"})
+
+ @pytest.mark.parametrize(
+ ("event", "match"),
+ [
+ (None, "did not return an event"),
+ ({"status": "cancelled"}, "unexpected status"),
+ ],
Review Comment:
Please add IDs to each of these test cases.
##########
providers/influxdb/tests/unit/influxdb/triggers/test_influxdb3.py:
##########
@@ -0,0 +1,66 @@
+# 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.influxdb.triggers.influxdb3 import InfluxDB3QueryTrigger
+from airflow.triggers.base import TriggerEvent
+
+SQL = 'SELECT "duration" FROM "pyexample"'
+
+
+class TestInfluxDB3QueryTrigger:
+ def test_serialization(self):
+ """Trigger round-trips its constructor arguments."""
+ trigger = InfluxDB3QueryTrigger(sql=SQL,
influxdb3_conn_id="influxdb3_default")
+ classpath, kwargs = trigger.serialize()
+
+ assert classpath ==
"airflow.providers.influxdb.triggers.influxdb3.InfluxDB3QueryTrigger"
+ assert kwargs == {"sql": SQL, "influxdb3_conn_id": "influxdb3_default"}
+
+ @pytest.mark.asyncio
+ @mock.patch("airflow.providers.influxdb.triggers.influxdb3.InfluxDB3Hook",
autospec=True)
+ async def test_run_success(self, mock_hook_class):
+ """A completed query emits a single success event carrying
JSON-serializable records."""
+ pd = pytest.importorskip("pandas")
Review Comment:
if pandas is a required dependency (correct me if am wrong here), why use
importorskip? why not import it at top level?
##########
providers/influxdb/tests/unit/influxdb/triggers/test_influxdb3.py:
##########
@@ -0,0 +1,66 @@
+# 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.influxdb.triggers.influxdb3 import InfluxDB3QueryTrigger
+from airflow.triggers.base import TriggerEvent
+
+SQL = 'SELECT "duration" FROM "pyexample"'
+
+
+class TestInfluxDB3QueryTrigger:
+ def test_serialization(self):
+ """Trigger round-trips its constructor arguments."""
+ trigger = InfluxDB3QueryTrigger(sql=SQL,
influxdb3_conn_id="influxdb3_default")
+ classpath, kwargs = trigger.serialize()
+
+ assert classpath ==
"airflow.providers.influxdb.triggers.influxdb3.InfluxDB3QueryTrigger"
+ assert kwargs == {"sql": SQL, "influxdb3_conn_id": "influxdb3_default"}
+
+ @pytest.mark.asyncio
+ @mock.patch("airflow.providers.influxdb.triggers.influxdb3.InfluxDB3Hook",
autospec=True)
+ async def test_run_success(self, mock_hook_class):
+ """A completed query emits a single success event carrying
JSON-serializable records."""
+ pd = pytest.importorskip("pandas")
+
+ dataframe = pd.DataFrame({"col1": [1, 2], "col2": [3, 4]})
+ records = [{"col1": 1, "col2": 3}, {"col1": 2, "col2": 4}]
+ mock_hook = mock_hook_class.return_value
+ mock_hook.query_async = mock.AsyncMock(return_value=dataframe)
+
+ trigger = InfluxDB3QueryTrigger(sql=SQL)
+ events = [event async for event in trigger.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": "success", "records":
records})]
+
+ @pytest.mark.asyncio
+ @mock.patch("airflow.providers.influxdb.triggers.influxdb3.InfluxDB3Hook",
autospec=True)
+ async def test_run_failure(self, mock_hook_class):
+ """A failing query is reported as an event, not raised out of the
triggerer."""
+ mock_hook = mock_hook_class.return_value
+ mock_hook.query_async = mock.AsyncMock(side_effect=ValueError("boom"))
+
+ trigger = InfluxDB3QueryTrigger(sql=SQL)
+ events = [event async for event in trigger.run()]
+
+ assert events == [TriggerEvent({"status": "error", "message": "boom"})]
Review Comment:
I would add a test to cover cancellation propagation too. Please see the
below:
```
@pytest.mark.asyncio
@mock.patch(
"airflow.providers.influxdb.triggers.influxdb3.InfluxDB3Hook",
autospec=True,
)
async def test_run_propagates_cancellation(self, mock_hook_class):
mock_hook_class.return_value.query_async = mock.AsyncMock(
side_effect=asyncio.CancelledError
)
trigger = InfluxDB3QueryTrigger(sql=SQL)
with pytest.raises(asyncio.CancelledError):
await anext(trigger.run())
```
##########
providers/influxdb/tests/unit/influxdb/triggers/test_influxdb3.py:
##########
@@ -0,0 +1,66 @@
+# 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.influxdb.triggers.influxdb3 import InfluxDB3QueryTrigger
+from airflow.triggers.base import TriggerEvent
+
+SQL = 'SELECT "duration" FROM "pyexample"'
+
+
+class TestInfluxDB3QueryTrigger:
+ def test_serialization(self):
+ """Trigger round-trips its constructor arguments."""
+ trigger = InfluxDB3QueryTrigger(sql=SQL,
influxdb3_conn_id="influxdb3_default")
+ classpath, kwargs = trigger.serialize()
+
+ assert classpath ==
"airflow.providers.influxdb.triggers.influxdb3.InfluxDB3QueryTrigger"
+ assert kwargs == {"sql": SQL, "influxdb3_conn_id": "influxdb3_default"}
+
+ @pytest.mark.asyncio
+ @mock.patch("airflow.providers.influxdb.triggers.influxdb3.InfluxDB3Hook",
autospec=True)
+ async def test_run_success(self, mock_hook_class):
+ """A completed query emits a single success event carrying
JSON-serializable records."""
+ pd = pytest.importorskip("pandas")
+
+ dataframe = pd.DataFrame({"col1": [1, 2], "col2": [3, 4]})
+ records = [{"col1": 1, "col2": 3}, {"col1": 2, "col2": 4}]
+ mock_hook = mock_hook_class.return_value
+ mock_hook.query_async = mock.AsyncMock(return_value=dataframe)
+
+ trigger = InfluxDB3QueryTrigger(sql=SQL)
+ events = [event async for event in trigger.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": "success", "records":
records})]
+
+ @pytest.mark.asyncio
+ @mock.patch("airflow.providers.influxdb.triggers.influxdb3.InfluxDB3Hook",
autospec=True)
+ async def test_run_failure(self, mock_hook_class):
+ """A failing query is reported as an event, not raised out of the
triggerer."""
+ mock_hook = mock_hook_class.return_value
+ mock_hook.query_async = mock.AsyncMock(side_effect=ValueError("boom"))
+
+ trigger = InfluxDB3QueryTrigger(sql=SQL)
+ events = [event async for event in trigger.run()]
+
+ assert events == [TriggerEvent({"status": "error", "message": "boom"})]
Review Comment:
I would add these 2 asserts here to verify that the query was attempted:
```
mock_hook_class.assert_called_once_with(conn_id="influxdb3_default")
mock_hook.query_async.assert_awaited_once_with(SQL)
```
--
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]
