ashb commented on code in PR #71976:
URL: https://github.com/apache/airflow/pull/71976#discussion_r3889320755
##########
providers/influxdb/docs/operators/index.rst:
##########
@@ -46,3 +46,26 @@ Example usage:
:language: python
:start-after: [START howto_operator_influxdb3]
:end-before: [END howto_operator_influxdb3]
+
+Deferrable mode
+^^^^^^^^^^^^^^^
+
+Set ``deferrable=True`` to release the worker slot while the query runs. The
task is resumed by the
+:class:`~airflow.providers.influxdb.triggers.influxdb3.InfluxDB3QueryTrigger`
once results are ready.
+
+.. exampleinclude:: /../../influxdb/tests/system/influxdb/example_influxdb3.py
+ :language: python
+ :start-after: [START howto_operator_influxdb3_deferrable]
+ :end-before: [END howto_operator_influxdb3_deferrable]
+
+.. note::
+
+ This implementation follows the upstream ``influxdb3-python`` client,
which documents querying
+ through the `Flight client
<https://influxdb3-python.readthedocs.io/en/latest/>`__ and exposes
+ `query_async()
<https://influxdb3-python.readthedocs.io/en/latest/api_reference/>`__ as a
+ single async query call. The trigger therefore awaits one query call
instead of polling at an
+ interval, and there is no ``poll_interval`` parameter.
Review Comment:
Not sure we need to mention there's no poll interval in the docs. It makes
sense in the pr as there issue talked about it, but just going to the docs most
people will have no context on this (and its not universal on deferrable
operators either )
##########
providers/influxdb/src/airflow/providers/influxdb/hooks/influxdb3.py:
##########
@@ -41,14 +43,23 @@
InfluxDBClient3 = None # type: ignore[assignment, misc]
Point = None # type: ignore[assignment, misc]
-from airflow.providers.common.compat.sdk import BaseHook
+from airflow.providers.common.compat.sdk import
AirflowOptionalProviderFeatureException, BaseHook
if TYPE_CHECKING:
import pandas as pd
from airflow.models import Connection
+class InfluxDB3AsyncQueryNotAvailableError(RuntimeError):
Review Comment:
Given we depend on an updated version of theof influx client is this
possible to hit?
##########
providers/influxdb/src/airflow/providers/influxdb/operators/influxdb3.py:
##########
@@ -23,13 +23,21 @@
from collections.abc import Sequence
from typing import TYPE_CHECKING, Any
-from airflow.providers.common.compat.sdk import BaseOperator
+from airflow.providers.common.compat.sdk import BaseOperator, conf
from airflow.providers.influxdb.hooks.influxdb3 import InfluxDB3Hook
+from airflow.providers.influxdb.triggers.influxdb3 import InfluxDB3QueryTrigger
if TYPE_CHECKING:
+ import pandas as pd
+
from airflow.sdk.definitions.context import Context
+def _convert_dataframe_to_records(dataframe: pd.DataFrame) -> list[dict[str,
Any]]:
Review Comment:
Duplication from the hook fn
##########
providers/influxdb/tests/system/influxdb/example_influxdb3.py:
##########
@@ -57,6 +57,15 @@ def write_to_influxdb3():
)
# [END howto_operator_influxdb3]
+# [START howto_operator_influxdb3_deferrable]
+deferrable_query_task = InfluxDB3Operator(
+ task_id="query_data_deferrable",
+ sql="SELECT * FROM \"temperature\" WHERE time > now() - INTERVAL '1 hour'",
+ influxdb3_conn_id="influxdb3_default",
+ deferrable=True,
Review Comment:
Deferrable is true for other operators that support it isn't it?
##########
providers/influxdb/src/airflow/providers/influxdb/hooks/influxdb3.py:
##########
@@ -205,6 +222,45 @@ def query(self, query: str) -> pd.DataFrame:
return result
+ async def query_async(self, query: str) -> list[dict[str, Any]]:
+ """
+ Run a SQL query from the triggerer and return JSON-serializable
records.
+
+ ``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 setup and DataFrame conversion
are offloaded
+ with ``asyncio.to_thread`` so that no step runs on the triggerer's
event loop.
+
+ :param query: SQL query string
+ :return: List of dictionaries representing query results
+ """
+ client = await asyncio.to_thread(self.get_conn)
+ if not hasattr(client, "query_async"):
+ raise InfluxDB3AsyncQueryNotAvailableError(
+ "Deferrable mode requires an InfluxDB 3 client that exposes "
+ "InfluxDBClient3.query_async(). Reinstall or upgrade the
provider "
+ "dependencies to use influxdb3-python>=0.12.0."
+ )
+
+ try:
+ import pandas as pd
+ except ImportError as e:
+ raise AirflowOptionalProviderFeatureException(
+ "pandas is required for InfluxDB 3 query results. Install it
with: "
+ "pip install 'apache-airflow-providers-influxdb[pandas]'"
+ ) from e
+
+ result = await client.query_async(query=query, language="sql",
mode="pandas")
+
+ if not isinstance(result, pd.DataFrame):
+ raise ValueError(
+ f"Query did not return a DataFrame. "
+ f"Result type:
{type(result).__module__}.{type(result).__name__}"
+ )
+
+ return await asyncio.to_thread(_convert_dataframe_to_records, result)
Review Comment:
Why the different return type on Async? The sync path returns the pandas
object directly doesn't it?
##########
providers/influxdb/src/airflow/providers/influxdb/hooks/influxdb3.py:
##########
@@ -205,6 +222,45 @@ def query(self, query: str) -> pd.DataFrame:
return result
+ async def query_async(self, query: str) -> list[dict[str, Any]]:
+ """
+ Run a SQL query from the triggerer and return JSON-serializable
records.
+
+ ``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 setup and DataFrame conversion
are offloaded
+ with ``asyncio.to_thread`` so that no step runs on the triggerer's
event loop.
+
+ :param query: SQL query string
+ :return: List of dictionaries representing query results
+ """
+ client = await asyncio.to_thread(self.get_conn)
Review Comment:
Those does move it off the main loop but isn't the best way of doing this.
Please look in other hooks for an `awit self.aget_conn()` or similar
##########
providers/influxdb/src/airflow/providers/influxdb/hooks/influxdb3.py:
##########
@@ -205,6 +222,45 @@ def query(self, query: str) -> pd.DataFrame:
return result
+ async def query_async(self, query: str) -> list[dict[str, Any]]:
+ """
+ Run a SQL query from the triggerer and return JSON-serializable
records.
+
+ ``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 setup and DataFrame conversion
are offloaded
+ with ``asyncio.to_thread`` so that no step runs on the triggerer's
event loop.
+
+ :param query: SQL query string
+ :return: List of dictionaries representing query results
+ """
+ client = await asyncio.to_thread(self.get_conn)
+ if not hasattr(client, "query_async"):
+ raise InfluxDB3AsyncQueryNotAvailableError(
+ "Deferrable mode requires an InfluxDB 3 client that exposes "
+ "InfluxDBClient3.query_async(). Reinstall or upgrade the
provider "
+ "dependencies to use influxdb3-python>=0.12.0."
+ )
Review Comment:
```suggestion
```
##########
providers/influxdb/tests/system/influxdb/example_influxdb3.py:
##########
@@ -57,6 +57,15 @@ def write_to_influxdb3():
)
# [END howto_operator_influxdb3]
+# [START howto_operator_influxdb3_deferrable]
+deferrable_query_task = InfluxDB3Operator(
+ task_id="query_data_deferrable",
+ sql="SELECT * FROM \"temperature\" WHERE time > now() - INTERVAL '1 hour'",
Review Comment:
```suggestion
sql="""SELECT * FROM "temperature" WHERE time > now() - INTERVAL '1
hour'""",
```
##########
providers/influxdb/tests/unit/influxdb/hooks/test_influxdb3.py:
##########
@@ -85,6 +85,46 @@ def test_query(self):
assert isinstance(result, pd.DataFrame)
assert len(result) == 2
+ @pytest.mark.asyncio
+ async def test_query_async(self):
+ """Test async query with InfluxDB 3.x."""
+ pd = pytest.importorskip("pandas")
+
+ self.influxdb3_hook.client = mock.Mock()
Review Comment:
If at all reasonable is rather we didn't mock the client, but instead the
http transport/reaponse underneath it.
At the very least, this mock needs a spec parameter to enforce its "shape"
--
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]
