dabla commented on code in PR #65618:
URL: https://github.com/apache/airflow/pull/65618#discussion_r4015589898
##########
providers/common/sql/src/airflow/providers/common/sql/hooks/sql.pyi:
##########
@@ -185,6 +186,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: ...
+ @overload
+ async def arun(
+ self,
+ sql: str | Iterable[str],
+ autocommit: bool = ...,
+ parameters: Iterable | Mapping[str, Any] | None = ...,
+ handler: None = ...,
+ split_statements: bool = ...,
+ return_last: bool = ...,
+ ) -> None: ...
+ @overload
+ async def arun(
+ self,
+ sql: str | Iterable[str],
+ autocommit: bool = ...,
+ parameters: Iterable | Mapping[str, Any] | None = ...,
Review Comment:
**[blocker]** `arun()` type stub is out of sync with the implementation —
missing `read_only`.
In `hooks/sql.py`, both `@overload` signatures of `arun` include `read_only:
bool = ...`, and the trigger calls it with that keyword (`triggers/sql.py`:
`await hook.arun(..., read_only=self.read_only)`). The `.pyi` stub (the
declared "public interface" per its own module docstring) only kept `sql`,
`autocommit`, `parameters`, `handler`, `split_statements`, `return_last` —
`read_only` was dropped from both overloads.
Since this package ships `.pyi` stubs that mypy resolves in preference to
the real implementation, any caller passing `read_only=` (including this PR's
own trigger) will fail static type checking (`breeze run mypy path/to/code` /
the `mypy-shared-*`-style CI hooks for providers). Please add `read_only: bool
= ...,` to both overloads in the stub, matching `hooks/sql.py`.
##########
providers/common/sql/src/airflow/providers/common/sql/hooks/sql.py:
##########
@@ -1171,3 +1185,220 @@ def get_db_log_messages(self, conn) -> None:
:param conn: Connection object
"""
+
+ def _translate_sql(self, sql: str) -> str:
+ """
+ Translate SQL to driver-specific paramstyle.
+
+ DB-specific hooks may override this to translate from a canonical style
+ to their driver's paramstyle if you want a unified SQL authoring style.
+ """
+ return sql
+
+ def _prepare_parameters(self, parameters: Iterable | Mapping[str, Any] |
None):
+ """DB hooks may override to adapt parameter style."""
+ return parameters
+
+ def _build_conn_kwargs_from_airflow_connection(self, db) -> dict:
+ """Build a DB-agnostic kwargs dict."""
+ extra = {}
+ extra_dejson = getattr(db, "extra_dejson", None)
+ if isinstance(extra_dejson, dict):
+ extra = extra_dejson
+ try:
+ port = int(db.port) if db.port else None
+ except (TypeError, ValueError):
+ port = None
+ return {
+ "host": db.host or "",
+ "port": port,
+ "username": db.login or "",
+ "password": db.password or "",
+ "database": extra.get("database") or extra.get("dbname") or "",
+ "schema": db.schema or "",
+ "conn_type": db.conn_type,
+ "extra": extra,
+ "raw_connection": db,
+ }
+
+ async def aget_conn(self) -> Any:
+ if not hasattr(self, "_conn_lock"):
+ self._conn_lock = asyncio.Lock()
+ async with self._conn_lock:
+ if not self._connection:
+ self._connection = await
get_async_connection(self.get_conn_id())
+ db = self._connection
+ if self.connector is None:
+ raise RuntimeError(f"{type(self).__name__} didn't have
`self.connector` set!")
+ conn_kwargs = self._build_conn_kwargs_from_airflow_connection(db)
+ return await self.connector.connect(**conn_kwargs)
+
+ async def _call(self, func: Callable, *args, **kwargs) -> Any:
+ if inspect.iscoroutinefunction(func):
+ return await func(*args, **kwargs)
+ return await sync_to_async(func)(*args, **kwargs)
+
+ @asynccontextmanager
+ async def _acreate_autocommit_connection(self, autocommit: bool = False):
+ conn = await self.aget_conn()
+ try:
+ if self.supports_autocommit:
+ set_autocommit = getattr(conn, "set_autocommit", None)
+ if set_autocommit and
inspect.iscoroutinefunction(set_autocommit):
+ await set_autocommit(autocommit)
+ else:
+ self.set_autocommit(conn, autocommit)
+ yield conn
+ finally:
+ close = getattr(conn, "aclose", None) or getattr(conn, "close",
None)
+ if close:
+ await self._call(close)
+
+ @asynccontextmanager
+ async def _aget_cursor(self, conn):
+ cursor = getattr(conn, "cursor", None)
+ if cursor is None:
+ raise TypeError(
+ f"{type(conn).__name__} has no cursor() method. "
+ "Override _aget_cursor in the DB-specific hook."
+ )
+ cur_or_cm = cursor()
+ if inspect.isawaitable(cur_or_cm):
+ cur_or_cm = await cur_or_cm
+ if hasattr(cur_or_cm, "__aenter__") and hasattr(cur_or_cm,
"__aexit__"):
+ async with cur_or_cm as cur:
+ yield cur
+ return
+ execute = getattr(cur_or_cm, "execute", None)
+ if execute and inspect.iscoroutinefunction(execute):
+ try:
+ yield cur_or_cm
+ finally:
+ close = getattr(cur_or_cm, "aclose", None) or
getattr(cur_or_cm, "close", None)
+ if close:
+ if inspect.iscoroutinefunction(close):
+ await close()
+ else:
+ close()
+ return
+ raise RuntimeError(
+ f"Unsupported cursor type returned by {type(conn).__name__}. "
+ "Override _aget_cursor in the DB-specific hook."
+ )
+
+ def supports_async_execution(self) -> bool:
Review Comment:
**[warning]** Capability detection should be an explicit override, not
implicit reflection on method identity.
```python
def supports_async_execution(self) -> bool:
"""Whether the DB-specific hook overrides :meth:`aget_conn` for a real
async driver."""
return type(self).aget_conn is not DbApiHook.aget_conn
def supports_readonly_execution(self) -> bool:
"""Whether the DB-specific hook overrides :meth:`_aenter_read_only`."""
return type(self)._aenter_read_only is not DbApiHook._aenter_read_only
```
Comparing `type(self).aget_conn is not DbApiHook.aget_conn` infers
capability from *whether a method was overridden*, rather than a hook stating
its capability directly. This is fragile: it silently breaks for any
indirection between the subclass and `DbApiHook` — a mixin providing
`aget_conn`, a decorator wrapping the method, `functools.wraps`-based
instrumentation, or simply an intermediate abstract base between `DbApiHook`
and the concrete hook — any of these makes the identity check `True` even when
the override doesn't actually provide a working async driver, or `False` when
it does. It also makes the contract implicit: a hook author has to know this
reflection trick exists rather than just declaring "I support async."
`DbApiHook` already has the right pattern a few lines above
(`hooks/sql.py:170-173`) for exactly this kind of capability negotiation —
plain class attributes, defaulting to `False`, that a DB-specific hook
overrides:
```python
# Override if this db supports autocommit.
supports_autocommit = False
# Override if this db supports executemany.
supports_executemany = False
```
`supports_async_execution`/`supports_readonly_execution` should follow the
same convention instead of introducing a different, reflection-based mechanism
for the same kind of question:
```python
# DbApiHook
# Override if this db supports a real async driver (aget_conn/_aget_cursor).
supports_async_execution = False
# Override if this db supports a read-only transaction (_aenter_read_only).
supports_readonly_execution = False
# PostgresHook
supports_async_execution = True
supports_readonly_execution = True
```
This is simpler, is robust regardless of how the subclass structures its
overrides (mixins, decorators, indirection), and keeps all of `DbApiHook`'s
capability flags consistent with one pattern instead of two.
##########
providers/common/sql/src/airflow/providers/common/sql/triggers/sql.py:
##########
@@ -57,49 +84,160 @@ def serialize(self) -> tuple[str, dict[str, Any]]:
{
"sql": self.sql,
"conn_id": self.conn_id,
+ "autocommit": self.autocommit,
+ "parameters": self.parameters,
+ "fetch_results": self.fetch_results,
+ "split_statements": self.split_statements,
+ "return_last": self.return_last,
+ "read_only": self.read_only,
"hook_params": self.hook_params,
},
)
- def get_hook(self) -> DbApiHook:
+ @staticmethod
+ def _jsonsafe_descriptions(
+ descriptions: list[Sequence[Sequence] | None],
+ ) -> list[list[list[Any]] | None]:
+ """
+ Normalise cursor descriptions into a JSON-serializable form for the
``TriggerEvent``.
+
+ ``cursor.description`` is a sequence of column 7-tuples whose
``type_code`` can be a
+ driver-specific object that is not JSON-serializable; such values are
stringified while
+ JSON-native fields (column names, sizes, precision, ...) are preserved.
+ """
+ safe: list[list[list[Any]] | None] = []
+ for description in descriptions:
+ if description is None:
+ safe.append(None)
+ continue
+ safe.append(
+ [
+ [
+ field if isinstance(field, (str, int, float, bool,
type(None))) else str(field)
+ for field in column
+ ]
+ for column in description
+ ]
+ )
+ return safe
+
+ async def aget_hook(self) -> DbApiHook:
"""
Return DbApiHook.
:return: DbApiHook for this connection
"""
- connection = BaseHook.get_connection(self.conn_id)
- hook = connection.get_hook(hook_params=self.hook_params)
+ hook = await get_async_hook(self.conn_id, hook_params=self.hook_params)
if not isinstance(hook, DbApiHook):
- raise AirflowException(
- f"You are trying to use `common-sql` with
{hook.__class__.__name__},"
- " but its provider does not support it. Please upgrade the
provider to a version that"
- " supports `common-sql`. The hook class should be a subclass
of"
- f" `{hook.__class__.__module__}.{hook.__class__.__name__}`."
- f" Got {hook.__class__.__name__} hook with class hierarchy:
{hook.__class__.mro()}"
+ raise TypeError(
+ f"You are trying to use the SqlExecuteQueryOperator in
deferrable mode with {hook.__class__.__name__},"
+ " but its provider does not support this. Please set
deferrable=False"
+ f" Got {hook.__class__.__name__} with class hierarchy:
{hook.__class__.mro()}"
+ )
+ if self.read_only and not (hook.supports_async_execution() and
hook.supports_readonly_execution()):
+ raise NotImplementedError(
+ f"{hook.__class__.__name__} does not support read-only
execution, so it cannot run a"
+ " deferred query safely (a triggerer restart could re-run it).
Set"
+ " enforce_read_only=False to run without the read-only guard
if the query is"
+ " idempotent, or deferrable=False to run it on the worker."
)
return hook
- async def _get_records(self) -> Any:
- from asgiref.sync import sync_to_async
-
- hook = self.get_hook()
+ async def _run_query(self, hook: DbApiHook, fetch_results: bool):
+ """
+ Run the query against ``hook``, using its native async driver if
available.
+ Hooks without a real async driver
(:meth:`DbApiHook.supports_async_execution`) fall back to
+ running the synchronous :meth:`DbApiHook.run` in a worker thread,
matching the compatibility the
+ `GenericTransfer` operator has always relied on for arbitrary
DB-specific hooks.
+ """
+ if hook.supports_async_execution():
+ if fetch_results:
+ return await hook.arun(
+ sql=self.sql,
+ autocommit=self.autocommit,
+ parameters=self.parameters,
+ handler=fetch_all_handler,
+ split_statements=self.split_statements,
+ return_last=self.return_last,
+ read_only=self.read_only,
+ )
+ return await hook.arun(
+ sql=self.sql,
+ autocommit=self.autocommit,
+ parameters=self.parameters,
+ handler=None,
+ split_statements=self.split_statements,
+ return_last=self.return_last,
+ read_only=self.read_only,
+ )
if AIRFLOW_V_3_2_PLUS:
# This is only supported from Airflow 3.2 or higher due to added
async support in CommsDecoder
- return await sync_to_async(hook.get_records)(self.sql)
- return hook.get_records(self.sql)
+ # `sync_to_async` erases `run`'s overloads, so it is cast to a
plain callable first.
+ arun_in_thread = sync_to_async(cast("Callable[..., Any]",
hook.run))
+ if fetch_results:
+ return await arun_in_thread(
+ sql=self.sql,
+ autocommit=self.autocommit,
+ parameters=self.parameters,
+ handler=fetch_all_handler,
+ split_statements=self.split_statements,
+ return_last=self.return_last,
+ )
+ return await arun_in_thread(
+ sql=self.sql,
+ autocommit=self.autocommit,
+ parameters=self.parameters,
+ handler=None,
+ split_statements=self.split_statements,
+ return_last=self.return_last,
+ )
+ if fetch_results:
+ return hook.run(
+ sql=self.sql,
+ autocommit=self.autocommit,
+ parameters=self.parameters,
+ handler=fetch_all_handler,
+ split_statements=self.split_statements,
+ return_last=self.return_last,
+ )
+ return hook.run(
+ sql=self.sql,
+ autocommit=self.autocommit,
+ parameters=self.parameters,
+ handler=None,
+ split_statements=self.split_statements,
+ return_last=self.return_last,
+ )
async def run(self) -> AsyncIterator[TriggerEvent]:
try:
+ hook = await self.aget_hook()
+
self.log.info("Extracting data from %s", self.conn_id)
self.log.info("Executing: \n %s", self.sql)
- self.log.info("Reading records from %s", self.conn_id)
- results = await self._get_records()
+ if self.fetch_results:
+ # Fetch the raw rows with the built-in handler and return them
with the cursor
+ # descriptions; the operator applies any user handler on the
worker.
+ results = await self._run_query(hook, fetch_results=True)
+
+ self.log.info("Executing query from %s done!", self.conn_id)
+ self.log.debug("results: %s", results)
+ yield TriggerEvent(
+ {
+ "status": "success",
+ "results": results,
+ "descriptions":
self._jsonsafe_descriptions(hook.descriptions),
+ }
+ )
+
+ else:
+ await self._run_query(hook, fetch_results=False)
+
+ self.log.info("Executing query from %s done!", self.conn_id)
Review Comment:
**[blocker]** Unbounded result set is written to the metadata database, with
no way to bound or detect it.
The rationale for replaying results via `_ReplayCursor` instead of running
the user `handler` directly in the triggerer is sound in principle: the
triggerer must not execute arbitrary user code (it has no DAG bundle access and
the Security Model keeps user code off that process), so the built-in
`fetch_all_handler` runs in the triggerer and the user's `handler` is deferred
to the worker in `execute_complete`.
However, the mechanism for getting the rows there is not an in-memory hop —
`yield TriggerEvent({"results": ..., "descriptions": ...})` is consumed by the
scheduler and persisted into `TaskInstance.next_kwargs`, a JSON column on the
metadata database (`airflow-core/src/airflow/models/taskinstance.py:672,1864`),
until the worker picks the task instance back up and calls `execute_complete`.
Every row of every deferred query is serialized to JSON and written as a
task-instance row in Airflow's own database, then read back out — unlike the
non-deferrable path, which streams rows straight from the DB-API cursor without
ever touching Airflow's metadata store.
The docstring's guidance ("very large results are best fetched with
deferrable=False") is not enforced anywhere, and there is no mechanism in this
PR for a user or the operator to determine in advance whether a given query's
result set is "large" — the row/byte count is only known after the query has
already run and been serialized into `next_kwargs`. A single deferred task with
an unexpectedly large result set writes that entire payload into the shared
metadata database that the scheduler, webserver, and every other task instance
also depend on, which is a shared-resource risk (bloated table/row size, backup
size, replication lag, storage pressure), not merely a per-task inefficiency.
This is a design gap, not a documentation gap: relying on docstring guidance
alone is not a safeguard against a task instance writing an arbitrarily large
payload into the metadata database.
A general fix would parallel how workers get a customizable XCom backend to
keep large payloads out of the metadata database — the trigger path currently
has no equivalent, so `TriggerEvent` payloads always go through the metadata
database regardless of size. That is a broader architectural change and out of
scope for this PR, but this PR introduces the first deferred code path in
`common-sql` that can put arbitrarily large query result sets through that
channel, and currently has no limit, no size check, and no fallback for the
case where the result set turns out to be large.
##########
providers/common/sql/src/airflow/providers/common/sql/operators/generic_transfer.py:
##########
@@ -206,7 +210,7 @@ def execute_complete(
event: dict[Any, Any] | None = None,
) -> Any:
if event:
- if event.get("status") == "failure":
+ if event.get("status") == "error":
Review Comment:
**[warning]** Avoidable breaking change to `SQLExecuteQueryTrigger`, and no
changelog/version-bump entry accompanies it.
`SQLExecuteQueryTrigger.__init__` dropped `**kwargs` and now requires
`autocommit`, `split_statements`, and `return_last` as required args with no
defaults, and the `TriggerEvent` status key changed from `"failure"` to
`"error"` (reflected here in `generic_transfer.py`'s `execute_complete`).
Neither change is required by the deferrable-execution feature itself:
- **Constructor:** the old trigger's implicit behavior was equivalent to
`autocommit=False, split_statements=False, return_last=True` (it just called
`hook.get_records(self.sql)`). Giving the three new params those same defaults
would let old serialized kwargs dicts (from tasks already deferred before an
upgrade) still reconstruct successfully via `cls(**kwargs)`, while new callers
can still pass them explicitly. As written, any `GenericTransferOperator` task
already deferred before the upgrade will fail to reconstruct its trigger after
upgrade with a missing-argument `TypeError`.
- **Status rename:** `"error"` is the dominant convention across other
providers' triggers (~409 occurrences vs. ~5 for `"failure"`), but the rename
is not required by this feature, and it breaks any code (including a
subclassed/custom operator) that checks `event["status"] == "failure"`.
Review thread 8/6 raised making this "purely additive", and the response was
to merge `GenericTransfer`'s trigger into `SQLExecuteQueryTrigger` so the class
name wouldn't change — but the constructor signature and status key still
changed underneath the same class name, so the break was moved rather than
removed. Defaulting the three new params and keeping (or aliasing) `"failure"`
would remove the breaking change entirely at no functional cost.
If the breaking change is kept regardless, there is still no corresponding
update to `providers/common/sql/docs/changelog.rst` (providers don't use
newsfragments — the changelog file is the right place per
`providers/AGENTS.md`) or a provider major-version bump reflected in this diff.
##########
providers/common/sql/src/airflow/providers/common/sql/hooks/sql.py:
##########
@@ -1171,3 +1185,220 @@ def get_db_log_messages(self, conn) -> None:
:param conn: Connection object
"""
+
+ def _translate_sql(self, sql: str) -> str:
+ """
+ Translate SQL to driver-specific paramstyle.
+
+ DB-specific hooks may override this to translate from a canonical style
+ to their driver's paramstyle if you want a unified SQL authoring style.
+ """
+ return sql
+
+ def _prepare_parameters(self, parameters: Iterable | Mapping[str, Any] |
None):
+ """DB hooks may override to adapt parameter style."""
+ return parameters
+
+ def _build_conn_kwargs_from_airflow_connection(self, db) -> dict:
+ """Build a DB-agnostic kwargs dict."""
+ extra = {}
+ extra_dejson = getattr(db, "extra_dejson", None)
+ if isinstance(extra_dejson, dict):
+ extra = extra_dejson
+ try:
+ port = int(db.port) if db.port else None
+ except (TypeError, ValueError):
+ port = None
+ return {
+ "host": db.host or "",
+ "port": port,
+ "username": db.login or "",
+ "password": db.password or "",
+ "database": extra.get("database") or extra.get("dbname") or "",
+ "schema": db.schema or "",
+ "conn_type": db.conn_type,
+ "extra": extra,
+ "raw_connection": db,
+ }
+
+ async def aget_conn(self) -> Any:
+ if not hasattr(self, "_conn_lock"):
+ self._conn_lock = asyncio.Lock()
+ async with self._conn_lock:
+ if not self._connection:
+ self._connection = await
get_async_connection(self.get_conn_id())
+ db = self._connection
+ if self.connector is None:
+ raise RuntimeError(f"{type(self).__name__} didn't have
`self.connector` set!")
+ conn_kwargs = self._build_conn_kwargs_from_airflow_connection(db)
+ return await self.connector.connect(**conn_kwargs)
+
+ async def _call(self, func: Callable, *args, **kwargs) -> Any:
+ if inspect.iscoroutinefunction(func):
+ return await func(*args, **kwargs)
+ return await sync_to_async(func)(*args, **kwargs)
+
+ @asynccontextmanager
+ async def _acreate_autocommit_connection(self, autocommit: bool = False):
Review Comment:
**[warning]** Connection close and autocommit handling guess at driver shape
via duck-typing instead of an explicit per-hook contract.
```python
@asynccontextmanager
async def _acreate_autocommit_connection(self, autocommit: bool = False):
conn = await self.aget_conn()
try:
if self.supports_autocommit:
set_autocommit = getattr(conn, "set_autocommit", None)
if set_autocommit and
inspect.iscoroutinefunction(set_autocommit):
await set_autocommit(autocommit)
else:
self.set_autocommit(conn, autocommit)
yield conn
finally:
close = getattr(conn, "aclose", None) or getattr(conn, "close", None)
if close:
await self._call(close)
```
There is no standard async equivalent of PEP 249 that all async DB drivers
share, so this code cannot assume a single connection-object shape. The current
implementation compensates with `getattr`/`inspect.iscoroutinefunction` probing
at the base-class level to guess whether close is named `close()` or
`aclose()`, and whether autocommit is a coroutine method on the connection or
must go through the hook's own synchronous `set_autocommit()`. This guessing is
itself unreliable: for the one driver this PR ships, `psycopg.AsyncConnection`
exposes `close()` as a coroutine but has no `aclose()` at all, so the
`aclose`-first check silently falls through to the correct branch only because
of the `or` fallback — a driver that has neither name, or that has a
synchronous `close()` alongside an unrelated async `aclose()`, would be
mishandled.
The read-only enforcement path in this same PR (`_aenter_read_only`) uses
the opposite, more robust pattern: the base class raises `NotImplementedError`
by default, and `PostgresHook` overrides it with driver-specific knowledge.
`_acreate_autocommit_connection` should follow that same contract instead of
guessing generically:
```python
# DbApiHook (base)
async def _aclose_conn(self, conn) -> None:
raise NotImplementedError(f"{type(self).__name__} does not implement
_aclose_conn.")
async def _aset_autocommit(self, conn, autocommit: bool) -> None:
raise NotImplementedError(f"{type(self).__name__} does not implement
_aset_autocommit.")
@asynccontextmanager
async def _acreate_autocommit_connection(self, autocommit: bool = False):
conn = await self.aget_conn()
try:
if self.supports_autocommit:
await self._aset_autocommit(conn, autocommit)
yield conn
finally:
await self._aclose_conn(conn)
# PostgresHook
async def _aclose_conn(self, conn) -> None:
await conn.close()
async def _aset_autocommit(self, conn, autocommit: bool) -> None:
await conn.set_autocommit(autocommit)
```
This states the actual driver contract explicitly per hook, the same fix as
comment [3]'s capability flags — one consistent override pattern instead of
base-class code inferring behavior it cannot know generically.
--
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]
