dabla commented on code in PR #65618:
URL: https://github.com/apache/airflow/pull/65618#discussion_r4510809585
##########
providers/common/sql/src/airflow/providers/common/sql/operators/sql.py:
##########
@@ -575,29 +645,78 @@ def _should_run_output_processing(self) -> bool:
def execute(self, context):
self.log.info("Executing: %s", self.sql)
- hook = self.get_db_hook()
- if self.split_statements is not None:
- extra_kwargs = {"split_statements": self.split_statements}
+ if self.deferrable:
+ if self.enforce_read_only:
+ is_write, reason = scan_for_writes(self.sql)
+ if is_write:
+ raise ValueError(
+ f"enforce_read_only=True but the SQL appears to
contain a write: {reason}. "
+ "Set enforce_read_only=False if the query is
idempotent, or deferrable=False "
+ "to run it on the worker."
+ )
+ self.defer(
+ trigger=SQLExecuteQueryTrigger(
+ sql=self.sql,
+ conn_id=self.conn_id,
Review Comment:
**[blocker]** In deferrable mode, `database` and `hook_params` never reach
the trigger, so the query runs against the connection's default database.
The non-deferrable path gets its hook from `get_db_hook()` →
`BaseSQLOperator._hook`, which applies `self.hook_params` (plus the
connection-type defaults from `get_hook`) and then overrides
`hook.database`/`hook.schema` with `self.database`. The trigger is built from
`conn_id` alone and calls `get_async_hook(self.conn_id,
hook_params=self.hook_params)` with its own `hook_params`, which is `None`
here. As a result:
```python
SQLExecuteQueryOperator(task_id="t", conn_id="pg", database="reporting",
sql="SELECT ...", deferrable=True)
```
runs against the connection's default database, not `reporting`. The task
still succeeds and returns rows from the wrong database. Any `hook_params`
(e.g. `options`, `enable_log_db_messages`) are dropped the same way.
Suggest passing `hook_params=self.hook_params` to the trigger and carrying
`database` as well (a `database` kwarg on the trigger that `aget_hook` applies
the way `_hook` does), plus one test per parameter that checks the hook the
trigger builds.
##########
providers/common/sql/src/airflow/providers/common/sql/triggers/sql.py:
##########
@@ -17,37 +17,65 @@
# under the License.
from __future__ import annotations
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, cast
-from airflow.providers.common.compat.sdk import AirflowException, BaseHook
+from asgiref.sync import sync_to_async
+
+from airflow.providers.common.compat.hook import get_async_hook
from airflow.providers.common.compat.version_compat import AIRFLOW_V_3_2_PLUS
+from airflow.providers.common.sql.hooks.handlers import fetch_all_handler
from airflow.providers.common.sql.hooks.sql import DbApiHook
from airflow.triggers.base import BaseTrigger, TriggerEvent
if TYPE_CHECKING:
- from collections.abc import AsyncIterator
+ from collections.abc import AsyncIterator, Callable
from typing import Any
+from collections.abc import (
+ Iterable,
+ Mapping,
+ Sequence,
+)
+
class SQLExecuteQueryTrigger(BaseTrigger):
"""
- A trigger that executes SQL code in async mode.
+ A SQL trigger that executes SQL code in async mode.
+
+ The query runs in the triggerer, but no user code does: when
``fetch_results`` is set the rows are
+ fetched with the built-in :func:`fetch_all_handler` and returned, together
with the cursor
+ descriptions, in the ``TriggerEvent``. Any user-provided ``handler`` is
applied on the worker in
+ ``SQLExecuteQueryOperator.execute_complete`` -- keeping user code out of
the triggerer's event loop
+ and out of its (bundle-less) import path.
:param sql: the sql statement to be executed (str) or a list of sql
statements to execute
:param conn_id: the connection ID used to connect to the database
+ :param fetch_results: whether the query results should be fetched and
returned to the worker
:param hook_params: hook parameters
"""
def __init__(
self,
- sql: str | list[str],
+ *,
+ sql: str | Iterable[str],
conn_id: str,
+ autocommit: bool = False,
+ split_statements: bool = False,
+ return_last: bool = True,
+ parameters: Iterable[Any] | Mapping[str, Any] | None = None,
+ fetch_results: bool = False,
Review Comment:
**[blocker]** `fetch_results=False` by default silently ends
`GenericTransfer` runs that were deferred before the upgrade.
Thanks for adding the defaults (fb4f274862). However, the old trigger always
fetched rows (`hook.get_records(self.sql)`), so the old behaviour corresponds
to `fetch_results=True`. A `GenericTransfer` trigger serialized by the previous
provider version stores only `sql`, `conn_id` and `hook_params`. After the
upgrade it is rebuilt with `fetch_results=False`, yields `{"status":
"success"}` without `results`, and `GenericTransfer.execute_complete` reads
`rows = event.get("results")` as `None`. It then logs "No more rows to fetch …
ending transfer" and the task succeeds with the remaining pages never copied.
Defaulting `fetch_results` to `True` restores the old behaviour for rebuilt
triggers. `SQLExecuteQueryOperator` always passes the value explicitly, so it
is not affected. A test that builds the trigger from the old three-key kwargs
dict and checks that rows come back would pin this down.
##########
providers/common/sql/src/airflow/providers/common/sql/hooks/sql.pyi:
##########
@@ -179,6 +180,27 @@ class DbApiHook(BaseHook):
@staticmethod
def get_openlineage_authority_part(connection, default_port: int | None =
None) -> str: ...
def get_db_log_messages(self, conn) -> None: ...
+ async def aget_conn(self) -> Any: ...
Review Comment:
**[warning]** The stub also lacks the capability checks the trigger calls.
Besides `read_only` on the `arun` overloads, `triggers/sql.py` calls
`hook.supports_async_execution()` and `hook.supports_readonly_execution()`. The
`.pyi` declares neither, so mypy, which reads the stub instead of `sql.py`,
reports `attr-defined` on those calls. Whatever form the capability check ends
up taking (methods, or class attributes next to `supports_autocommit`), it
needs an entry here too.
##########
providers/postgres/src/airflow/providers/postgres/hooks/postgres.py:
##########
@@ -316,16 +316,43 @@ def get_conn(self) -> CompatConnection:
if arg_name not in self.ignored_extra_options:
conn_args[arg_name] = arg_val
- raw_cursor = conn.extra_dejson.get("cursor")
+ return conn_args, conn
+ def get_conn(self) -> CompatConnection:
+ """Establish a connection to a postgres database."""
+ conn_args, conn = self._build_conn_args()
+ raw_cursor = conn.extra_dejson.get("cursor")
if raw_cursor:
key, value = self._get_cursor_config(raw_cursor)
conn_args[key] = value
-
self.conn = self._create_connection(conn_args)
-
return self.conn
+ async def aget_conn(self) -> Any:
+ """Establish an async connection to a postgres database."""
+ if not USE_PSYCOPG3:
+ raise NotImplementedError("Async connections for PostgresHook
require psycopg3.")
+ from psycopg import AsyncConnection
+
+ conn_args, conn = self._build_conn_args()
Review Comment:
**[warning]** `aget_conn` makes a sync connection lookup on the triggerer's
event loop.
`_build_conn_args()` starts with `deepcopy(self.connection)`. The hook comes
from `get_async_hook` → `BaseHook.aget_hook`, which fetches the connection
asynchronously but builds the hook with only the conn id
(`hook_class(**{conn_id_attr: self.conn_id}, **hook_params)` in
`Connection.get_hook`). So `self.connection` is still unset and
`DbApiHook.connection` calls the sync `get_connection()` from inside `async def
run()`.
On `main` this works only through the greenback portal in
`TriggerCommsDecoder.send`
(`airflow-core/src/airflow/jobs/triggerer_job_runner.py`), and it raises
`RuntimeError` when `AIRFLOW_DISABLE_GREENBACK_PORTAL` is set. That portal
arrived in 3.3.0 with #66412 ("Fix triggerer crash when multiple triggers call
sync SDK methods concurrently"). The postgres provider still supports
`apache-airflow>=2.11.0`, so on 3.0–3.2 several of these triggers running at
once hit the crash #66412 fixed. The IAM branch of `_build_conn_args`
(`get_iam_token`, sync AWS calls) also blocks the loop.
Suggest fetching the connection asynchronously before building the args,
e.g. `self.connection = await get_async_connection(self.get_conn_id())` from
`common.compat`, the same helper the trigger was moved to.
--
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]
