From 70704b799dc31c848b49b31f9944ca2a062daeb1 Mon Sep 17 00:00:00 2001 From: Karen Braganza Date: Tue, 30 Jun 2026 13:32:23 -0400 Subject: [PATCH 1/9] Add async methods to DbApiHook to support deferrable SQL operators --- generated/known_airflow_exceptions.txt | 4 +- providers/common/sql/docs/operators.rst | 1 + .../airflow/providers/common/sql/hooks/sql.py | 220 ++++++- .../providers/common/sql/hooks/sql.pyi | 24 +- .../common/sql/operators/generic_transfer.py | 6 +- .../providers/common/sql/operators/sql.py | 157 ++++- .../providers/common/sql/triggers/sql.py | 156 ++++- .../providers/common/sql/triggers/sql.pyi | 2 +- .../tests/unit/common/sql/hooks/test_sql.py | 566 +++++++++++++++++- .../common/sql/operators/test_sql_execute.py | 107 ++++ .../unit/common/sql/triggers/test_sql.py | 107 +++- .../providers/postgres/hooks/postgres.py | 37 +- .../unit/postgres/hooks/test_postgres.py | 138 ++++- .../snowflake/operators/snowflake.py | 2 +- scripts/ci/prek/check_deferrable_default.py | 37 +- 15 files changed, 1512 insertions(+), 52 deletions(-) diff --git a/generated/known_airflow_exceptions.txt b/generated/known_airflow_exceptions.txt index 1fb461854905d..d75fe492d6f83 100644 --- a/generated/known_airflow_exceptions.txt +++ b/generated/known_airflow_exceptions.txt @@ -170,9 +170,9 @@ providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/triggers/pod.py: providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/utils/pod_manager.py::2 providers/common/sql/src/airflow/providers/common/sql/hooks/sql.py::1 providers/common/sql/src/airflow/providers/common/sql/operators/generic_transfer.py::1 -providers/common/sql/src/airflow/providers/common/sql/operators/sql.py::7 +providers/common/sql/src/airflow/providers/common/sql/operators/sql.py::9 providers/common/sql/src/airflow/providers/common/sql/sensors/sql.py::6 -providers/common/sql/src/airflow/providers/common/sql/triggers/sql.py::1 +providers/common/sql/src/airflow/providers/common/sql/triggers/sql.py::2 providers/databricks/src/airflow/providers/databricks/hooks/databricks.py::8 providers/databricks/src/airflow/providers/databricks/hooks/databricks_base.py::44 providers/databricks/src/airflow/providers/databricks/hooks/databricks_sql.py::2 diff --git a/providers/common/sql/docs/operators.rst b/providers/common/sql/docs/operators.rst index 88e326a25fe5b..7afe6e055745d 100644 --- a/providers/common/sql/docs/operators.rst +++ b/providers/common/sql/docs/operators.rst @@ -35,6 +35,7 @@ different databases. Parameters of the operators are: - ``handler`` (optional) the function that will be applied to the cursor. If it's ``None`` results won't returned (default: fetch_all_handler). - ``split_statements`` (optional) if split single SQL string into statements and run separately (default: False). - ``return_last`` (optional) depends ``split_statements`` and if it's ``True`` this parameter is used to return the result of only last statement or all split statements (default: True). +- ``deferrable`` (optional) defaults to False and does not use the ``default_deferrable`` Airflow configuration. Currently all hooks used by this operator except the PostgresHook are incompatible with the deferrable mode. Using ``default_deferrable`` could result in failures for tasks using the incompatible hooks. The example below shows how to instantiate the SQLExecuteQueryOperator task. diff --git a/providers/common/sql/src/airflow/providers/common/sql/hooks/sql.py b/providers/common/sql/src/airflow/providers/common/sql/hooks/sql.py index 7a02f42c62ea1..6ed9749adca6c 100644 --- a/providers/common/sql/src/airflow/providers/common/sql/hooks/sql.py +++ b/providers/common/sql/src/airflow/providers/common/sql/hooks/sql.py @@ -16,22 +16,33 @@ # under the License. from __future__ import annotations +import asyncio import contextlib +import inspect import warnings -from collections.abc import Callable, Generator, Iterable, Mapping, MutableMapping, Sequence -from contextlib import closing, contextmanager, suppress +from collections.abc import ( + Awaitable, + Callable, + Generator, + Iterable, + Mapping, + MutableMapping, + Sequence, +) +from contextlib import asynccontextmanager, closing, contextmanager, suppress from datetime import datetime from functools import cached_property from typing import TYPE_CHECKING, Any, Literal, Protocol, TypeVar, cast, overload from urllib.parse import urlparse import sqlparse -from deprecated import deprecated +from asgiref.sync import sync_to_async +from deprecated import deprecated # type: ignore[import-untyped] from methodtools import lru_cache from more_itertools import chunked try: - from sqlalchemy import create_engine, inspect + from sqlalchemy import create_engine, inspect as sa_inspect from sqlalchemy.engine import make_url from sqlalchemy.exc import ArgumentError, NoSuchModuleError except ImportError: @@ -42,7 +53,7 @@ NoSuchModuleError = Exception # type: ignore[misc,assignment] -from airflow.exceptions import AirflowProviderDeprecationWarning +from airflow.exceptions import AirflowNotFoundException, AirflowProviderDeprecationWarning from airflow.providers.common.compat.module_loading import import_string from airflow.providers.common.compat.sdk import ( AirflowException, @@ -65,6 +76,8 @@ T = TypeVar("T") +HANDLER = Callable[[Any], Any | Awaitable[Any]] +ROW = tuple[Any, ...] SQL_PLACEHOLDERS = frozenset({"%s", "?"}) WARNING_MESSAGE = """Import of {} from the 'airflow.providers.common.sql.hooks' module is deprecated and will be removed in the future. Please import it from 'airflow.providers.common.sql.hooks.handlers'.""" @@ -343,7 +356,7 @@ def inspector(self) -> Inspector: "SQLAlchemy is required for database inspection. " "Install it with: pip install 'apache-airflow-providers-common-sql[sqlalchemy]'" ) - return inspect(self.get_sqlalchemy_engine()) + return sa_inspect(self.get_sqlalchemy_engine()) def get_table_schema(self, table_name: str, schema: str | None = None) -> list[dict[str, str]]: """ @@ -815,7 +828,8 @@ def run( with closing(conn.cursor()) as cur: results = [] for sql_statement in sql_list: - self._run_command(cur, sql_statement, parameters) + sql_to_run = self._translate_sql(sql_statement) + self._run_command(cur, sql_to_run, parameters) if handler is not None: result = self._make_common_data_structure(handler(cur)) @@ -1136,3 +1150,195 @@ 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 async_get_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 sync_to_async(self.get_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 _async_create_autocommit_connection(self, autocommit: bool = False): + conn = await self.async_get_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 _async_get_cursor(self, conn): + cursor = getattr(conn, "cursor", None) + if cursor is None: + raise AirflowNotFoundException( + f"{type(conn).__name__} has no cursor() method. " + "Override _async_get_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 _async_get_cursor in the DB-specific hook." + ) + + async def _async_run_command(self, cur, sql_statement, parameters): + """Run a statement using an already open cursor.""" + if self.log_sql: + self.log.info("Running statement: %s, parameters: %s", sql_statement, parameters) + prepared_params = await self._call(self._prepare_parameters, parameters) + if prepared_params: + await cur.execute(sql_statement, prepared_params) + else: + await cur.execute(sql_statement) + if cur.rowcount >= 0: + self.log.info("Rows affected: %s", cur.rowcount) + + @overload + async def run_async( + 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 run_async( + self, + sql: str | Iterable[str], + autocommit: bool = ..., + parameters: Iterable | Mapping[str, Any] | None = ..., + handler: HANDLER = ..., + split_statements: bool = ..., + return_last: bool = ..., + ) -> tuple | list | list[tuple] | list[list[tuple] | tuple] | None: ... + + async def run_async( + self, + sql: str | Iterable[str], + autocommit: bool = False, + parameters: Iterable | Mapping[str, Any] | None = None, + handler: HANDLER | None = None, + split_statements: bool = False, + return_last: bool = True, + ) -> tuple | list | list[tuple] | list[list[tuple] | tuple] | None: + self.descriptions = [] + if isinstance(sql, str): + if split_statements: + sql_list: Iterable[str] = self.split_sql_string(sql=sql, strip_semicolon=self.strip_semicolon) + else: + sql_list = [sql] if sql.strip() else [] + else: + sql_list = sql + if sql_list: + self.log.debug("Executing following statements against DB: %s", sql_list) + else: + raise ValueError("List of SQL statements is empty") + _last_result = None + _last_description = None + async with self._async_create_autocommit_connection(autocommit) as conn: + try: + async with self._async_get_cursor(conn) as cur: + results = [] + for sql_statement in sql_list: + sql_to_run = await self._call(self._translate_sql, sql_statement) + await self._async_run_command(cur, sql_to_run, parameters) + if handler is not None: + handled = await self._call(handler, cur) + if inspect.isawaitable(handled): + handled = await handled + result = self._make_common_data_structure(handled) + if handlers.return_single_query_results(sql, return_last, split_statements): + _last_result = result + _last_description = cur.description + else: + results.append(result) + self.descriptions.append(cur.description) + if not self.get_autocommit(conn): + await self._call(conn.commit) + except Exception: + rb = getattr(conn, "rollback", None) + if rb and not self.get_autocommit(conn): + await self._call(rb) + raise + finally: + await self._call(self.get_db_log_messages, conn) + if handler is None: + return None + if handlers.return_single_query_results(sql, return_last, split_statements): + self.descriptions = [_last_description] + return _last_result + return results diff --git a/providers/common/sql/src/airflow/providers/common/sql/hooks/sql.pyi b/providers/common/sql/src/airflow/providers/common/sql/hooks/sql.pyi index ebdcd58cf6608..59f74e06d149f 100644 --- a/providers/common/sql/src/airflow/providers/common/sql/hooks/sql.pyi +++ b/providers/common/sql/src/airflow/providers/common/sql/hooks/sql.pyi @@ -32,7 +32,7 @@ Definition of the public interface for airflow.providers.common.sql.src.airflow.providers.common.sql.hooks.sql. """ -from collections.abc import Callable, Generator, Iterable, Mapping, MutableMapping, Sequence +from collections.abc import Awaitable, Callable, Generator, Iterable, Mapping, MutableMapping, Sequence from functools import cached_property as cached_property from typing import Any, Literal, Protocol, TypeVar, overload @@ -48,6 +48,7 @@ from airflow.providers.openlineage.extractors import OperatorLineage as Operator from airflow.providers.openlineage.sqlparser import DatabaseInfo as DatabaseInfo T = TypeVar("T") +HANDLER = Callable[[Any], Any | Awaitable[Any]] SQL_PLACEHOLDERS: Incomplete WARNING_MESSAGE: str @@ -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 async_get_conn(self) -> Any: ... + @overload + async def run_async( + 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 run_async( + self, + sql: str | Iterable[str], + autocommit: bool = ..., + parameters: Iterable | Mapping[str, Any] | None = ..., + handler: HANDLER = ..., + split_statements: bool = ..., + return_last: bool = ..., + ) -> tuple | list | list[tuple] | list[list[tuple] | tuple] | None: ... @overload def run( self, diff --git a/providers/common/sql/src/airflow/providers/common/sql/operators/generic_transfer.py b/providers/common/sql/src/airflow/providers/common/sql/operators/generic_transfer.py index 3fd7b330c44b7..158997545b1e4 100644 --- a/providers/common/sql/src/airflow/providers/common/sql/operators/generic_transfer.py +++ b/providers/common/sql/src/airflow/providers/common/sql/operators/generic_transfer.py @@ -23,7 +23,7 @@ from airflow.providers.common.compat.sdk import AirflowException, BaseHook, BaseOperator, conf from airflow.providers.common.sql.hooks.sql import DbApiHook -from airflow.providers.common.sql.triggers.sql import SQLExecuteQueryTrigger +from airflow.providers.common.sql.triggers.sql import SQLGenericTransferTrigger if TYPE_CHECKING: import jinja2 @@ -165,7 +165,7 @@ def execute(self, context: Context): if self.page_size and isinstance(self.sql, str): if self.deferrable: self.defer( - trigger=SQLExecuteQueryTrigger( + trigger=SQLGenericTransferTrigger( conn_id=self.source_conn_id, hook_params=self.source_hook_params, sql=self.get_paginated_sql(0), @@ -230,7 +230,7 @@ def execute_complete( self._insert_rows(rows=rows, context=context) self.defer( - trigger=SQLExecuteQueryTrigger( + trigger=SQLGenericTransferTrigger( conn_id=self.source_conn_id, hook_params=self.source_hook_params, sql=self.get_paginated_sql(offset), diff --git a/providers/common/sql/src/airflow/providers/common/sql/operators/sql.py b/providers/common/sql/src/airflow/providers/common/sql/operators/sql.py index 21c8b6502b0bb..a627631d67537 100644 --- a/providers/common/sql/src/airflow/providers/common/sql/operators/sql.py +++ b/providers/common/sql/src/airflow/providers/common/sql/operators/sql.py @@ -18,6 +18,7 @@ from __future__ import annotations import ast +import inspect import re from collections.abc import Callable, Iterable, Mapping, Sequence from dataclasses import dataclass @@ -34,9 +35,11 @@ BaseOperator, SkipMixin, XComArg, + conf, ) from airflow.providers.common.sql.hooks.handlers import fetch_all_handler, return_single_query_results from airflow.providers.common.sql.hooks.sql import DbApiHook +from airflow.providers.common.sql.triggers.sql import SQLExecuteQueryTrigger from airflow.utils.helpers import merge_dicts if TYPE_CHECKING: @@ -456,6 +459,57 @@ def _attach_check_facets(self, operator_lineage: OperatorLineage) -> OperatorLin return operator_lineage +class _ReplayCursor: + """ + Read-only stand-in for a DBAPI cursor used by deferrable :class:`SQLExecuteQueryOperator` handlers. + + In deferrable mode the query runs in the triggerer; the rows it fetched (and the cursor + descriptions) are returned in the ``TriggerEvent`` and replayed here so the handler runs on the + worker, where it is defined, rather than in the triggerer's async event loop. Only the read side of + a cursor is available: live-connection features (server-side/streaming cursors, LOB locators, + ``.connection``, ``.scroll``, ``.nextset``, ``.copy_expert``, follow-up queries) are unsupported and + raise :class:`AttributeError` pointing at ``deferrable=False``. + """ + + def __init__( + self, + rows: list[Any] | None, + description: Sequence[Sequence] | None = None, + rowcount: int | None = None, + ) -> None: + self._rows = list(rows) if rows is not None else [] + self.description = description + self.rowcount = rowcount if rowcount is not None else len(self._rows) + self.arraysize = 1 + + def fetchall(self) -> list[Any]: + rows, self._rows = self._rows, [] + return rows + + def fetchone(self) -> Any | None: + return self._rows.pop(0) if self._rows else None + + def fetchmany(self, size: int | None = None) -> list[Any]: + size = self.arraysize if size is None else size + chunk, self._rows = self._rows[:size], self._rows[size:] + return chunk + + def __iter__(self): + while self._rows: + yield self._rows.pop(0) + + def close(self) -> None: + self._rows = [] + + def __getattr__(self, name: str) -> NoReturn: + raise AttributeError( + f"'{name}' is not available on a deferrable SQL handler's cursor. In deferrable mode the " + f"result set is replayed on the worker, so live-cursor features (a live connection, " + f"server-side/streaming cursors, LOB locators, .scroll, .nextset, .copy_expert, follow-up " + f"queries) are unsupported. Set deferrable=False to use a handler that needs a live cursor." + ) + + class SQLExecuteQueryOperator(BaseSQLOperator): """ Executes SQL code in a specific database. @@ -470,6 +524,13 @@ class SQLExecuteQueryOperator(BaseSQLOperator): :param autocommit: (optional) if True, each command is automatically committed (default: False). :param parameters: (optional) the parameters to render the SQL query with. :param handler: (optional) the function that will be applied to the cursor (default: fetch_all_handler). + In deferrable mode the query runs in the triggerer and the handler is applied on the worker over a + replayed, read-only cursor (``fetchall``/``fetchone``/``fetchmany``/``description``/``rowcount``). + Handlers that need a live cursor -- server-side/streaming cursors, LOB locators, + ``cursor.connection``, ``.scroll``, ``.nextset``, ``.copy_expert`` or follow-up queries on the same + connection -- are not supported with ``deferrable=True``; use ``deferrable=False`` for those. The + whole result set is materialized into the ``TriggerEvent``, so very large results are best fetched + with ``deferrable=False`` as well. :param output_processor: (optional) the function that will be applied to the result (default: default_output_processor). :param split_statements: (optional) if split single SQL string into statements. By default, defers @@ -482,6 +543,7 @@ class SQLExecuteQueryOperator(BaseSQLOperator): :param requires_result_fetch: (optional) if True, ensures that query results are fetched before completing execution. If `do_xcom_push` is True, results are fetched automatically, making this parameter redundant. (default: False). + :param deferrable: (optional) Run operator in the deferrable mode. .. seealso:: For more information on how to use this operator, take a look at the guide: @@ -499,7 +561,7 @@ def __init__( sql: str | list[str], autocommit: bool = False, parameters: Mapping | Iterable | None = None, - handler: Callable[[Any], list[tuple] | None] = fetch_all_handler, + handler: Callable[[Any], list[tuple] | None] | None = fetch_all_handler, output_processor: ( Callable[ [list[Any], list[Sequence[Sequence] | None]], @@ -513,6 +575,7 @@ def __init__( return_last: bool = True, show_return_value_in_logs: bool = False, requires_result_fetch: bool = False, + deferrable: bool = False, **kwargs, ) -> None: super().__init__(conn_id=conn_id, database=database, **kwargs) @@ -525,6 +588,7 @@ def __init__( self.return_last = return_last self.show_return_value_in_logs = show_return_value_in_logs self.requires_result_fetch = requires_result_fetch + self.deferrable = deferrable def _process_output( self, results: list[Any], descriptions: list[Sequence[Sequence] | None] @@ -552,31 +616,82 @@ def _process_output( def _should_run_output_processing(self) -> bool: return self.do_xcom_push + def _get_handler_import_path(self) -> str: + if not inspect.isfunction(self.handler): + raise ValueError("The handler must be a function object.") + module = getattr(self.handler, "__module__", None) + qualname = getattr(self.handler, "__qualname__", None) + if not module or not qualname: + raise ValueError("handler must have __module__ and __qualname__") + if "" in qualname or "" in qualname: + raise ValueError("The handler must not be a nested/local or lambda function.") + return f"{module}:{qualname}" + 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: + self.defer( + trigger=SQLExecuteQueryTrigger( + sql=self.sql, + conn_id=self.conn_id, + autocommit=self.autocommit, + parameters=self.parameters, + fetch_results=self._should_run_output_processing() or self.requires_result_fetch, + split_statements=self.split_statements, + return_last=self.return_last, + ), + method_name="execute_complete", + ) else: - extra_kwargs = {} - output = hook.run( - sql=self.sql, - autocommit=self.autocommit, - parameters=self.parameters, - handler=self.handler - if self._should_run_output_processing() or self.requires_result_fetch - else None, - return_last=self.return_last, - **extra_kwargs, - ) - if not self._should_run_output_processing(): + hook = self.get_db_hook() + if self.split_statements is not None: + extra_kwargs = {"split_statements": self.split_statements} + else: + extra_kwargs = {} + output = hook.run( + sql=self.sql, + autocommit=self.autocommit, + parameters=self.parameters, + handler=self.handler + if self._should_run_output_processing() or self.requires_result_fetch + else None, + return_last=self.return_last, + **extra_kwargs, + ) + if not self._should_run_output_processing(): + return None + if return_single_query_results(self.sql, self.return_last, self.split_statements): + # For simplicity, we pass always list as input to _process_output, regardless if + # single query results are going to be returned, and we return the first element + # of the list in this case from the (always) list returned by _process_output + return self._process_output([output], hook.descriptions)[-1] + result = self._process_output(output, hook.descriptions) + self.log.info("result: %s", result) + return result + + def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> Any: + if event is None: + raise AirflowException("Unknown error in SQLExecuteQueryTrigger") + if event.get("status") == "error": + raise AirflowException(event.get("message", "Unknown error in SQLExecuteQueryTrigger")) + self.log.info("SQL query executed successfully.") + results = event.get("results") + if not self._should_run_output_processing() or results is None: return None + descriptions: list[Sequence[Sequence] | None] = event.get("descriptions") or [] if return_single_query_results(self.sql, self.return_last, self.split_statements): - # For simplicity, we pass always list as input to _process_output, regardless if - # single query results are going to be returned, and we return the first element - # of the list in this case from the (always) list returned by _process_output - return self._process_output([output], hook.descriptions)[-1] - result = self._process_output(output, hook.descriptions) + # The trigger returns the rows of a single statement; apply the handler on the worker, where + # it is defined, over a replayed cursor rather than in the triggerer's event loop. + rows = [results] if isinstance(results, str) else results + description = descriptions[0] if descriptions else None + output = self.handler(_ReplayCursor(rows, description)) if self.handler else rows + result = self._process_output([output], [description])[-1] + else: + outputs = [] + for index, rows in enumerate(results): + description = descriptions[index] if index < len(descriptions) else None + outputs.append(self.handler(_ReplayCursor(rows, description)) if self.handler else rows) + result = self._process_output(outputs, descriptions or [None] * len(outputs)) self.log.info("result: %s", result) return result diff --git a/providers/common/sql/src/airflow/providers/common/sql/triggers/sql.py b/providers/common/sql/src/airflow/providers/common/sql/triggers/sql.py index f9d8980f332ad..78fddac31244c 100644 --- a/providers/common/sql/src/airflow/providers/common/sql/triggers/sql.py +++ b/providers/common/sql/src/airflow/providers/common/sql/triggers/sql.py @@ -17,10 +17,15 @@ # under the License. from __future__ import annotations +import importlib +import sys from typing import TYPE_CHECKING +from asgiref.sync import sync_to_async + from airflow.providers.common.compat.sdk import AirflowException, BaseHook 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 @@ -28,10 +33,16 @@ from collections.abc import AsyncIterator from typing import Any +from collections.abc import ( + Iterable, + Mapping, + Sequence, +) -class SQLExecuteQueryTrigger(BaseTrigger): + +class SQLGenericTransferTrigger(BaseTrigger): """ - A trigger that executes SQL code in async mode. + A SQL trigger that executes SQL to get records in async mode. :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 @@ -51,7 +62,7 @@ def __init__( self.hook_params = hook_params def serialize(self) -> tuple[str, dict[str, Any]]: - """Serialize the SQLExecuteQueryTrigger arguments and classpath.""" + """Serialize the SQLGenericTransferTrigger arguments and classpath.""" return ( f"{self.__class__.__module__}.{self.__class__.__name__}", { @@ -103,3 +114,142 @@ async def run(self) -> AsyncIterator[TriggerEvent]: except Exception as e: self.log.exception("An error occurred: %s", e) yield TriggerEvent({"status": "failure", "message": str(e)}) + + +class SQLExecuteQueryTrigger(BaseTrigger): + """ + 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 + """ + + def __init__( + self, + sql: str | Iterable[str], + conn_id: str, + autocommit: bool, + split_statements: bool, + return_last: bool, + parameters: Iterable[Any] | Mapping[str, Any] | None = None, + fetch_results: bool = False, + ): + super().__init__() + self.sql = sql + self.conn_id = conn_id + self.autocommit = autocommit + self.parameters = parameters + self.fetch_results = fetch_results + self.split_statements = split_statements + self.return_last = return_last + + def serialize(self) -> tuple[str, dict[str, Any]]: + """Serialize the SQLExecuteQueryTrigger arguments and classpath.""" + return ( + f"{self.__class__.__module__}.{self.__class__.__name__}", + { + "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, + }, + ) + + @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 get_hook(self) -> DbApiHook: + """ + Return DbApiHook. + + :return: DbApiHook for this connection + """ + connection = await sync_to_async(BaseHook.get_connection)(self.conn_id) + hook = await sync_to_async(connection.get_hook)() + if not isinstance(hook, DbApiHook) or not hasattr(hook, "run_async"): + raise AirflowException( + f"You are trying to use the SqlExecuteQueryOperator in deferrable mode with {hook.__class__.__name__}," + f" but its provider does not support this. Please set deferrable=False" + f" Got {hook.__class__.__name__} with class hierarchy: {hook.__class__.mro()}" + ) + return hook + + async def run(self) -> AsyncIterator[TriggerEvent]: + try: + hook = await self.get_hook() + + self.log.info("Extracting data from %s", self.conn_id) + self.log.info("Executing: \n %s", self.sql) + + 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 hook.run_async( + sql=self.sql, + autocommit=self.autocommit, + parameters=self.parameters, + handler=fetch_all_handler, + split_statements=self.split_statements, + return_last=self.return_last, + ) + + 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 hook.run_async( + sql=self.sql, + autocommit=self.autocommit, + parameters=self.parameters, + handler=None, + split_statements=self.split_statements, + return_last=self.return_last, + ) + + self.log.info("Executing query from %s done!", self.conn_id) + yield TriggerEvent({"status": "success"}) + + except Exception as e: + self.log.error("status: error, message: %s", str(e)) + yield TriggerEvent({"status": "error", "message": str(e)}) diff --git a/providers/common/sql/src/airflow/providers/common/sql/triggers/sql.pyi b/providers/common/sql/src/airflow/providers/common/sql/triggers/sql.pyi index 3a14c65998d20..850ec1cdb1247 100644 --- a/providers/common/sql/src/airflow/providers/common/sql/triggers/sql.pyi +++ b/providers/common/sql/src/airflow/providers/common/sql/triggers/sql.pyi @@ -35,7 +35,7 @@ from typing import Any from airflow.providers.common.sql.hooks.sql import DbApiHook as DbApiHook from airflow.triggers.base import BaseTrigger as BaseTrigger, TriggerEvent as TriggerEvent -class SQLExecuteQueryTrigger(BaseTrigger): +class SQLGenericTransferTrigger(BaseTrigger): def __init__( self, sql: str | list[str], conn_id: str, hook_params: dict | None = None, **kwargs ) -> None: ... diff --git a/providers/common/sql/tests/unit/common/sql/hooks/test_sql.py b/providers/common/sql/tests/unit/common/sql/hooks/test_sql.py index d88cc3abbc8cf..44e5d1f75f476 100644 --- a/providers/common/sql/tests/unit/common/sql/hooks/test_sql.py +++ b/providers/common/sql/tests/unit/common/sql/hooks/test_sql.py @@ -20,13 +20,14 @@ import inspect import logging -from unittest.mock import MagicMock, PropertyMock, patch +from contextlib import asynccontextmanager +from unittest.mock import AsyncMock, MagicMock, PropertyMock, patch import pandas as pd import polars as pl import pytest -from airflow.exceptions import AirflowProviderDeprecationWarning +from airflow.exceptions import AirflowNotFoundException, AirflowProviderDeprecationWarning from airflow.models import Connection from airflow.providers.common.sql.dialects.dialect import Dialect from airflow.providers.common.sql.hooks.handlers import fetch_all_handler @@ -64,7 +65,21 @@ def get_cursor_descriptions(fields: list[str]) -> list[tuple[str]]: return [(field,) for field in fields] -index = 0 +def _make_async_cursor(descriptions_per_call, fetchall_results): + """Cursor mock whose description advances on each async execute call.""" + cur = MagicMock() + cur.rowcount = 2 + cur.description = None + call_idx = 0 + + async def _execute(*args, **kwargs): + nonlocal call_idx + cur.description = descriptions_per_call[call_idx] + call_idx += 1 + + cur.execute = _execute + cur.fetchall.side_effect = fetchall_results + return cur @pytest.mark.db_test @@ -333,6 +348,551 @@ def test_get_df_with_df_type(db, df_type, expected_type): df = dbapi_hook.get_df("SQL", df_type=df_type) assert isinstance(df, expected_type) + @pytest.mark.db_test + def test_build_conn_kwargs_basic_mapping(self): + hook = mock_db_hook(DbApiHook) + db = Connection( + conn_id="c", + conn_type="test", + host="dbhost", + login="user", + password="secret", + schema="mydb", + port=5432, + ) + result = hook._build_conn_kwargs_from_airflow_connection(db) + assert result["host"] == "dbhost" + assert result["port"] == 5432 + assert result["username"] == "user" + assert result["password"] == "secret" + assert result["schema"] == "mydb" + assert result["raw_connection"] is db + + @pytest.mark.db_test + def test_build_conn_kwargs_none_values_become_empty(self): + hook = mock_db_hook(DbApiHook) + db = Connection(conn_id="c", conn_type="test") + result = hook._build_conn_kwargs_from_airflow_connection(db) + assert result["host"] == "" + assert result["username"] == "" + assert result["schema"] == "" + assert result["port"] is None + + @pytest.mark.db_test + def test_build_conn_kwargs_uses_extra_database(self): + hook = mock_db_hook(DbApiHook) + db = Connection(conn_id="c", conn_type="test", extra='{"database": "myschema"}') + assert hook._build_conn_kwargs_from_airflow_connection(db)["database"] == "myschema" + + @pytest.mark.db_test + def test_build_conn_kwargs_falls_back_to_dbname(self): + hook = mock_db_hook(DbApiHook) + db = Connection(conn_id="c", conn_type="test", extra='{"dbname": "altdb"}') + assert hook._build_conn_kwargs_from_airflow_connection(db)["database"] == "altdb" + + @pytest.mark.db_test + @pytest.mark.asyncio + async def test_async_get_conn_success(self): + hook = mock_db_hook(DbApiHook) + mock_db_conn = MagicMock() + hook.connector = AsyncMock() + hook.connector.connect.return_value = mock_db_conn + assert await hook.async_get_conn() is mock_db_conn + hook.connector.connect.assert_awaited_once() + + @pytest.mark.db_test + @pytest.mark.asyncio + async def test_async_get_conn_raises_without_connector(self): + hook = mock_db_hook(DbApiHook) + hook.connector = None + with pytest.raises(RuntimeError, match="didn't have `self.connector` set"): + await hook.async_get_conn() + + @pytest.mark.db_test + @pytest.mark.asyncio + async def test_async_get_conn_caches_airflow_connection(self): + hook = mock_db_hook(DbApiHook) + hook.connector = AsyncMock() + hook.connector.connect.return_value = MagicMock() + await hook.async_get_conn() + await hook.async_get_conn() + assert hook.connection_invocations == 1 + + @pytest.mark.db_test + @pytest.mark.asyncio + async def test_call_awaits_coroutine_function(self): + hook = mock_db_hook(DbApiHook) + + async def fn(x): + return x * 3 + + assert await hook._call(fn, 4) == 12 + + @pytest.mark.db_test + @pytest.mark.asyncio + async def test_call_runs_sync_function(self): + hook = mock_db_hook(DbApiHook) + + def fn(x): + return x + 1 + + assert await hook._call(fn, 9) == 10 + + @pytest.mark.db_test + @pytest.mark.asyncio + async def test_async_create_connection_yields_conn(self): + hook = mock_db_hook(DbApiHook) + mock_conn = MagicMock() + mock_conn.close = MagicMock() + with patch.object(hook, "async_get_conn", AsyncMock(return_value=mock_conn)): + async with hook._async_create_autocommit_connection() as conn: + assert conn is mock_conn + + @pytest.mark.db_test + @pytest.mark.asyncio + async def test_async_create_connection_calls_close(self): + hook = mock_db_hook(DbApiHook) + mock_conn = MagicMock() + mock_conn.close = MagicMock() + del mock_conn.aclose + with patch.object(hook, "async_get_conn", AsyncMock(return_value=mock_conn)): + async with hook._async_create_autocommit_connection(): + pass + mock_conn.close.assert_called_once() + + @pytest.mark.db_test + @pytest.mark.asyncio + async def test_async_create_connection_prefers_aclose(self): + hook = mock_db_hook(DbApiHook) + mock_conn = MagicMock() + mock_conn.aclose = AsyncMock() + mock_conn.close = MagicMock() + with patch.object(hook, "async_get_conn", AsyncMock(return_value=mock_conn)): + async with hook._async_create_autocommit_connection(): + pass + mock_conn.aclose.assert_awaited_once() + mock_conn.close.assert_not_called() + + @pytest.mark.db_test + @pytest.mark.asyncio + async def test_async_create_connection_sets_sync_autocommit(self): + hook = mock_db_hook(DbApiHook) + hook.supports_autocommit = True + mock_conn = MagicMock() + mock_conn.close = MagicMock() + del mock_conn.set_autocommit + with patch.object(hook, "async_get_conn", AsyncMock(return_value=mock_conn)): + async with hook._async_create_autocommit_connection(autocommit=True): + pass + assert mock_conn.autocommit is True + + @pytest.mark.db_test + @pytest.mark.asyncio + async def test_async_create_connection_sets_async_autocommit(self): + hook = mock_db_hook(DbApiHook) + hook.supports_autocommit = True + mock_conn = MagicMock() + mock_conn.close = MagicMock() + mock_conn.set_autocommit = AsyncMock() + with patch.object(hook, "async_get_conn", AsyncMock(return_value=mock_conn)): + async with hook._async_create_autocommit_connection(autocommit=True): + pass + mock_conn.set_autocommit.assert_awaited_once_with(True) + + @pytest.mark.db_test + @pytest.mark.asyncio + async def test_async_get_cursor_raises_without_cursor_method(self): + hook = mock_db_hook(DbApiHook) + with pytest.raises(AirflowNotFoundException, match="has no cursor\\(\\) method"): + async with hook._async_get_cursor(MagicMock(spec=[])): + pass + + @pytest.mark.db_test + @pytest.mark.asyncio + async def test_async_get_cursor_uses_async_context_manager(self): + hook = mock_db_hook(DbApiHook) + mock_cur = MagicMock() + cursor_cm = AsyncMock() + cursor_cm.__aenter__ = AsyncMock(return_value=mock_cur) + cursor_cm.__aexit__ = AsyncMock(return_value=False) + conn = MagicMock() + conn.cursor.return_value = cursor_cm + async with hook._async_get_cursor(conn) as cur: + assert cur is mock_cur + + @pytest.mark.db_test + @pytest.mark.asyncio + async def test_async_get_cursor_awaits_cursor_then_uses_async_cm(self): + hook = mock_db_hook(DbApiHook) + mock_cur = MagicMock() + cursor_cm = AsyncMock() + cursor_cm.__aenter__ = AsyncMock(return_value=mock_cur) + cursor_cm.__aexit__ = AsyncMock(return_value=False) + + async def async_cursor(): + return cursor_cm + + conn = MagicMock() + conn.cursor = async_cursor + async with hook._async_get_cursor(conn) as cur: + assert cur is mock_cur + + @pytest.mark.db_test + @pytest.mark.asyncio + async def test_async_get_cursor_yields_cursor_with_async_execute(self): + hook = mock_db_hook(DbApiHook) + mock_cur = MagicMock() + mock_cur.execute = AsyncMock() + mock_cur.close = MagicMock() + del mock_cur.__aenter__ + del mock_cur.__aexit__ + del mock_cur.aclose + conn = MagicMock() + conn.cursor.return_value = mock_cur + async with hook._async_get_cursor(conn) as cur: + assert cur is mock_cur + mock_cur.close.assert_called_once() + + @pytest.mark.db_test + @pytest.mark.asyncio + async def test_async_get_cursor_raises_for_unsupported_type(self): + hook = mock_db_hook(DbApiHook) + mock_cur = MagicMock() + mock_cur.execute = MagicMock() + del mock_cur.__aenter__ + del mock_cur.__aexit__ + conn = MagicMock() + conn.cursor.return_value = mock_cur + with pytest.raises(RuntimeError, match="Unsupported cursor type"): + async with hook._async_get_cursor(conn) as _: + pass + + @pytest.mark.db_test + @pytest.mark.asyncio + async def test_async_run_command_no_params(self): + hook = mock_db_hook(DbApiHook) + cur = AsyncMock() + cur.rowcount = 0 + await hook._async_run_command(cur, "SELECT 1", None) + cur.execute.assert_awaited_once_with("SELECT 1") + + @pytest.mark.db_test + @pytest.mark.asyncio + async def test_async_run_command_with_params(self): + hook = mock_db_hook(DbApiHook) + cur = AsyncMock() + cur.rowcount = 3 + await hook._async_run_command(cur, "SELECT * FROM t WHERE id=%s", (42,)) + cur.execute.assert_awaited_once_with("SELECT * FROM t WHERE id=%s", (42,)) + + @pytest.mark.db_test + @pytest.mark.asyncio + async def test_async_run_command_logs_rowcount(self): + hook = mock_db_hook(DbApiHook) + cur = AsyncMock() + cur.rowcount = 5 + mock_log = MagicMock() + with patch.object(hook, "_log", mock_log): + await hook._async_run_command(cur, "DELETE FROM t", None) + mock_log.info.assert_any_call("Rows affected: %s", 5) + + @pytest.mark.db_test + @pytest.mark.asyncio + async def test_async_run_command_skips_rowcount_when_negative(self): + hook = mock_db_hook(DbApiHook) + cur = AsyncMock() + cur.rowcount = -1 + mock_log = MagicMock() + with patch.object(hook, "_log", mock_log): + await hook._async_run_command(cur, "SELECT 1", None) + assert not any(call.args[0] == "Rows affected: %s" for call in mock_log.info.call_args_list) + + @pytest.mark.db_test + @pytest.mark.asyncio + @pytest.mark.parametrize("empty_sql", [[], "", "\n"]) + async def test_run_async_empty_sql_raises(self, empty_sql): + hook = mock_db_hook(DbApiHook) + mock_conn = MagicMock() + mock_conn.autocommit = False + + @asynccontextmanager + async def _conn_cm(autocommit=False): + yield mock_conn + + @asynccontextmanager + async def _cursor_cm(conn): + yield AsyncMock() + + with ( + patch.object(hook, "_async_create_autocommit_connection", _conn_cm), + patch.object(hook, "_async_get_cursor", _cursor_cm), + ): + with pytest.raises(ValueError, match="List of SQL statements is empty"): + await hook.run_async(sql=empty_sql) + + @pytest.mark.db_test + @pytest.mark.asyncio + async def test_run_async_no_handler_returns_none(self): + hook = mock_db_hook(DbApiHook) + cur = AsyncMock() + cur.rowcount = 0 + mock_conn = MagicMock() + mock_conn.autocommit = False + mock_conn.commit = MagicMock() + + @asynccontextmanager + async def _conn_cm(autocommit=False): + yield mock_conn + + @asynccontextmanager + async def _cursor_cm(conn): + yield cur + + with ( + patch.object(hook, "_async_create_autocommit_connection", _conn_cm), + patch.object(hook, "_async_get_cursor", _cursor_cm), + ): + assert await hook.run_async(sql="INSERT INTO t VALUES (1)") is None + + @pytest.mark.db_test + @pytest.mark.asyncio + @pytest.mark.parametrize( + ( + "return_last", + "split_statements", + "sql", + "cursor_descriptions", + "cursor_results", + "hook_descriptions", + "hook_results", + ), + [ + pytest.param( + True, + False, + "select * from test.test", + [["id", "value"]], + ([[1, 2], [11, 12]],), + [[("id",), ("value",)]], + [[1, 2], [11, 12]], + id="The return_last set and no split statements set on single query in string", + ), + pytest.param( + False, + False, + "select * from test.test;", + [["id", "value"]], + ([[1, 2], [11, 12]],), + [[("id",), ("value",)]], + [[1, 2], [11, 12]], + id="The return_last not set and no split statements set on single query in string", + ), + pytest.param( + True, + True, + "select * from test.test;", + [["id", "value"]], + ([[1, 2], [11, 12]],), + [[("id",), ("value",)]], + [[1, 2], [11, 12]], + id="The return_last set and split statements set on single query in string", + ), + pytest.param( + False, + True, + "select * from test.test;", + [["id", "value"]], + ([[1, 2], [11, 12]],), + [[("id",), ("value",)]], + [[[1, 2], [11, 12]]], + id="The return_last not set and split statements set on single query in string", + ), + pytest.param( + True, + True, + "select * from test.test;select * from test.test2;", + [["id", "value"], ["id2", "value2"]], + ([[1, 2], [11, 12]], [[3, 4], [13, 14]]), + [[("id2",), ("value2",)]], + [[3, 4], [13, 14]], + id="The return_last set and split statements set on multiple queries in string", + ), + pytest.param( + False, + True, + "select * from test.test;select * from test.test2;", + [["id", "value"], ["id2", "value2"]], + ([[1, 2], [11, 12]], [[3, 4], [13, 14]]), + [[("id",), ("value",)], [("id2",), ("value2",)]], + [[[1, 2], [11, 12]], [[3, 4], [13, 14]]], + id="The return_last not set and split statements set on multiple queries in string", + ), + pytest.param( + True, + True, + ["select * from test.test;"], + [["id", "value"]], + ([[1, 2], [11, 12]],), + [[("id",), ("value",)]], + [[[1, 2], [11, 12]]], + id="The return_last set on single query in list", + ), + pytest.param( + False, + True, + ["select * from test.test;"], + [["id", "value"]], + ([[1, 2], [11, 12]],), + [[("id",), ("value",)]], + [[[1, 2], [11, 12]]], + id="The return_last not set on single query in list", + ), + pytest.param( + True, + True, + "select * from test.test;select * from test.test2;", + [["id", "value"], ["id2", "value2"]], + ([[1, 2], [11, 12]], [[3, 4], [13, 14]]), + [[("id2",), ("value2",)]], + [[3, 4], [13, 14]], + id="The return_last set on multiple queries in list", + ), + pytest.param( + False, + True, + "select * from test.test;select * from test.test2;", + [["id", "value"], ["id2", "value2"]], + ([[1, 2], [11, 12]], [[3, 4], [13, 14]]), + [[("id",), ("value",)], [("id2",), ("value2",)]], + [[[1, 2], [11, 12]], [[3, 4], [13, 14]]], + id="The return_last not set on multiple queries not set", + ), + ], + ) + async def test_run_async_query( + self, + return_last, + split_statements, + sql, + cursor_descriptions, + cursor_results, + hook_descriptions, + hook_results, + ): + modified = [get_cursor_descriptions(d) for d in cursor_descriptions] + cur = _make_async_cursor(modified, cursor_results) + mock_conn = MagicMock() + mock_conn.autocommit = False + mock_conn.commit = MagicMock() + hook = mock_db_hook(DbApiHook) + + @asynccontextmanager + async def _conn_cm(autocommit=False): + yield mock_conn + + @asynccontextmanager + async def _cursor_cm(conn): + yield cur + + with ( + patch.object(hook, "_async_create_autocommit_connection", _conn_cm), + patch.object(hook, "_async_get_cursor", _cursor_cm), + ): + results = await hook.run_async( + sql=sql, + handler=fetch_all_handler, + return_last=return_last, + split_statements=split_statements, + ) + + assert hook.descriptions == hook_descriptions + assert hook.last_description == hook_descriptions[-1] + assert results == hook_results + + @pytest.mark.db_test + @pytest.mark.asyncio + async def test_run_async_commits_on_success(self): + hook = mock_db_hook(DbApiHook) + cur = AsyncMock() + cur.rowcount = 0 + mock_conn = MagicMock() + mock_conn.autocommit = False + mock_conn.commit = MagicMock() + + @asynccontextmanager + async def _conn_cm(autocommit=False): + yield mock_conn + + @asynccontextmanager + async def _cursor_cm(conn): + yield cur + + with ( + patch.object(hook, "_async_create_autocommit_connection", _conn_cm), + patch.object(hook, "_async_get_cursor", _cursor_cm), + ): + await hook.run_async(sql="INSERT INTO t VALUES (1)") + + mock_conn.commit.assert_called_once() + + @pytest.mark.db_test + @pytest.mark.asyncio + async def test_run_async_rollback_on_exception(self): + hook = mock_db_hook(DbApiHook) + cur = AsyncMock() + cur.execute.side_effect = RuntimeError("DB error") + mock_conn = MagicMock() + mock_conn.autocommit = False + mock_conn.rollback = MagicMock() + + @asynccontextmanager + async def _conn_cm(autocommit=False): + yield mock_conn + + @asynccontextmanager + async def _cursor_cm(conn): + yield cur + + with ( + patch.object(hook, "_async_create_autocommit_connection", _conn_cm), + patch.object(hook, "_async_get_cursor", _cursor_cm), + ): + with pytest.raises(RuntimeError, match="DB error"): + await hook.run_async(sql="SELECT fail") + + mock_conn.rollback.assert_called_once() + + @pytest.mark.db_test + @pytest.mark.asyncio + async def test_run_async_awaits_async_handler(self): + hook = mock_db_hook(DbApiHook) + rows = [(1, "a"), (2, "b")] + cur = AsyncMock() + cur.rowcount = 2 + cur.description = [("id",), ("name",)] + mock_conn = MagicMock() + mock_conn.autocommit = False + mock_conn.commit = MagicMock() + + async def async_handler(cursor): + return rows + + @asynccontextmanager + async def _conn_cm(autocommit=False): + yield mock_conn + + @asynccontextmanager + async def _cursor_cm(conn): + yield cur + + with ( + patch.object(hook, "_async_create_autocommit_connection", _conn_cm), + patch.object(hook, "_async_get_cursor", _cursor_cm), + ): + result = await hook.run_async(sql="SELECT 1", handler=async_handler) + + assert result == rows + class TestDbApiHookGetTableSchema: @pytest.mark.db_test diff --git a/providers/common/sql/tests/unit/common/sql/operators/test_sql_execute.py b/providers/common/sql/tests/unit/common/sql/operators/test_sql_execute.py index 564cfabd1a617..be56119bf29d1 100644 --- a/providers/common/sql/tests/unit/common/sql/operators/test_sql_execute.py +++ b/providers/common/sql/tests/unit/common/sql/operators/test_sql_execute.py @@ -24,6 +24,7 @@ import pytest +from airflow.exceptions import AirflowException, TaskDeferred from airflow.models import Connection from airflow.providers.common.compat.openlineage.facet import ( Dataset, @@ -34,6 +35,7 @@ from airflow.providers.common.sql.hooks.handlers import fetch_all_handler from airflow.providers.common.sql.hooks.sql import DbApiHook from airflow.providers.common.sql.operators.sql import SQLExecuteQueryOperator +from airflow.providers.common.sql.triggers.sql import SQLExecuteQueryTrigger from airflow.providers.openlineage.extractors.base import OperatorLineage DATE = "2017-04-20" @@ -380,3 +382,108 @@ def mock__import__(name, globals_=None, locals_=None, fromlist=(), level=0): op = SQLExecuteQueryOperator(task_id=TASK_ID, sql="SELECT 1;") assert op.get_openlineage_facets_on_start() is None assert op.get_openlineage_facets_on_complete(None) is None + + +class TestSQLExecuteQueryOperatorDeferrable: + def test_execute_defers(self): + op = SQLExecuteQueryOperator(task_id=TASK_ID, sql="SELECT 1", deferrable=True) + with pytest.raises(TaskDeferred) as exc: + op.execute({}) + assert isinstance(exc.value.trigger, SQLExecuteQueryTrigger) + assert exc.value.method_name == "execute_complete" + + def test_execute_complete_raises_on_none_event(self): + op = SQLExecuteQueryOperator(task_id=TASK_ID, sql="SELECT 1", deferrable=True) + with pytest.raises(AirflowException, match="Unknown error in SQLExecuteQueryTrigger"): + op.execute_complete(context={}, event=None) + + def test_execute_complete_raises_on_error_event(self): + op = SQLExecuteQueryOperator(task_id=TASK_ID, sql="SELECT 1", deferrable=True) + with pytest.raises(AirflowException, match="something went wrong"): + op.execute_complete(context={}, event={"status": "error", "message": "something went wrong"}) + + def test_execute_complete_returns_none_when_no_xcom_push(self): + op = SQLExecuteQueryOperator(task_id=TASK_ID, sql="SELECT 1", deferrable=True, do_xcom_push=False) + result = op.execute_complete(context={}, event={"status": "success", "results": [("row1",)]}) + assert result is None + + def test_execute_complete_returns_none_when_no_results(self): + op = SQLExecuteQueryOperator(task_id=TASK_ID, sql="SELECT 1", deferrable=True, do_xcom_push=True) + result = op.execute_complete(context={}, event={"status": "success", "results": None}) + assert result is None + + def test_execute_sets_fetch_results_from_output_processing(self): + op = SQLExecuteQueryOperator(task_id=TASK_ID, sql="SELECT 1", deferrable=True, do_xcom_push=True) + with pytest.raises(TaskDeferred) as exc: + op.execute({}) + assert exc.value.trigger.fetch_results is True + + def test_execute_sets_fetch_results_false_without_output_processing(self): + op = SQLExecuteQueryOperator(task_id=TASK_ID, sql="SELECT 1", deferrable=True, do_xcom_push=False) + with pytest.raises(TaskDeferred) as exc: + op.execute({}) + assert exc.value.trigger.fetch_results is False + + def test_execute_complete_applies_default_handler(self): + op = SQLExecuteQueryOperator(task_id=TASK_ID, sql="SELECT 1", deferrable=True, do_xcom_push=True) + result = op.execute_complete( + context={}, + event={ + "status": "success", + "results": [("a",), ("b",)], + "descriptions": [[["col", 23, None, None, None, None, None]]], + }, + ) + assert result == [("a",), ("b",)] + + def test_execute_complete_applies_custom_handler_on_worker(self): + # The handler runs on the worker over a replayed cursor and can read rows and descriptions. + def handler(cursor): + return {"columns": [column[0] for column in cursor.description], "rows": cursor.fetchall()} + + op = SQLExecuteQueryOperator( + task_id=TASK_ID, sql="SELECT 1", deferrable=True, do_xcom_push=True, handler=handler + ) + result = op.execute_complete( + context={}, + event={ + "status": "success", + "results": [("x",)], + "descriptions": [[["col", 23, None, None, None, None, None]]], + }, + ) + assert result == {"columns": ["col"], "rows": [("x",)]} + + def test_execute_complete_applies_handler_per_statement(self): + op = SQLExecuteQueryOperator( + task_id=TASK_ID, + sql=["SELECT 1", "SELECT 2"], + deferrable=True, + do_xcom_push=True, + split_statements=True, + return_last=False, + handler=fetch_all_handler, + ) + description = [["col", 23, None, None, None, None, None]] + result = op.execute_complete( + context={}, + event={ + "status": "success", + "results": [[("a",)], [("b",)]], + "descriptions": [description, description], + }, + ) + assert result == [[("a",)], [("b",)]] + + def test_execute_complete_replay_cursor_rejects_live_cursor_features(self): + def handler(cursor): + return cursor.connection + + op = SQLExecuteQueryOperator( + task_id=TASK_ID, sql="SELECT 1", deferrable=True, do_xcom_push=True, handler=handler + ) + with pytest.raises(AttributeError, match="deferrable=False"): + op.execute_complete( + context={}, + event={"status": "success", "results": [("a",)], "descriptions": [None]}, + ) diff --git a/providers/common/sql/tests/unit/common/sql/triggers/test_sql.py b/providers/common/sql/tests/unit/common/sql/triggers/test_sql.py index e09fba590be73..6d90b793be125 100644 --- a/providers/common/sql/tests/unit/common/sql/triggers/test_sql.py +++ b/providers/common/sql/tests/unit/common/sql/triggers/test_sql.py @@ -19,8 +19,9 @@ from unittest import mock from airflow.models.connection import Connection +from airflow.providers.common.sql.hooks.handlers import fetch_all_handler from airflow.providers.common.sql.hooks.sql import DbApiHook -from airflow.providers.common.sql.triggers.sql import SQLExecuteQueryTrigger +from airflow.providers.common.sql.triggers.sql import SQLExecuteQueryTrigger, SQLGenericTransferTrigger from airflow.triggers.base import TriggerEvent try: @@ -35,7 +36,7 @@ from tests_common.test_utils.operators.run_deferrable import run_trigger -class TestSQLExecuteQueryTrigger: +class TestSQLGenericTransferTrigger: @mock.patch(f"{BASEHOOK_PATCH_PATH}.get_connection") def test_run(self, mock_get_connection): data = [(1, "Alice"), (2, "Bob")] @@ -45,10 +46,110 @@ def test_run(self, mock_get_connection): mock_get_connection.return_value = mock_connection mock_connection.get_hook.side_effect = lambda hook_params: mock_hook - trigger = SQLExecuteQueryTrigger(sql="SELECT * FROM users;", conn_id="test_conn_id") + trigger = SQLGenericTransferTrigger(sql="SELECT * FROM users;", conn_id="test_conn_id") actual = run_trigger(trigger) assert len(actual) == 1 assert isinstance(actual[0], TriggerEvent) assert actual[0].payload["status"] == "success" assert actual[0].payload["results"] == data + + +class TestSQLExecuteQueryTrigger: + def _make_trigger(self, **kwargs): + defaults = { + "sql": "SELECT 1", + "conn_id": "test_conn", + "autocommit": False, + "split_statements": False, + "return_last": True, + } + defaults.update(kwargs) + return SQLExecuteQueryTrigger(**defaults) + + def _make_mock_hook(self): + mock_hook = mock.MagicMock(spec=DbApiHook) + mock_hook.run_async = mock.AsyncMock() + return mock_hook + + def test_serialize(self): + trigger = self._make_trigger( + autocommit=True, + parameters={"p": 1}, + fetch_results=True, + split_statements=True, + return_last=False, + ) + classpath, kwargs = trigger.serialize() + assert classpath == "airflow.providers.common.sql.triggers.sql.SQLExecuteQueryTrigger" + assert kwargs == { + "sql": "SELECT 1", + "conn_id": "test_conn", + "autocommit": True, + "parameters": {"p": 1}, + "fetch_results": True, + "split_statements": True, + "return_last": False, + } + + def test_run_fetch_results_returns_rows_and_descriptions(self): + rows = [("val1",), ("val2",)] + descriptions = [(("col1", 23, None, None, None, None, None),)] + mock_hook = self._make_mock_hook() + mock_hook.run_async.return_value = rows + mock_hook.descriptions = descriptions + trigger = self._make_trigger(fetch_results=True) + with mock.patch.object(trigger, "get_hook", new=mock.AsyncMock(return_value=mock_hook)): + events = run_trigger(trigger) + + # The built-in fetch handler is used; no user handler runs in the triggerer. + assert mock_hook.run_async.await_args.kwargs["handler"] is fetch_all_handler + assert len(events) == 1 + assert events[0].payload == { + "status": "success", + "results": rows, + "descriptions": [[["col1", 23, None, None, None, None, None]]], + } + + def test_run_without_fetch_results_returns_no_results(self): + mock_hook = self._make_mock_hook() + trigger = self._make_trigger(fetch_results=False) + with mock.patch.object(trigger, "get_hook", new=mock.AsyncMock(return_value=mock_hook)): + events = run_trigger(trigger) + + assert mock_hook.run_async.await_args.kwargs["handler"] is None + assert len(events) == 1 + assert events[0].payload == {"status": "success"} + + def test_jsonsafe_descriptions_stringifies_non_native_type_codes(self): + class _Type: + def __str__(self): + return "CUSTOM_TYPE" + + descriptions = [(("col", _Type(), None, None, None, None, None),), None] + assert SQLExecuteQueryTrigger._jsonsafe_descriptions(descriptions) == [ + [["col", "CUSTOM_TYPE", None, None, None, None, None]], + None, + ] + + def test_run_yields_error_event_on_exception(self): + mock_hook = self._make_mock_hook() + mock_hook.run_async.side_effect = Exception("DB error") + trigger = self._make_trigger(fetch_results=True) + with mock.patch.object(trigger, "get_hook", new=mock.AsyncMock(return_value=mock_hook)): + events = run_trigger(trigger) + + assert len(events) == 1 + assert events[0].payload["status"] == "error" + assert "DB error" in events[0].payload["message"] + + @mock.patch(f"{BASEHOOK_PATCH_PATH}.get_connection") + def test_get_hook_raises_when_hook_is_not_dbapihook(self, mock_get_connection): + mock_connection = mock.MagicMock(spec=Connection) + mock_connection.get_hook.return_value = mock.MagicMock() # not a DbApiHook + mock_get_connection.return_value = mock_connection + trigger = self._make_trigger() + events = run_trigger(trigger) + + assert len(events) == 1 + assert events[0].payload["status"] == "error" diff --git a/providers/postgres/src/airflow/providers/postgres/hooks/postgres.py b/providers/postgres/src/airflow/providers/postgres/hooks/postgres.py index 77ca0537b934d..234e3ff2550e1 100644 --- a/providers/postgres/src/airflow/providers/postgres/hooks/postgres.py +++ b/providers/postgres/src/airflow/providers/postgres/hooks/postgres.py @@ -259,8 +259,8 @@ def _generate_cursor_name(self): return f"airflow_cursor_{uuid.uuid4().hex}" - def get_conn(self) -> CompatConnection: - """Establish a connection to a postgres database.""" + def _build_conn_args(self) -> tuple[dict[str, Any], Any]: + """Build base connection arguments from the Airflow connection.""" conn = deepcopy(self.connection) if conn.extra_dejson.get("iam", False): @@ -285,16 +285,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 async_get_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() + + raw_cursor = conn.extra_dejson.get("cursor") + if raw_cursor: + conn_args["row_factory"] = self._get_cursor(raw_cursor) + + # Use Any type for the connection args to avoid type conflicts + connection = await AsyncConnection.connect(**cast("Any", conn_args)) + + # Register JSON handlers for both json and jsonb types + # This ensures JSON data is properly decoded from bytes to Python objects + register_default_adapters(connection) + + # Add the notice handler AFTER the connection is established + if self.enable_log_db_messages and hasattr(connection, "add_notice_handler"): + connection.add_notice_handler(self._notice_handler) + + return connection + @overload def get_df( self, diff --git a/providers/postgres/tests/unit/postgres/hooks/test_postgres.py b/providers/postgres/tests/unit/postgres/hooks/test_postgres.py index e8c1dd5eff9d0..1bac52926e303 100644 --- a/providers/postgres/tests/unit/postgres/hooks/test_postgres.py +++ b/providers/postgres/tests/unit/postgres/hooks/test_postgres.py @@ -611,6 +611,143 @@ def test_get_conn_cursor(self, mocker): port=None, ) + @pytest.mark.asyncio + async def test_async_get_conn_raises_not_implemented_without_psycopg3(self, mocker): + mocker.patch("airflow.providers.postgres.hooks.postgres.USE_PSYCOPG3", False) + with pytest.raises(NotImplementedError, match="Async connections for PostgresHook require psycopg3"): + await self.db_hook.async_get_conn() + + @pytest.mark.asyncio + async def test_async_get_conn(self, mocker): + mock_async_connect = mocker.patch("psycopg.AsyncConnection.connect", new_callable=mock.AsyncMock) + mock_register = mocker.patch("airflow.providers.postgres.hooks.postgres.register_default_adapters") + mock_connection = mocker.MagicMock() + mock_async_connect.return_value = mock_connection + + result = await self.db_hook.async_get_conn() + + mock_async_connect.assert_awaited_once_with( + user="login", + password="password", + host="host", + dbname="database", + port=None, + ) + mock_register.assert_called_once_with(mock_connection) + assert result is mock_connection + + @pytest.mark.asyncio + async def test_async_get_conn_with_cursor(self, mocker): + self.connection.extra = '{"cursor": "dictcursor"}' + mock_async_connect = mocker.patch("psycopg.AsyncConnection.connect", new_callable=mock.AsyncMock) + mocker.patch("airflow.providers.postgres.hooks.postgres.register_default_adapters") + + await self.db_hook.async_get_conn() + + mock_async_connect.assert_awaited_once_with( + row_factory=psycopg.rows.dict_row, + user="login", + password="password", + host="host", + dbname="database", + port=None, + ) + + @pytest.mark.asyncio + async def test_async_get_conn_with_invalid_cursor(self, mocker): + self.connection.extra = '{"cursor": "mycursor"}' + + with pytest.raises(ValueError, match="Invalid cursor passed mycursor"): + await self.db_hook.async_get_conn() + + @pytest.mark.asyncio + async def test_async_get_conn_adds_notice_handler_when_enabled(self, mocker): + mock_async_connect = mocker.patch("psycopg.AsyncConnection.connect", new_callable=mock.AsyncMock) + mocker.patch("airflow.providers.postgres.hooks.postgres.register_default_adapters") + mock_connection = mocker.MagicMock() + mock_async_connect.return_value = mock_connection + + hook = PostgresHook(enable_log_db_messages=True) + hook.get_connection = mock.Mock(return_value=self.connection) + + await hook.async_get_conn() + + mock_connection.add_notice_handler.assert_called_once_with(hook._notice_handler) + + @pytest.mark.asyncio + async def test_async_get_conn_no_notice_handler_when_disabled(self, mocker): + mock_async_connect = mocker.patch("psycopg.AsyncConnection.connect", new_callable=mock.AsyncMock) + mocker.patch("airflow.providers.postgres.hooks.postgres.register_default_adapters") + mock_connection = mocker.MagicMock() + mock_async_connect.return_value = mock_connection + + await self.db_hook.async_get_conn() + + mock_connection.add_notice_handler.assert_not_called() + + @pytest.mark.asyncio + async def test_async_get_conn_with_namedtuple_cursor(self, mocker): + self.connection.extra = '{"cursor": "namedtuplecursor"}' + mock_async_connect = mocker.patch("psycopg.AsyncConnection.connect", new_callable=mock.AsyncMock) + mocker.patch("airflow.providers.postgres.hooks.postgres.register_default_adapters") + + await self.db_hook.async_get_conn() + + mock_async_connect.assert_awaited_once_with( + row_factory=psycopg.rows.namedtuple_row, + user="login", + password="password", + host="host", + dbname="database", + port=None, + ) + + @pytest.mark.asyncio + async def test_async_get_conn_with_realdictcursor_raises(self, mocker): + self.connection.extra = '{"cursor": "realdictcursor"}' + mocker.patch("psycopg.AsyncConnection.connect", new_callable=mock.AsyncMock) + + from airflow.providers.common.compat.sdk import AirflowException + + with pytest.raises(AirflowException, match="realdictcursor is not supported with psycopg3"): + await self.db_hook.async_get_conn() + + @pytest.mark.asyncio + async def test_async_get_conn_with_extra(self, mocker): + self.connection.extra = '{"connect_timeout": 3}' + mock_async_connect = mocker.patch("psycopg.AsyncConnection.connect", new_callable=mock.AsyncMock) + mocker.patch("airflow.providers.postgres.hooks.postgres.register_default_adapters") + + await self.db_hook.async_get_conn() + + mock_async_connect.assert_awaited_once_with( + user="login", + password="password", + host="host", + dbname="database", + port=None, + connect_timeout=3, + ) + + @pytest.mark.asyncio + async def test_async_get_conn_with_options(self, mocker): + mock_async_connect = mocker.patch("psycopg.AsyncConnection.connect", new_callable=mock.AsyncMock) + mocker.patch("airflow.providers.postgres.hooks.postgres.register_default_adapters") + + hook = PostgresHook(options="-c statement_timeout=3000ms") + hook.get_connection = mock.Mock(return_value=self.connection) + + await hook.async_get_conn() + + mock_async_connect.assert_awaited_once_with( + user="login", + password="password", + host="host", + dbname="database", + port=None, + options="-c statement_timeout=3000ms", + ) + @pytest.mark.backend("postgres") class TestPostgresHook: @@ -913,7 +1050,6 @@ def test_insert_rows_replace_all_index(self): @mock.patch("airflow.providers.postgres.hooks.postgres.PostgresHook.insert_rows") def test_upsert_rows(self, mock_insert_rows): - rows = [(1, "hello")] table = "table" diff --git a/providers/snowflake/src/airflow/providers/snowflake/operators/snowflake.py b/providers/snowflake/src/airflow/providers/snowflake/operators/snowflake.py index eec2e39ecc2c6..4184b27b1f6c3 100644 --- a/providers/snowflake/src/airflow/providers/snowflake/operators/snowflake.py +++ b/providers/snowflake/src/airflow/providers/snowflake/operators/snowflake.py @@ -409,7 +409,7 @@ def __init__( "session_parameters": session_parameters, **hook_params, } - super().__init__(conn_id=snowflake_conn_id, **kwargs) # pragma: no cover + super().__init__(conn_id=snowflake_conn_id, deferrable=deferrable, **kwargs) # pragma: no cover @cached_property def _hook(self): diff --git a/scripts/ci/prek/check_deferrable_default.py b/scripts/ci/prek/check_deferrable_default.py index 8a9252d76441f..00b0ba9f42814 100755 --- a/scripts/ci/prek/check_deferrable_default.py +++ b/scripts/ci/prek/check_deferrable_default.py @@ -41,13 +41,27 @@ "authoring-and-scheduling/deferring.rst#writing-deferrable-operators" ) +# Operators allowed to default ``deferrable`` to a plain ``False`` instead of reading the +# ``operators/default_deferrable`` config. SQLExecuteQueryOperator opts out because its deferrable +# mode runs the query in the triggerer behind a replay cursor that drops live-cursor semantics, so +# silently honouring a global ``default_deferrable=True`` would change behaviour for existing users. +DEFERRABLE_DEFAULT_EXEMPT_CLASSES = frozenset({"SQLExecuteQueryOperator"}) + class DefaultDeferrableVisitor(ast.NodeVisitor): def __init__(self, *args, **kwargs) -> None: super().__init__(*args, *kwargs) self.error_linenos: list[int] = [] + self._class_stack: list[str] = [] + + def visit_ClassDef(self, node: ast.ClassDef) -> None: + self._class_stack.append(node.name) + self.generic_visit(node) + self._class_stack.pop() def visit_FunctionDef(self, node: ast.FunctionDef) -> ast.FunctionDef: + if self._class_stack and self._class_stack[-1] in DEFERRABLE_DEFAULT_EXEMPT_CLASSES: + return node if node.name == "__init__": args = node.args arguments = reversed([*args.args, *args.posonlyargs, *args.kwonlyargs]) @@ -69,7 +83,21 @@ def visit_FunctionDef(self, node: ast.FunctionDef) -> ast.FunctionDef: class DefaultDeferrableTransformer(cst.CSTTransformer): + def __init__(self) -> None: + super().__init__() + self._class_stack: list[str] = [] + + def visit_ClassDef(self, node: cst.ClassDef) -> bool: + self._class_stack.append(node.name.value) + return True + + def leave_ClassDef(self, original_node: cst.ClassDef, updated_node: cst.ClassDef) -> cst.ClassDef: + self._class_stack.pop() + return updated_node + def leave_Param(self, original_node: cst.Param, updated_node: cst.Param) -> cst.Param: + if self._class_stack and self._class_stack[-1] in DEFERRABLE_DEFAULT_EXEMPT_CLASSES: + return updated_node if original_node.name.value == "deferrable": expected_default_cst = cst.parse_expression( 'conf.getboolean("operators", "default_deferrable", fallback=False)' @@ -97,9 +125,16 @@ def iter_check_deferrable_default_errors(module_filename: str) -> Iterator[str]: yield from (f"{module_filename}:{lineno}" for lineno in visitor.error_linenos) +def _conf_import_module(module_filename: str) -> str: + """Provider files must import ``conf`` from the compat SDK; everything else from core.""" + if f"{os.sep}providers{os.sep}" in os.path.abspath(module_filename): + return "airflow.providers.common.compat.sdk" + return "airflow.configuration" + + def _fix_invalid_deferrable_default_value(module_filename: str) -> None: context = CodemodContext(filename=module_filename) - AddImportsVisitor.add_needed_import(context, "airflow.configuration", "conf") + AddImportsVisitor.add_needed_import(context, _conf_import_module(module_filename), "conf") transformer = DefaultDeferrableTransformer() source_cst_tree = cst.parse_module(open(module_filename).read()) From d055fae9a11ef62c8f98483b7e9aec5446eeed9e Mon Sep 17 00:00:00 2001 From: Karen Braganza Date: Tue, 30 Jun 2026 17:10:50 -0400 Subject: [PATCH 2/9] Remove dead code for importing handler --- .../src/airflow/providers/common/sql/hooks/sql.py | 4 +--- .../airflow/providers/common/sql/operators/sql.py | 13 ------------- .../airflow/providers/common/sql/triggers/sql.py | 2 -- 3 files changed, 1 insertion(+), 18 deletions(-) diff --git a/providers/common/sql/src/airflow/providers/common/sql/hooks/sql.py b/providers/common/sql/src/airflow/providers/common/sql/hooks/sql.py index 6ed9749adca6c..91bb626f4b09b 100644 --- a/providers/common/sql/src/airflow/providers/common/sql/hooks/sql.py +++ b/providers/common/sql/src/airflow/providers/common/sql/hooks/sql.py @@ -77,7 +77,6 @@ T = TypeVar("T") HANDLER = Callable[[Any], Any | Awaitable[Any]] -ROW = tuple[Any, ...] SQL_PLACEHOLDERS = frozenset({"%s", "?"}) WARNING_MESSAGE = """Import of {} from the 'airflow.providers.common.sql.hooks' module is deprecated and will be removed in the future. Please import it from 'airflow.providers.common.sql.hooks.handlers'.""" @@ -1158,7 +1157,6 @@ def _translate_sql(self, sql: str) -> str: 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): @@ -1317,7 +1315,7 @@ async def run_async( sql_to_run = await self._call(self._translate_sql, sql_statement) await self._async_run_command(cur, sql_to_run, parameters) if handler is not None: - handled = await self._call(handler, cur) + handled = handler(cur) if inspect.isawaitable(handled): handled = await handled result = self._make_common_data_structure(handled) diff --git a/providers/common/sql/src/airflow/providers/common/sql/operators/sql.py b/providers/common/sql/src/airflow/providers/common/sql/operators/sql.py index a627631d67537..3588616f34482 100644 --- a/providers/common/sql/src/airflow/providers/common/sql/operators/sql.py +++ b/providers/common/sql/src/airflow/providers/common/sql/operators/sql.py @@ -18,7 +18,6 @@ from __future__ import annotations import ast -import inspect import re from collections.abc import Callable, Iterable, Mapping, Sequence from dataclasses import dataclass @@ -35,7 +34,6 @@ BaseOperator, SkipMixin, XComArg, - conf, ) from airflow.providers.common.sql.hooks.handlers import fetch_all_handler, return_single_query_results from airflow.providers.common.sql.hooks.sql import DbApiHook @@ -616,17 +614,6 @@ def _process_output( def _should_run_output_processing(self) -> bool: return self.do_xcom_push - def _get_handler_import_path(self) -> str: - if not inspect.isfunction(self.handler): - raise ValueError("The handler must be a function object.") - module = getattr(self.handler, "__module__", None) - qualname = getattr(self.handler, "__qualname__", None) - if not module or not qualname: - raise ValueError("handler must have __module__ and __qualname__") - if "" in qualname or "" in qualname: - raise ValueError("The handler must not be a nested/local or lambda function.") - return f"{module}:{qualname}" - def execute(self, context): self.log.info("Executing: %s", self.sql) if self.deferrable: diff --git a/providers/common/sql/src/airflow/providers/common/sql/triggers/sql.py b/providers/common/sql/src/airflow/providers/common/sql/triggers/sql.py index 78fddac31244c..15eea107c93ac 100644 --- a/providers/common/sql/src/airflow/providers/common/sql/triggers/sql.py +++ b/providers/common/sql/src/airflow/providers/common/sql/triggers/sql.py @@ -17,8 +17,6 @@ # under the License. from __future__ import annotations -import importlib -import sys from typing import TYPE_CHECKING from asgiref.sync import sync_to_async From a2b5fe6cded1e81c9f9d5e642306c369e7d676aa Mon Sep 17 00:00:00 2001 From: Karen Braganza Date: Tue, 7 Jul 2026 14:20:23 -0400 Subject: [PATCH 3/9] Rename run_async to arun --- .../sql/src/airflow/providers/common/sql/hooks/sql.py | 6 +++--- .../sql/src/airflow/providers/common/sql/hooks/sql.pyi | 4 ++-- .../src/airflow/providers/common/sql/triggers/sql.py | 6 +++--- .../sql/tests/unit/common/sql/triggers/test_sql.py | 10 +++++----- .../src/airflow/providers/postgres/hooks/postgres.py | 2 +- 5 files changed, 14 insertions(+), 14 deletions(-) diff --git a/providers/common/sql/src/airflow/providers/common/sql/hooks/sql.py b/providers/common/sql/src/airflow/providers/common/sql/hooks/sql.py index 91bb626f4b09b..7810c8cf9544e 100644 --- a/providers/common/sql/src/airflow/providers/common/sql/hooks/sql.py +++ b/providers/common/sql/src/airflow/providers/common/sql/hooks/sql.py @@ -1263,7 +1263,7 @@ async def _async_run_command(self, cur, sql_statement, parameters): self.log.info("Rows affected: %s", cur.rowcount) @overload - async def run_async( + async def arun( self, sql: str | Iterable[str], autocommit: bool = ..., @@ -1274,7 +1274,7 @@ async def run_async( ) -> None: ... @overload - async def run_async( + async def arun( self, sql: str | Iterable[str], autocommit: bool = ..., @@ -1284,7 +1284,7 @@ async def run_async( return_last: bool = ..., ) -> tuple | list | list[tuple] | list[list[tuple] | tuple] | None: ... - async def run_async( + async def arun( self, sql: str | Iterable[str], autocommit: bool = False, diff --git a/providers/common/sql/src/airflow/providers/common/sql/hooks/sql.pyi b/providers/common/sql/src/airflow/providers/common/sql/hooks/sql.pyi index 59f74e06d149f..991ee195c86dd 100644 --- a/providers/common/sql/src/airflow/providers/common/sql/hooks/sql.pyi +++ b/providers/common/sql/src/airflow/providers/common/sql/hooks/sql.pyi @@ -188,7 +188,7 @@ class DbApiHook(BaseHook): def get_db_log_messages(self, conn) -> None: ... async def async_get_conn(self) -> Any: ... @overload - async def run_async( + async def arun( self, sql: str | Iterable[str], autocommit: bool = ..., @@ -198,7 +198,7 @@ class DbApiHook(BaseHook): return_last: bool = ..., ) -> None: ... @overload - async def run_async( + async def arun( self, sql: str | Iterable[str], autocommit: bool = ..., diff --git a/providers/common/sql/src/airflow/providers/common/sql/triggers/sql.py b/providers/common/sql/src/airflow/providers/common/sql/triggers/sql.py index 15eea107c93ac..a05b6985b27b2 100644 --- a/providers/common/sql/src/airflow/providers/common/sql/triggers/sql.py +++ b/providers/common/sql/src/airflow/providers/common/sql/triggers/sql.py @@ -198,7 +198,7 @@ async def get_hook(self) -> DbApiHook: """ connection = await sync_to_async(BaseHook.get_connection)(self.conn_id) hook = await sync_to_async(connection.get_hook)() - if not isinstance(hook, DbApiHook) or not hasattr(hook, "run_async"): + if not isinstance(hook, DbApiHook) or not hasattr(hook, "arun"): raise AirflowException( f"You are trying to use the SqlExecuteQueryOperator in deferrable mode with {hook.__class__.__name__}," f" but its provider does not support this. Please set deferrable=False" @@ -216,7 +216,7 @@ async def run(self) -> AsyncIterator[TriggerEvent]: 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 hook.run_async( + results = await hook.arun( sql=self.sql, autocommit=self.autocommit, parameters=self.parameters, @@ -236,7 +236,7 @@ async def run(self) -> AsyncIterator[TriggerEvent]: ) else: - await hook.run_async( + await hook.arun( sql=self.sql, autocommit=self.autocommit, parameters=self.parameters, diff --git a/providers/common/sql/tests/unit/common/sql/triggers/test_sql.py b/providers/common/sql/tests/unit/common/sql/triggers/test_sql.py index 6d90b793be125..823d6c46ae150 100644 --- a/providers/common/sql/tests/unit/common/sql/triggers/test_sql.py +++ b/providers/common/sql/tests/unit/common/sql/triggers/test_sql.py @@ -69,7 +69,7 @@ def _make_trigger(self, **kwargs): def _make_mock_hook(self): mock_hook = mock.MagicMock(spec=DbApiHook) - mock_hook.run_async = mock.AsyncMock() + mock_hook.arun = mock.AsyncMock() return mock_hook def test_serialize(self): @@ -96,14 +96,14 @@ def test_run_fetch_results_returns_rows_and_descriptions(self): rows = [("val1",), ("val2",)] descriptions = [(("col1", 23, None, None, None, None, None),)] mock_hook = self._make_mock_hook() - mock_hook.run_async.return_value = rows + mock_hook.arun.return_value = rows mock_hook.descriptions = descriptions trigger = self._make_trigger(fetch_results=True) with mock.patch.object(trigger, "get_hook", new=mock.AsyncMock(return_value=mock_hook)): events = run_trigger(trigger) # The built-in fetch handler is used; no user handler runs in the triggerer. - assert mock_hook.run_async.await_args.kwargs["handler"] is fetch_all_handler + assert mock_hook.arun.await_args.kwargs["handler"] is fetch_all_handler assert len(events) == 1 assert events[0].payload == { "status": "success", @@ -117,7 +117,7 @@ def test_run_without_fetch_results_returns_no_results(self): with mock.patch.object(trigger, "get_hook", new=mock.AsyncMock(return_value=mock_hook)): events = run_trigger(trigger) - assert mock_hook.run_async.await_args.kwargs["handler"] is None + assert mock_hook.arun.await_args.kwargs["handler"] is None assert len(events) == 1 assert events[0].payload == {"status": "success"} @@ -134,7 +134,7 @@ def __str__(self): def test_run_yields_error_event_on_exception(self): mock_hook = self._make_mock_hook() - mock_hook.run_async.side_effect = Exception("DB error") + mock_hook.arun.side_effect = Exception("DB error") trigger = self._make_trigger(fetch_results=True) with mock.patch.object(trigger, "get_hook", new=mock.AsyncMock(return_value=mock_hook)): events = run_trigger(trigger) diff --git a/providers/postgres/src/airflow/providers/postgres/hooks/postgres.py b/providers/postgres/src/airflow/providers/postgres/hooks/postgres.py index 234e3ff2550e1..fca205d46caa5 100644 --- a/providers/postgres/src/airflow/providers/postgres/hooks/postgres.py +++ b/providers/postgres/src/airflow/providers/postgres/hooks/postgres.py @@ -127,7 +127,7 @@ class PostgresHook(DbApiHook): reference to a specific postgres database. :param options: Optional. Specifies command-line options to send to the server at connection start. For example, setting this to ``-c search_path=myschema`` - sets the session's value of the ``search_path`` to ``myschema``. + sets the session's value of the ``search_path`` to ``myschema``.z :param enable_log_db_messages: Optional. If enabled logs database messages sent to the client during the session. To avoid a memory leak psycopg2 only saves the last 50 messages. For details, see: `PostgreSQL logging configuration parameters From 1df931a1900cb07d679fe1b5da163f446d0ce18f Mon Sep 17 00:00:00 2001 From: Karen Braganza Date: Tue, 7 Jul 2026 15:02:31 -0400 Subject: [PATCH 4/9] Change naming convention for async functions --- .../airflow/providers/common/sql/hooks/sql.py | 20 +++++----- .../providers/common/sql/hooks/sql.pyi | 2 +- .../providers/common/sql/triggers/sql.py | 4 +- .../providers/postgres/hooks/postgres.py | 2 +- .../unit/postgres/hooks/test_postgres.py | 40 +++++++++---------- 5 files changed, 34 insertions(+), 34 deletions(-) diff --git a/providers/common/sql/src/airflow/providers/common/sql/hooks/sql.py b/providers/common/sql/src/airflow/providers/common/sql/hooks/sql.py index 7810c8cf9544e..dcea7642b9a4a 100644 --- a/providers/common/sql/src/airflow/providers/common/sql/hooks/sql.py +++ b/providers/common/sql/src/airflow/providers/common/sql/hooks/sql.py @@ -1185,7 +1185,7 @@ def _build_conn_kwargs_from_airflow_connection(self, db) -> dict: "raw_connection": db, } - async def async_get_conn(self) -> Any: + async def aget_conn(self) -> Any: if not hasattr(self, "_conn_lock"): self._conn_lock = asyncio.Lock() async with self._conn_lock: @@ -1203,8 +1203,8 @@ async def _call(self, func: Callable, *args, **kwargs) -> Any: return await sync_to_async(func)(*args, **kwargs) @asynccontextmanager - async def _async_create_autocommit_connection(self, autocommit: bool = False): - conn = await self.async_get_conn() + 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) @@ -1219,12 +1219,12 @@ async def _async_create_autocommit_connection(self, autocommit: bool = False): await self._call(close) @asynccontextmanager - async def _async_get_cursor(self, conn): + async def _aget_cursor(self, conn): cursor = getattr(conn, "cursor", None) if cursor is None: raise AirflowNotFoundException( f"{type(conn).__name__} has no cursor() method. " - "Override _async_get_cursor in the DB-specific hook." + "Override _aget_cursor in the DB-specific hook." ) cur_or_cm = cursor() if inspect.isawaitable(cur_or_cm): @@ -1247,10 +1247,10 @@ async def _async_get_cursor(self, conn): return raise RuntimeError( f"Unsupported cursor type returned by {type(conn).__name__}. " - "Override _async_get_cursor in the DB-specific hook." + "Override _aget_cursor in the DB-specific hook." ) - async def _async_run_command(self, cur, sql_statement, parameters): + async def _arun_command(self, cur, sql_statement, parameters): """Run a statement using an already open cursor.""" if self.log_sql: self.log.info("Running statement: %s, parameters: %s", sql_statement, parameters) @@ -1307,13 +1307,13 @@ async def arun( raise ValueError("List of SQL statements is empty") _last_result = None _last_description = None - async with self._async_create_autocommit_connection(autocommit) as conn: + async with self._acreate_autocommit_connection(autocommit) as conn: try: - async with self._async_get_cursor(conn) as cur: + async with self._aget_cursor(conn) as cur: results = [] for sql_statement in sql_list: sql_to_run = await self._call(self._translate_sql, sql_statement) - await self._async_run_command(cur, sql_to_run, parameters) + await self._arun_command(cur, sql_to_run, parameters) if handler is not None: handled = handler(cur) if inspect.isawaitable(handled): diff --git a/providers/common/sql/src/airflow/providers/common/sql/hooks/sql.pyi b/providers/common/sql/src/airflow/providers/common/sql/hooks/sql.pyi index 991ee195c86dd..65f6007361714 100644 --- a/providers/common/sql/src/airflow/providers/common/sql/hooks/sql.pyi +++ b/providers/common/sql/src/airflow/providers/common/sql/hooks/sql.pyi @@ -186,7 +186,7 @@ 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 async_get_conn(self) -> Any: ... + async def aget_conn(self) -> Any: ... @overload async def arun( self, diff --git a/providers/common/sql/src/airflow/providers/common/sql/triggers/sql.py b/providers/common/sql/src/airflow/providers/common/sql/triggers/sql.py index a05b6985b27b2..77b2927c04511 100644 --- a/providers/common/sql/src/airflow/providers/common/sql/triggers/sql.py +++ b/providers/common/sql/src/airflow/providers/common/sql/triggers/sql.py @@ -190,7 +190,7 @@ def _jsonsafe_descriptions( ) return safe - async def get_hook(self) -> DbApiHook: + async def aget_hook(self) -> DbApiHook: """ Return DbApiHook. @@ -208,7 +208,7 @@ async def get_hook(self) -> DbApiHook: async def run(self) -> AsyncIterator[TriggerEvent]: try: - hook = await self.get_hook() + hook = await self.aget_hook() self.log.info("Extracting data from %s", self.conn_id) self.log.info("Executing: \n %s", self.sql) diff --git a/providers/postgres/src/airflow/providers/postgres/hooks/postgres.py b/providers/postgres/src/airflow/providers/postgres/hooks/postgres.py index fca205d46caa5..a584d6979d8b5 100644 --- a/providers/postgres/src/airflow/providers/postgres/hooks/postgres.py +++ b/providers/postgres/src/airflow/providers/postgres/hooks/postgres.py @@ -297,7 +297,7 @@ def get_conn(self) -> CompatConnection: self.conn = self._create_connection(conn_args) return self.conn - async def async_get_conn(self) -> Any: + 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.") diff --git a/providers/postgres/tests/unit/postgres/hooks/test_postgres.py b/providers/postgres/tests/unit/postgres/hooks/test_postgres.py index 1bac52926e303..8becfad3575d3 100644 --- a/providers/postgres/tests/unit/postgres/hooks/test_postgres.py +++ b/providers/postgres/tests/unit/postgres/hooks/test_postgres.py @@ -612,19 +612,19 @@ def test_get_conn_cursor(self, mocker): ) @pytest.mark.asyncio - async def test_async_get_conn_raises_not_implemented_without_psycopg3(self, mocker): + async def test_aget_conn_raises_not_implemented_without_psycopg3(self, mocker): mocker.patch("airflow.providers.postgres.hooks.postgres.USE_PSYCOPG3", False) with pytest.raises(NotImplementedError, match="Async connections for PostgresHook require psycopg3"): - await self.db_hook.async_get_conn() + await self.db_hook.aget_conn() @pytest.mark.asyncio - async def test_async_get_conn(self, mocker): + async def test_aget_conn(self, mocker): mock_async_connect = mocker.patch("psycopg.AsyncConnection.connect", new_callable=mock.AsyncMock) mock_register = mocker.patch("airflow.providers.postgres.hooks.postgres.register_default_adapters") mock_connection = mocker.MagicMock() mock_async_connect.return_value = mock_connection - result = await self.db_hook.async_get_conn() + result = await self.db_hook.aget_conn() mock_async_connect.assert_awaited_once_with( user="login", @@ -637,12 +637,12 @@ async def test_async_get_conn(self, mocker): assert result is mock_connection @pytest.mark.asyncio - async def test_async_get_conn_with_cursor(self, mocker): + async def test_aget_conn_with_cursor(self, mocker): self.connection.extra = '{"cursor": "dictcursor"}' mock_async_connect = mocker.patch("psycopg.AsyncConnection.connect", new_callable=mock.AsyncMock) mocker.patch("airflow.providers.postgres.hooks.postgres.register_default_adapters") - await self.db_hook.async_get_conn() + await self.db_hook.aget_conn() mock_async_connect.assert_awaited_once_with( row_factory=psycopg.rows.dict_row, @@ -654,14 +654,14 @@ async def test_async_get_conn_with_cursor(self, mocker): ) @pytest.mark.asyncio - async def test_async_get_conn_with_invalid_cursor(self, mocker): + async def test_aget_conn_with_invalid_cursor(self, mocker): self.connection.extra = '{"cursor": "mycursor"}' with pytest.raises(ValueError, match="Invalid cursor passed mycursor"): - await self.db_hook.async_get_conn() + await self.db_hook.aget_conn() @pytest.mark.asyncio - async def test_async_get_conn_adds_notice_handler_when_enabled(self, mocker): + async def test_aget_conn_adds_notice_handler_when_enabled(self, mocker): mock_async_connect = mocker.patch("psycopg.AsyncConnection.connect", new_callable=mock.AsyncMock) mocker.patch("airflow.providers.postgres.hooks.postgres.register_default_adapters") mock_connection = mocker.MagicMock() @@ -670,28 +670,28 @@ async def test_async_get_conn_adds_notice_handler_when_enabled(self, mocker): hook = PostgresHook(enable_log_db_messages=True) hook.get_connection = mock.Mock(return_value=self.connection) - await hook.async_get_conn() + await hook.aget_conn() mock_connection.add_notice_handler.assert_called_once_with(hook._notice_handler) @pytest.mark.asyncio - async def test_async_get_conn_no_notice_handler_when_disabled(self, mocker): + async def test_aget_conn_no_notice_handler_when_disabled(self, mocker): mock_async_connect = mocker.patch("psycopg.AsyncConnection.connect", new_callable=mock.AsyncMock) mocker.patch("airflow.providers.postgres.hooks.postgres.register_default_adapters") mock_connection = mocker.MagicMock() mock_async_connect.return_value = mock_connection - await self.db_hook.async_get_conn() + await self.db_hook.aget_conn() mock_connection.add_notice_handler.assert_not_called() @pytest.mark.asyncio - async def test_async_get_conn_with_namedtuple_cursor(self, mocker): + async def test_aget_conn_with_namedtuple_cursor(self, mocker): self.connection.extra = '{"cursor": "namedtuplecursor"}' mock_async_connect = mocker.patch("psycopg.AsyncConnection.connect", new_callable=mock.AsyncMock) mocker.patch("airflow.providers.postgres.hooks.postgres.register_default_adapters") - await self.db_hook.async_get_conn() + await self.db_hook.aget_conn() mock_async_connect.assert_awaited_once_with( row_factory=psycopg.rows.namedtuple_row, @@ -703,22 +703,22 @@ async def test_async_get_conn_with_namedtuple_cursor(self, mocker): ) @pytest.mark.asyncio - async def test_async_get_conn_with_realdictcursor_raises(self, mocker): + async def test_aget_conn_with_realdictcursor_raises(self, mocker): self.connection.extra = '{"cursor": "realdictcursor"}' mocker.patch("psycopg.AsyncConnection.connect", new_callable=mock.AsyncMock) from airflow.providers.common.compat.sdk import AirflowException with pytest.raises(AirflowException, match="realdictcursor is not supported with psycopg3"): - await self.db_hook.async_get_conn() + await self.db_hook.aget_conn() @pytest.mark.asyncio - async def test_async_get_conn_with_extra(self, mocker): + async def test_aget_conn_with_extra(self, mocker): self.connection.extra = '{"connect_timeout": 3}' mock_async_connect = mocker.patch("psycopg.AsyncConnection.connect", new_callable=mock.AsyncMock) mocker.patch("airflow.providers.postgres.hooks.postgres.register_default_adapters") - await self.db_hook.async_get_conn() + await self.db_hook.aget_conn() mock_async_connect.assert_awaited_once_with( user="login", @@ -730,14 +730,14 @@ async def test_async_get_conn_with_extra(self, mocker): ) @pytest.mark.asyncio - async def test_async_get_conn_with_options(self, mocker): + async def test_aget_conn_with_options(self, mocker): mock_async_connect = mocker.patch("psycopg.AsyncConnection.connect", new_callable=mock.AsyncMock) mocker.patch("airflow.providers.postgres.hooks.postgres.register_default_adapters") hook = PostgresHook(options="-c statement_timeout=3000ms") hook.get_connection = mock.Mock(return_value=self.connection) - await hook.async_get_conn() + await hook.aget_conn() mock_async_connect.assert_awaited_once_with( user="login", From e86bc2f4ed6b6d06d20bd439e6cf157740d11d76 Mon Sep 17 00:00:00 2001 From: Karen Braganza Date: Tue, 7 Jul 2026 16:49:29 -0400 Subject: [PATCH 5/9] Change fucntion names in test_sql.py --- .../tests/unit/common/sql/hooks/test_sql.py | 128 +++++++++--------- 1 file changed, 64 insertions(+), 64 deletions(-) diff --git a/providers/common/sql/tests/unit/common/sql/hooks/test_sql.py b/providers/common/sql/tests/unit/common/sql/hooks/test_sql.py index 44e5d1f75f476..439501d0479b8 100644 --- a/providers/common/sql/tests/unit/common/sql/hooks/test_sql.py +++ b/providers/common/sql/tests/unit/common/sql/hooks/test_sql.py @@ -392,30 +392,30 @@ def test_build_conn_kwargs_falls_back_to_dbname(self): @pytest.mark.db_test @pytest.mark.asyncio - async def test_async_get_conn_success(self): + async def test_aget_conn_success(self): hook = mock_db_hook(DbApiHook) mock_db_conn = MagicMock() hook.connector = AsyncMock() hook.connector.connect.return_value = mock_db_conn - assert await hook.async_get_conn() is mock_db_conn + assert await hook.aget_conn() is mock_db_conn hook.connector.connect.assert_awaited_once() @pytest.mark.db_test @pytest.mark.asyncio - async def test_async_get_conn_raises_without_connector(self): + async def test_aget_conn_raises_without_connector(self): hook = mock_db_hook(DbApiHook) hook.connector = None with pytest.raises(RuntimeError, match="didn't have `self.connector` set"): - await hook.async_get_conn() + await hook.aget_conn() @pytest.mark.db_test @pytest.mark.asyncio - async def test_async_get_conn_caches_airflow_connection(self): + async def test_aget_conn_caches_airflow_connection(self): hook = mock_db_hook(DbApiHook) hook.connector = AsyncMock() hook.connector.connect.return_value = MagicMock() - await hook.async_get_conn() - await hook.async_get_conn() + await hook.aget_conn() + await hook.aget_conn() assert hook.connection_invocations == 1 @pytest.mark.db_test @@ -440,76 +440,76 @@ def fn(x): @pytest.mark.db_test @pytest.mark.asyncio - async def test_async_create_connection_yields_conn(self): + async def test_acreate_connection_yields_conn(self): hook = mock_db_hook(DbApiHook) mock_conn = MagicMock() mock_conn.close = MagicMock() - with patch.object(hook, "async_get_conn", AsyncMock(return_value=mock_conn)): - async with hook._async_create_autocommit_connection() as conn: + with patch.object(hook, "aget_conn", AsyncMock(return_value=mock_conn)): + async with hook._acreate_autocommit_connection() as conn: assert conn is mock_conn @pytest.mark.db_test @pytest.mark.asyncio - async def test_async_create_connection_calls_close(self): + async def test_acreate_connection_calls_close(self): hook = mock_db_hook(DbApiHook) mock_conn = MagicMock() mock_conn.close = MagicMock() del mock_conn.aclose - with patch.object(hook, "async_get_conn", AsyncMock(return_value=mock_conn)): - async with hook._async_create_autocommit_connection(): + with patch.object(hook, "aget_conn", AsyncMock(return_value=mock_conn)): + async with hook._acreate_autocommit_connection(): pass mock_conn.close.assert_called_once() @pytest.mark.db_test @pytest.mark.asyncio - async def test_async_create_connection_prefers_aclose(self): + async def test_acreate_connection_prefers_aclose(self): hook = mock_db_hook(DbApiHook) mock_conn = MagicMock() mock_conn.aclose = AsyncMock() mock_conn.close = MagicMock() - with patch.object(hook, "async_get_conn", AsyncMock(return_value=mock_conn)): - async with hook._async_create_autocommit_connection(): + with patch.object(hook, "aget_conn", AsyncMock(return_value=mock_conn)): + async with hook._acreate_autocommit_connection(): pass mock_conn.aclose.assert_awaited_once() mock_conn.close.assert_not_called() @pytest.mark.db_test @pytest.mark.asyncio - async def test_async_create_connection_sets_sync_autocommit(self): + async def test_acreate_connection_sets_sync_autocommit(self): hook = mock_db_hook(DbApiHook) hook.supports_autocommit = True mock_conn = MagicMock() mock_conn.close = MagicMock() del mock_conn.set_autocommit - with patch.object(hook, "async_get_conn", AsyncMock(return_value=mock_conn)): - async with hook._async_create_autocommit_connection(autocommit=True): + with patch.object(hook, "aget_conn", AsyncMock(return_value=mock_conn)): + async with hook._acreate_autocommit_connection(autocommit=True): pass assert mock_conn.autocommit is True @pytest.mark.db_test @pytest.mark.asyncio - async def test_async_create_connection_sets_async_autocommit(self): + async def test_acreate_connection_sets_async_autocommit(self): hook = mock_db_hook(DbApiHook) hook.supports_autocommit = True mock_conn = MagicMock() mock_conn.close = MagicMock() mock_conn.set_autocommit = AsyncMock() - with patch.object(hook, "async_get_conn", AsyncMock(return_value=mock_conn)): - async with hook._async_create_autocommit_connection(autocommit=True): + with patch.object(hook, "aget_conn", AsyncMock(return_value=mock_conn)): + async with hook._acreate_autocommit_connection(autocommit=True): pass mock_conn.set_autocommit.assert_awaited_once_with(True) @pytest.mark.db_test @pytest.mark.asyncio - async def test_async_get_cursor_raises_without_cursor_method(self): + async def test_aget_cursor_raises_without_cursor_method(self): hook = mock_db_hook(DbApiHook) with pytest.raises(AirflowNotFoundException, match="has no cursor\\(\\) method"): - async with hook._async_get_cursor(MagicMock(spec=[])): + async with hook._aget_cursor(MagicMock(spec=[])): pass @pytest.mark.db_test @pytest.mark.asyncio - async def test_async_get_cursor_uses_async_context_manager(self): + async def test_aget_cursor_uses_async_context_manager(self): hook = mock_db_hook(DbApiHook) mock_cur = MagicMock() cursor_cm = AsyncMock() @@ -517,12 +517,12 @@ async def test_async_get_cursor_uses_async_context_manager(self): cursor_cm.__aexit__ = AsyncMock(return_value=False) conn = MagicMock() conn.cursor.return_value = cursor_cm - async with hook._async_get_cursor(conn) as cur: + async with hook._aget_cursor(conn) as cur: assert cur is mock_cur @pytest.mark.db_test @pytest.mark.asyncio - async def test_async_get_cursor_awaits_cursor_then_uses_async_cm(self): + async def test_aget_cursor_awaits_cursor_then_uses_async_cm(self): hook = mock_db_hook(DbApiHook) mock_cur = MagicMock() cursor_cm = AsyncMock() @@ -534,12 +534,12 @@ async def async_cursor(): conn = MagicMock() conn.cursor = async_cursor - async with hook._async_get_cursor(conn) as cur: + async with hook._aget_cursor(conn) as cur: assert cur is mock_cur @pytest.mark.db_test @pytest.mark.asyncio - async def test_async_get_cursor_yields_cursor_with_async_execute(self): + async def test_aget_cursor_yields_cursor_with_async_execute(self): hook = mock_db_hook(DbApiHook) mock_cur = MagicMock() mock_cur.execute = AsyncMock() @@ -549,13 +549,13 @@ async def test_async_get_cursor_yields_cursor_with_async_execute(self): del mock_cur.aclose conn = MagicMock() conn.cursor.return_value = mock_cur - async with hook._async_get_cursor(conn) as cur: + async with hook._aget_cursor(conn) as cur: assert cur is mock_cur mock_cur.close.assert_called_once() @pytest.mark.db_test @pytest.mark.asyncio - async def test_async_get_cursor_raises_for_unsupported_type(self): + async def test_aget_cursor_raises_for_unsupported_type(self): hook = mock_db_hook(DbApiHook) mock_cur = MagicMock() mock_cur.execute = MagicMock() @@ -564,53 +564,53 @@ async def test_async_get_cursor_raises_for_unsupported_type(self): conn = MagicMock() conn.cursor.return_value = mock_cur with pytest.raises(RuntimeError, match="Unsupported cursor type"): - async with hook._async_get_cursor(conn) as _: + async with hook._aget_cursor(conn) as _: pass @pytest.mark.db_test @pytest.mark.asyncio - async def test_async_run_command_no_params(self): + async def test_arun_command_no_params(self): hook = mock_db_hook(DbApiHook) cur = AsyncMock() cur.rowcount = 0 - await hook._async_run_command(cur, "SELECT 1", None) + await hook._arun_command(cur, "SELECT 1", None) cur.execute.assert_awaited_once_with("SELECT 1") @pytest.mark.db_test @pytest.mark.asyncio - async def test_async_run_command_with_params(self): + async def test_arun_command_with_params(self): hook = mock_db_hook(DbApiHook) cur = AsyncMock() cur.rowcount = 3 - await hook._async_run_command(cur, "SELECT * FROM t WHERE id=%s", (42,)) + await hook._arun_command(cur, "SELECT * FROM t WHERE id=%s", (42,)) cur.execute.assert_awaited_once_with("SELECT * FROM t WHERE id=%s", (42,)) @pytest.mark.db_test @pytest.mark.asyncio - async def test_async_run_command_logs_rowcount(self): + async def test_arun_command_logs_rowcount(self): hook = mock_db_hook(DbApiHook) cur = AsyncMock() cur.rowcount = 5 mock_log = MagicMock() with patch.object(hook, "_log", mock_log): - await hook._async_run_command(cur, "DELETE FROM t", None) + await hook._arun_command(cur, "DELETE FROM t", None) mock_log.info.assert_any_call("Rows affected: %s", 5) @pytest.mark.db_test @pytest.mark.asyncio - async def test_async_run_command_skips_rowcount_when_negative(self): + async def test_arun_command_skips_rowcount_when_negative(self): hook = mock_db_hook(DbApiHook) cur = AsyncMock() cur.rowcount = -1 mock_log = MagicMock() with patch.object(hook, "_log", mock_log): - await hook._async_run_command(cur, "SELECT 1", None) + await hook._arun_command(cur, "SELECT 1", None) assert not any(call.args[0] == "Rows affected: %s" for call in mock_log.info.call_args_list) @pytest.mark.db_test @pytest.mark.asyncio @pytest.mark.parametrize("empty_sql", [[], "", "\n"]) - async def test_run_async_empty_sql_raises(self, empty_sql): + async def test_arun_empty_sql_raises(self, empty_sql): hook = mock_db_hook(DbApiHook) mock_conn = MagicMock() mock_conn.autocommit = False @@ -624,15 +624,15 @@ async def _cursor_cm(conn): yield AsyncMock() with ( - patch.object(hook, "_async_create_autocommit_connection", _conn_cm), - patch.object(hook, "_async_get_cursor", _cursor_cm), + patch.object(hook, "_acreate_autocommit_connection", _conn_cm), + patch.object(hook, "_aget_cursor", _cursor_cm), ): with pytest.raises(ValueError, match="List of SQL statements is empty"): - await hook.run_async(sql=empty_sql) + await hook.arun(sql=empty_sql) @pytest.mark.db_test @pytest.mark.asyncio - async def test_run_async_no_handler_returns_none(self): + async def test_arun_no_handler_returns_none(self): hook = mock_db_hook(DbApiHook) cur = AsyncMock() cur.rowcount = 0 @@ -649,10 +649,10 @@ async def _cursor_cm(conn): yield cur with ( - patch.object(hook, "_async_create_autocommit_connection", _conn_cm), - patch.object(hook, "_async_get_cursor", _cursor_cm), + patch.object(hook, "_acreate_autocommit_connection", _conn_cm), + patch.object(hook, "_aget_cursor", _cursor_cm), ): - assert await hook.run_async(sql="INSERT INTO t VALUES (1)") is None + assert await hook.arun(sql="INSERT INTO t VALUES (1)") is None @pytest.mark.db_test @pytest.mark.asyncio @@ -769,7 +769,7 @@ async def _cursor_cm(conn): ), ], ) - async def test_run_async_query( + async def test_arun_query( self, return_last, split_statements, @@ -795,10 +795,10 @@ async def _cursor_cm(conn): yield cur with ( - patch.object(hook, "_async_create_autocommit_connection", _conn_cm), - patch.object(hook, "_async_get_cursor", _cursor_cm), + patch.object(hook, "_acreate_autocommit_connection", _conn_cm), + patch.object(hook, "_aget_cursor", _cursor_cm), ): - results = await hook.run_async( + results = await hook.arun( sql=sql, handler=fetch_all_handler, return_last=return_last, @@ -811,7 +811,7 @@ async def _cursor_cm(conn): @pytest.mark.db_test @pytest.mark.asyncio - async def test_run_async_commits_on_success(self): + async def test_arun_commits_on_success(self): hook = mock_db_hook(DbApiHook) cur = AsyncMock() cur.rowcount = 0 @@ -828,16 +828,16 @@ async def _cursor_cm(conn): yield cur with ( - patch.object(hook, "_async_create_autocommit_connection", _conn_cm), - patch.object(hook, "_async_get_cursor", _cursor_cm), + patch.object(hook, "_acreate_autocommit_connection", _conn_cm), + patch.object(hook, "_aget_cursor", _cursor_cm), ): - await hook.run_async(sql="INSERT INTO t VALUES (1)") + await hook.arun(sql="INSERT INTO t VALUES (1)") mock_conn.commit.assert_called_once() @pytest.mark.db_test @pytest.mark.asyncio - async def test_run_async_rollback_on_exception(self): + async def test_arun_rollback_on_exception(self): hook = mock_db_hook(DbApiHook) cur = AsyncMock() cur.execute.side_effect = RuntimeError("DB error") @@ -854,17 +854,17 @@ async def _cursor_cm(conn): yield cur with ( - patch.object(hook, "_async_create_autocommit_connection", _conn_cm), - patch.object(hook, "_async_get_cursor", _cursor_cm), + patch.object(hook, "_acreate_autocommit_connection", _conn_cm), + patch.object(hook, "_aget_cursor", _cursor_cm), ): with pytest.raises(RuntimeError, match="DB error"): - await hook.run_async(sql="SELECT fail") + await hook.arun(sql="SELECT fail") mock_conn.rollback.assert_called_once() @pytest.mark.db_test @pytest.mark.asyncio - async def test_run_async_awaits_async_handler(self): + async def test_arun_awaits_async_handler(self): hook = mock_db_hook(DbApiHook) rows = [(1, "a"), (2, "b")] cur = AsyncMock() @@ -886,10 +886,10 @@ async def _cursor_cm(conn): yield cur with ( - patch.object(hook, "_async_create_autocommit_connection", _conn_cm), - patch.object(hook, "_async_get_cursor", _cursor_cm), + patch.object(hook, "_acreate_autocommit_connection", _conn_cm), + patch.object(hook, "_aget_cursor", _cursor_cm), ): - result = await hook.run_async(sql="SELECT 1", handler=async_handler) + result = await hook.arun(sql="SELECT 1", handler=async_handler) assert result == rows From 272e38f60c1a26c58d06912fc54f8ffaf5d29c2f Mon Sep 17 00:00:00 2001 From: Karen Braganza Date: Tue, 7 Jul 2026 17:58:01 -0400 Subject: [PATCH 6/9] Fix failing tests --- .../common/sql/tests/unit/common/sql/triggers/test_sql.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/providers/common/sql/tests/unit/common/sql/triggers/test_sql.py b/providers/common/sql/tests/unit/common/sql/triggers/test_sql.py index 823d6c46ae150..e79cb69f33b15 100644 --- a/providers/common/sql/tests/unit/common/sql/triggers/test_sql.py +++ b/providers/common/sql/tests/unit/common/sql/triggers/test_sql.py @@ -99,7 +99,7 @@ def test_run_fetch_results_returns_rows_and_descriptions(self): mock_hook.arun.return_value = rows mock_hook.descriptions = descriptions trigger = self._make_trigger(fetch_results=True) - with mock.patch.object(trigger, "get_hook", new=mock.AsyncMock(return_value=mock_hook)): + with mock.patch.object(trigger, "aget_hook", new=mock.AsyncMock(return_value=mock_hook)): events = run_trigger(trigger) # The built-in fetch handler is used; no user handler runs in the triggerer. @@ -114,7 +114,7 @@ def test_run_fetch_results_returns_rows_and_descriptions(self): def test_run_without_fetch_results_returns_no_results(self): mock_hook = self._make_mock_hook() trigger = self._make_trigger(fetch_results=False) - with mock.patch.object(trigger, "get_hook", new=mock.AsyncMock(return_value=mock_hook)): + with mock.patch.object(trigger, "aget_hook", new=mock.AsyncMock(return_value=mock_hook)): events = run_trigger(trigger) assert mock_hook.arun.await_args.kwargs["handler"] is None @@ -136,7 +136,7 @@ def test_run_yields_error_event_on_exception(self): mock_hook = self._make_mock_hook() mock_hook.arun.side_effect = Exception("DB error") trigger = self._make_trigger(fetch_results=True) - with mock.patch.object(trigger, "get_hook", new=mock.AsyncMock(return_value=mock_hook)): + with mock.patch.object(trigger, "aget_hook", new=mock.AsyncMock(return_value=mock_hook)): events = run_trigger(trigger) assert len(events) == 1 From e7ee08fea987ecd7faa09c2400ec1b9425ca63fa Mon Sep 17 00:00:00 2001 From: Karen Braganza Date: Wed, 8 Jul 2026 19:49:16 -0400 Subject: [PATCH 7/9] Use correct async methods to get hook and conn --- .../sql/src/airflow/providers/common/sql/hooks/sql.py | 3 ++- .../sql/src/airflow/providers/common/sql/triggers/sql.py | 6 ++---- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/providers/common/sql/src/airflow/providers/common/sql/hooks/sql.py b/providers/common/sql/src/airflow/providers/common/sql/hooks/sql.py index a10297e396361..3786744a1cdd4 100644 --- a/providers/common/sql/src/airflow/providers/common/sql/hooks/sql.py +++ b/providers/common/sql/src/airflow/providers/common/sql/hooks/sql.py @@ -54,6 +54,7 @@ from airflow.exceptions import AirflowNotFoundException, AirflowProviderDeprecationWarning +from airflow.providers.common.compat.connection import get_async_connection from airflow.providers.common.compat.module_loading import import_string from airflow.providers.common.compat.sdk import ( AirflowException, @@ -1190,7 +1191,7 @@ async def aget_conn(self) -> Any: self._conn_lock = asyncio.Lock() async with self._conn_lock: if not self._connection: - self._connection = await sync_to_async(self.get_connection)(self.get_conn_id()) + self._connection = await get_async_connection(self.conn_id) db = self._connection if self.connector is None: raise RuntimeError(f"{type(self).__name__} didn't have `self.connector` set!") diff --git a/providers/common/sql/src/airflow/providers/common/sql/triggers/sql.py b/providers/common/sql/src/airflow/providers/common/sql/triggers/sql.py index 77b2927c04511..638d4a29de8f6 100644 --- a/providers/common/sql/src/airflow/providers/common/sql/triggers/sql.py +++ b/providers/common/sql/src/airflow/providers/common/sql/triggers/sql.py @@ -21,6 +21,7 @@ from asgiref.sync import sync_to_async +from airflow.providers.common.compat.hook import get_async_hook from airflow.providers.common.compat.sdk import AirflowException, BaseHook from airflow.providers.common.compat.version_compat import AIRFLOW_V_3_2_PLUS from airflow.providers.common.sql.hooks.handlers import fetch_all_handler @@ -89,8 +90,6 @@ def get_hook(self) -> DbApiHook: return hook async def _get_records(self) -> Any: - from asgiref.sync import sync_to_async - hook = self.get_hook() if AIRFLOW_V_3_2_PLUS: @@ -196,8 +195,7 @@ async def aget_hook(self) -> DbApiHook: :return: DbApiHook for this connection """ - connection = await sync_to_async(BaseHook.get_connection)(self.conn_id) - hook = await sync_to_async(connection.get_hook)() + hook = await get_async_hook(self.conn_id) if not isinstance(hook, DbApiHook) or not hasattr(hook, "arun"): raise AirflowException( f"You are trying to use the SqlExecuteQueryOperator in deferrable mode with {hook.__class__.__name__}," From 179c0d2e1fc5546970aeb6a8b12ecb7f73ea8b40 Mon Sep 17 00:00:00 2001 From: Karen Braganza Date: Wed, 8 Jul 2026 22:04:46 -0400 Subject: [PATCH 8/9] Fix failing tests --- .../src/airflow/providers/common/sql/hooks/sql.py | 2 +- .../sql/tests/unit/common/sql/hooks/test_sql.py | 14 ++++++++++---- 2 files changed, 11 insertions(+), 5 deletions(-) diff --git a/providers/common/sql/src/airflow/providers/common/sql/hooks/sql.py b/providers/common/sql/src/airflow/providers/common/sql/hooks/sql.py index 3786744a1cdd4..41cc9fccd0dd8 100644 --- a/providers/common/sql/src/airflow/providers/common/sql/hooks/sql.py +++ b/providers/common/sql/src/airflow/providers/common/sql/hooks/sql.py @@ -1191,7 +1191,7 @@ async def aget_conn(self) -> Any: self._conn_lock = asyncio.Lock() async with self._conn_lock: if not self._connection: - self._connection = await get_async_connection(self.conn_id) + 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!") diff --git a/providers/common/sql/tests/unit/common/sql/hooks/test_sql.py b/providers/common/sql/tests/unit/common/sql/hooks/test_sql.py index 439501d0479b8..dae8f9468dde0 100644 --- a/providers/common/sql/tests/unit/common/sql/hooks/test_sql.py +++ b/providers/common/sql/tests/unit/common/sql/hooks/test_sql.py @@ -392,8 +392,10 @@ def test_build_conn_kwargs_falls_back_to_dbname(self): @pytest.mark.db_test @pytest.mark.asyncio - async def test_aget_conn_success(self): + @patch("airflow.providers.common.sql.hooks.sql.get_async_connection", new_callable=AsyncMock) + async def test_aget_conn_success(self, mock_get_async_connection): hook = mock_db_hook(DbApiHook) + mock_get_async_connection.return_value = Connection(conn_id="c", conn_type="test") mock_db_conn = MagicMock() hook.connector = AsyncMock() hook.connector.connect.return_value = mock_db_conn @@ -402,21 +404,25 @@ async def test_aget_conn_success(self): @pytest.mark.db_test @pytest.mark.asyncio - async def test_aget_conn_raises_without_connector(self): + @patch("airflow.providers.common.sql.hooks.sql.get_async_connection", new_callable=AsyncMock) + async def test_aget_conn_raises_without_connector(self, mock_get_async_connection): hook = mock_db_hook(DbApiHook) + mock_get_async_connection.return_value = Connection(conn_id="c", conn_type="test") hook.connector = None with pytest.raises(RuntimeError, match="didn't have `self.connector` set"): await hook.aget_conn() @pytest.mark.db_test @pytest.mark.asyncio - async def test_aget_conn_caches_airflow_connection(self): + @patch("airflow.providers.common.sql.hooks.sql.get_async_connection", new_callable=AsyncMock) + async def test_aget_conn_caches_airflow_connection(self, mock_get_async_connection): hook = mock_db_hook(DbApiHook) + mock_get_async_connection.return_value = Connection(conn_id="c", conn_type="test") hook.connector = AsyncMock() hook.connector.connect.return_value = MagicMock() await hook.aget_conn() await hook.aget_conn() - assert hook.connection_invocations == 1 + mock_get_async_connection.assert_awaited_once() @pytest.mark.db_test @pytest.mark.asyncio From 966fad0983da00011499a7e9a77935669e7131ba Mon Sep 17 00:00:00 2001 From: Karen Braganza Date: Tue, 14 Jul 2026 11:30:00 -0400 Subject: [PATCH 9/9] Modify dependency in pyproject.toml --- providers/common/sql/pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/providers/common/sql/pyproject.toml b/providers/common/sql/pyproject.toml index 18e9cbe7fdb65..b90d35a94d5c3 100644 --- a/providers/common/sql/pyproject.toml +++ b/providers/common/sql/pyproject.toml @@ -60,7 +60,7 @@ requires-python = ">=3.10" # After you modify the dependencies, and rebuild your Breeze CI image with ``breeze ci-image build`` dependencies = [ "apache-airflow>=2.11.0", - "apache-airflow-providers-common-compat>=1.14.1", + "apache-airflow-providers-common-compat>=1.14.1", # use next version "sqlparse>=0.5.1", "more-itertools>=9.0.0", # The methodtools dependency is necessary since the introduction of dialects: