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/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: 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 86fb4a466f9c0..c4180d9ca8b2a 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: @@ -43,6 +54,7 @@ from airflow.exceptions import 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, @@ -65,6 +77,7 @@ T = TypeVar("T") +HANDLER = Callable[[Any], Any | Awaitable[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,216 @@ def get_db_log_messages(self, conn) -> None: :param conn: Connection object """ + + def _translate_sql(self, sql: str) -> str: + """ + Translate SQL to driver-specific paramstyle. + + DB-specific hooks may override this to translate from a canonical style + to their driver's paramstyle if you want a unified SQL authoring style. + """ + return sql + + def _prepare_parameters(self, parameters: Iterable | Mapping[str, Any] | None): + """DB hooks may override to adapt parameter style.""" + return parameters + + def _build_conn_kwargs_from_airflow_connection(self, db) -> dict: + """Build a DB-agnostic kwargs dict.""" + extra = {} + extra_dejson = getattr(db, "extra_dejson", None) + if isinstance(extra_dejson, dict): + extra = extra_dejson + try: + port = int(db.port) if db.port else None + except (TypeError, ValueError): + port = None + return { + "host": db.host or "", + "port": port, + "username": db.login or "", + "password": db.password or "", + "database": extra.get("database") or extra.get("dbname") or "", + "schema": db.schema or "", + "conn_type": db.conn_type, + "extra": extra, + "raw_connection": db, + } + + async def aget_conn(self) -> Any: + if not hasattr(self, "_conn_lock"): + self._conn_lock = asyncio.Lock() + async with self._conn_lock: + if not self._connection: + self._connection = await get_async_connection(self.get_conn_id()) + db = self._connection + if self.connector is None: + raise RuntimeError(f"{type(self).__name__} didn't have `self.connector` set!") + conn_kwargs = self._build_conn_kwargs_from_airflow_connection(db) + return await self.connector.connect(**conn_kwargs) + + async def _call(self, func: Callable, *args, **kwargs) -> Any: + if inspect.iscoroutinefunction(func): + return await func(*args, **kwargs) + return await sync_to_async(func)(*args, **kwargs) + + @asynccontextmanager + async def _acreate_autocommit_connection(self, autocommit: bool = False): + conn = await self.aget_conn() + try: + if self.supports_autocommit: + set_autocommit = getattr(conn, "set_autocommit", None) + if set_autocommit and inspect.iscoroutinefunction(set_autocommit): + await set_autocommit(autocommit) + else: + self.set_autocommit(conn, autocommit) + yield conn + finally: + close = getattr(conn, "aclose", None) or getattr(conn, "close", None) + if close: + await self._call(close) + + @asynccontextmanager + async def _aget_cursor(self, conn): + cursor = getattr(conn, "cursor", None) + if cursor is None: + raise TypeError( + f"{type(conn).__name__} has no cursor() method. " + "Override _aget_cursor in the DB-specific hook." + ) + cur_or_cm = cursor() + if inspect.isawaitable(cur_or_cm): + cur_or_cm = await cur_or_cm + if hasattr(cur_or_cm, "__aenter__") and hasattr(cur_or_cm, "__aexit__"): + async with cur_or_cm as cur: + yield cur + return + execute = getattr(cur_or_cm, "execute", None) + if execute and inspect.iscoroutinefunction(execute): + try: + yield cur_or_cm + finally: + close = getattr(cur_or_cm, "aclose", None) or getattr(cur_or_cm, "close", None) + if close: + if inspect.iscoroutinefunction(close): + await close() + else: + close() + return + raise RuntimeError( + f"Unsupported cursor type returned by {type(conn).__name__}. " + "Override _aget_cursor in the DB-specific hook." + ) + + def supports_readonly_execution(self) -> bool: + """Whether the DB-specific hook overrides :meth:`_aenter_read_only`.""" + return type(self)._aenter_read_only is not DbApiHook._aenter_read_only + + async def _aenter_read_only(self, conn): + """ + Put the opened async connection in read-only mode. + + Called by :meth:`arun` before any statement runs, whenever the operator + requests read-only execution. The base implementation raises an exception + if this method is not implemented in the DB-specific hook. + """ + raise NotImplementedError( + f"{type(self).__name__} does not implement read-only execution. Override " + "_aenter_read_only to support deferrable read-only queries for this database," + ) + + 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) + 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 arun( + self, + sql: str | Iterable[str], + autocommit: bool = ..., + parameters: Iterable | Mapping[str, Any] | None = ..., + handler: None = ..., + split_statements: bool = ..., + return_last: bool = ..., + read_only: bool = ..., + ) -> None: ... + + @overload + async def arun( + self, + sql: str | Iterable[str], + autocommit: bool = ..., + parameters: Iterable | Mapping[str, Any] | None = ..., + handler: HANDLER = ..., + split_statements: bool = ..., + return_last: bool = ..., + read_only: bool = ..., + ) -> tuple | list | list[tuple] | list[list[tuple] | tuple] | None: ... + + async def arun( + 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, + read_only: bool = False, + ) -> 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._acreate_autocommit_connection(autocommit) as conn: + try: + if read_only: + await self._aenter_read_only(conn) + 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._arun_command(cur, sql_to_run, parameters) + if handler is not None: + handled = 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..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 @@ -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 aget_conn(self) -> Any: ... + @overload + async def arun( + self, + sql: str | Iterable[str], + autocommit: bool = ..., + parameters: Iterable | Mapping[str, Any] | None = ..., + handler: None = ..., + split_statements: bool = ..., + return_last: bool = ..., + ) -> None: ... + @overload + async def arun( + self, + sql: str | Iterable[str], + autocommit: bool = ..., + parameters: Iterable | Mapping[str, Any] | None = ..., + 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/read_only_guard.py b/providers/common/sql/src/airflow/providers/common/sql/operators/read_only_guard.py new file mode 100644 index 0000000000000..e546cb6186c9e --- /dev/null +++ b/providers/common/sql/src/airflow/providers/common/sql/operators/read_only_guard.py @@ -0,0 +1,84 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import logging + +log = logging.getLogger(__name__) + +try: + import sqlglot # this is an optional dependency + from sqlglot import exp + from sqlglot.errors import ParseError, TokenError + + _SQLGLOT_AVAILABLE = True +except ImportError: # pragma: no cover + _SQLGLOT_AVAILABLE = False + +# Postgres only matches the current scope of the deferrable SQLExecuteQueryOperator +_DIALECT = "postgres" + +if _SQLGLOT_AVAILABLE: + _WRITE_TYPES = (exp.Insert, exp.Update, exp.Delete, exp.Merge, exp.Create, exp.Alter, exp.Drop) + for name in ("TruncateTable", "Rename"): + cls = getattr(exp, name, None) + if cls is not None: + _WRITE_TYPES += (cls,) +else: + _WRITE_TYPES = () + + +def _statement_write_node(stmt: exp.Expression): + """Return the first node that proves this statement writes, or None.""" + if isinstance(stmt, _WRITE_TYPES): + return stmt + return next(stmt.find_all(*_WRITE_TYPES), None) + + +def scan_for_writes(sql: str | list[str]) -> tuple[bool, str]: + """Pre-defer check for write statements``.""" + if not _SQLGLOT_AVAILABLE: + return False, "sqlglot not installed; read-only transaction will enforce" + + if isinstance(sql, (list, tuple)): + sql_text = ";\n".join(str(s) for s in sql) + else: + sql_text = str(sql) + + try: + statements = sqlglot.parse(sql_text, read=_DIALECT) + except (ParseError, TokenError) as exc: + log.info( + "read-only guard: could not parse SQL (%s); deferring judgement to the read-only transaction", + exc.__class__.__name__, + ) + return False, "unparseable; read-only transaction will enforce" + + for idx, stmt in enumerate(statements): + if stmt is None: # empty fragment from a trailing ';' + continue + node = _statement_write_node(stmt) + if node is not None: + kind = type(node).__name__.upper() + return True, ( + f"statement #{idx + 1} is a proven write ({kind}): {stmt.sql(dialect=_DIALECT)[:120]!r}" + ) + + return False, ( + "no write detected by parse (writes inside functions/procedures/dynamic " + "SQL are invisible here and are enforced by the read-only transaction)" + ) 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..d824eb051ea10 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 @@ -37,6 +37,8 @@ ) 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.operators.read_only_guard import scan_for_writes +from airflow.providers.common.sql.triggers.sql import SQLExecuteQueryTrigger from airflow.utils.helpers import merge_dicts if TYPE_CHECKING: @@ -456,6 +458,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 +523,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 +542,12 @@ 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. + :param enforce_read_only: (optional) Only effective when ``deferrable=True``. If the triggerer restarts + mid-query the query is re-run from the start. By default (``True``) the operator locks the session + to read-only, so a restart can only ever repeat a read. Set to ``False`` to permit writes. Do this + **only** if your statements are idempotent. This flag does not make writes safe if they are not + idempotent. (default: True). .. seealso:: For more information on how to use this operator, take a look at the guide: @@ -499,7 +565,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 +579,8 @@ def __init__( return_last: bool = True, show_return_value_in_logs: bool = False, requires_result_fetch: bool = False, + deferrable: bool = False, + enforce_read_only: bool = True, **kwargs, ) -> None: super().__init__(conn_id=conn_id, database=database, **kwargs) @@ -525,6 +593,8 @@ 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 + self.enforce_read_only = enforce_read_only def _process_output( self, results: list[Any], descriptions: list[Sequence[Sequence] | None] @@ -554,29 +624,78 @@ def _should_run_output_processing(self) -> bool: def execute(self, context): self.log.info("Executing: %s", self.sql) - hook = self.get_db_hook() - if self.split_statements is not None: - extra_kwargs = {"split_statements": self.split_statements} + if self.deferrable: + if self.enforce_read_only: + is_write, reason = scan_for_writes(self.sql) + if is_write: + raise ValueError( + f"enforce_read_only=True but the SQL appears to contain a write: {reason}. " + "Set enforce_read_only=False if the query is idempotent, or deferrable=False " + "to run it on the worker." + ) + self.defer( + trigger=SQLExecuteQueryTrigger( + sql=self.sql, + conn_id=self.conn_id, + 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, + read_only=self.enforce_read_only, + ), + 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 RuntimeError("Unknown error in SQLExecuteQueryTrigger") + if event.get("status") == "error": + raise RuntimeError(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..68a9a5543764c 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 @@ -19,8 +19,12 @@ from typing import TYPE_CHECKING +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 from airflow.providers.common.sql.hooks.sql import DbApiHook from airflow.triggers.base import BaseTrigger, TriggerEvent @@ -28,10 +32,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 +61,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__}", { @@ -80,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: @@ -103,3 +111,153 @@ 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, + read_only: 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 + self.read_only = read_only + + 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, + "read_only": self.read_only, + }, + ) + + @staticmethod + def _jsonsafe_descriptions( + descriptions: list[Sequence[Sequence] | None], + ) -> list[list[list[Any]] | None]: + """ + Normalise cursor descriptions into a JSON-serializable form for the ``TriggerEvent``. + + ``cursor.description`` is a sequence of column 7-tuples whose ``type_code`` can be a + driver-specific object that is not JSON-serializable; such values are stringified while + JSON-native fields (column names, sizes, precision, ...) are preserved. + """ + safe: list[list[list[Any]] | None] = [] + for description in descriptions: + if description is None: + safe.append(None) + continue + safe.append( + [ + [ + field if isinstance(field, (str, int, float, bool, type(None))) else str(field) + for field in column + ] + for column in description + ] + ) + return safe + + async def aget_hook(self) -> DbApiHook: + """ + Return DbApiHook. + + :return: DbApiHook for this connection + """ + hook = await get_async_hook(self.conn_id) + if not isinstance(hook, DbApiHook) or not hasattr(hook, "arun"): + raise TypeError( + 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()}" + ) + if self.read_only and not hook.supports_readonly_execution(): + raise NotImplementedError( + f"{hook.__class__.__name__} does not support read-only execution, so it cannot run a" + " deferred query safely (a triggerer restart could re-run it). Set" + " enforce_read_only=False to run without the read-only guard if the query is" + " idempotent, or deferrable=False to run it on the worker." + ) + return hook + + async def run(self) -> AsyncIterator[TriggerEvent]: + try: + hook = await self.aget_hook() + + self.log.info("Extracting data from %s", self.conn_id) + self.log.info("Executing: \n %s", self.sql) + + 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.arun( + sql=self.sql, + autocommit=self.autocommit, + parameters=self.parameters, + handler=fetch_all_handler, + split_statements=self.split_statements, + return_last=self.return_last, + read_only=self.read_only, + ) + + 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.arun( + sql=self.sql, + autocommit=self.autocommit, + parameters=self.parameters, + handler=None, + split_statements=self.split_statements, + return_last=self.return_last, + read_only=self.read_only, + ) + + 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..63c74e4d2fbf2 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,7 +20,8 @@ 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 @@ -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,557 @@ 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 + @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 + assert await hook.aget_conn() is mock_db_conn + hook.connector.connect.assert_awaited_once() + + @pytest.mark.db_test + @pytest.mark.asyncio + @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 + @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() + mock_get_async_connection.assert_awaited_once() + + @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_acreate_connection_yields_conn(self): + hook = mock_db_hook(DbApiHook) + mock_conn = MagicMock() + mock_conn.close = MagicMock() + 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_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, "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_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, "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_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, "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_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, "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_aget_cursor_raises_without_cursor_method(self): + hook = mock_db_hook(DbApiHook) + with pytest.raises(TypeError, match="has no cursor\\(\\) method"): + async with hook._aget_cursor(MagicMock(spec=[])): + pass + + @pytest.mark.db_test + @pytest.mark.asyncio + async def test_aget_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._aget_cursor(conn) as cur: + assert cur is mock_cur + + @pytest.mark.db_test + @pytest.mark.asyncio + async def test_aget_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._aget_cursor(conn) as cur: + assert cur is mock_cur + + @pytest.mark.db_test + @pytest.mark.asyncio + async def test_aget_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._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_aget_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._aget_cursor(conn) as _: + pass + + @pytest.mark.db_test + @pytest.mark.asyncio + async def test_arun_command_no_params(self): + hook = mock_db_hook(DbApiHook) + cur = AsyncMock() + cur.rowcount = 0 + 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_arun_command_with_params(self): + hook = mock_db_hook(DbApiHook) + cur = AsyncMock() + cur.rowcount = 3 + 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_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._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_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._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_arun_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, "_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.arun(sql=empty_sql) + + @pytest.mark.db_test + @pytest.mark.asyncio + async def test_arun_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, "_acreate_autocommit_connection", _conn_cm), + patch.object(hook, "_aget_cursor", _cursor_cm), + ): + assert await hook.arun(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_arun_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, "_acreate_autocommit_connection", _conn_cm), + patch.object(hook, "_aget_cursor", _cursor_cm), + ): + results = await hook.arun( + 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_arun_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, "_acreate_autocommit_connection", _conn_cm), + patch.object(hook, "_aget_cursor", _cursor_cm), + ): + 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_arun_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, "_acreate_autocommit_connection", _conn_cm), + patch.object(hook, "_aget_cursor", _cursor_cm), + ): + with pytest.raises(RuntimeError, match="DB error"): + await hook.arun(sql="SELECT fail") + + mock_conn.rollback.assert_called_once() + + @pytest.mark.db_test + @pytest.mark.asyncio + async def test_arun_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, "_acreate_autocommit_connection", _conn_cm), + patch.object(hook, "_aget_cursor", _cursor_cm), + ): + result = await hook.arun(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..cfabaf28e9d39 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 TaskDeferred from airflow.models import Connection from airflow.providers.common.compat.openlineage.facet import ( Dataset, @@ -33,7 +34,9 @@ ) 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 import read_only_guard 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 +383,195 @@ 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(RuntimeError, 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(RuntimeError, 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]}, + ) + + def test_execute_defers_read_only_by_default(self): + op = SQLExecuteQueryOperator(task_id=TASK_ID, sql="SELECT 1", deferrable=True) + with pytest.raises(TaskDeferred) as exc: + op.execute({}) + assert exc.value.trigger.read_only is True + + def test_execute_passes_enforce_read_only_false_to_trigger(self): + op = SQLExecuteQueryOperator( + task_id=TASK_ID, sql="INSERT INTO foo VALUES (1)", deferrable=True, enforce_read_only=False + ) + with pytest.raises(TaskDeferred) as exc: + op.execute({}) + assert exc.value.trigger.read_only is False + + def test_execute_raises_before_deferring_on_proven_write(self): + op = SQLExecuteQueryOperator( + task_id=TASK_ID, sql="INSERT INTO foo VALUES (1)", deferrable=True, enforce_read_only=True + ) + with pytest.raises(ValueError, match="enforce_read_only=True but the SQL appears to contain a write"): + op.execute({}) + + +@pytest.mark.parametrize( + ("sql", "expected_kind"), + [ + pytest.param("SELECT * FROM foo", None, id="select"), + pytest.param("SELECT 1;", None, id="select-trailing-semicolon"), + pytest.param("INSERT INTO foo VALUES (1)", "INSERT", id="insert"), + pytest.param("UPDATE foo SET x = 1", "UPDATE", id="update"), + pytest.param("DELETE FROM foo", "DELETE", id="delete"), + pytest.param( + "MERGE INTO foo USING bar ON foo.id = bar.id WHEN MATCHED THEN UPDATE SET x = 1", + "MERGE", + id="merge", + ), + pytest.param("CREATE TABLE foo (id int)", "CREATE", id="create"), + pytest.param("ALTER TABLE foo ADD COLUMN y int", "ALTER", id="alter"), + pytest.param("DROP TABLE foo", "DROP", id="drop"), + pytest.param("TRUNCATE TABLE foo", "TRUNCATETABLE", id="truncate"), + pytest.param( + "WITH cte AS (INSERT INTO foo VALUES (1) RETURNING *) SELECT * FROM cte", + "INSERT", + id="write-inside-cte", + ), + ], +) +def test_scan_for_writes_detects_write_kind(sql, expected_kind): + is_write, reason = read_only_guard.scan_for_writes(sql) + if expected_kind is None: + assert is_write is False + assert "no write detected" in reason + else: + assert is_write is True + assert f"proven write ({expected_kind})" in reason + + +def test_scan_for_writes_detects_write_in_second_of_multiple_statements(): + is_write, reason = read_only_guard.scan_for_writes("SELECT 1; INSERT INTO foo VALUES (1)") + assert is_write is True + assert "statement #2" in reason + assert "INSERT" in reason + + +def test_scan_for_writes_list_of_read_only_statements(): + is_write, reason = read_only_guard.scan_for_writes(["SELECT 1", "SELECT 2"]) + assert is_write is False + assert "no write detected" in reason + + +def test_scan_for_writes_detects_write_across_list_of_statements(): + is_write, reason = read_only_guard.scan_for_writes(["SELECT 1", "INSERT INTO foo VALUES (1)"]) + assert is_write is True + assert "statement #2" in reason + + +def test_scan_for_writes_unparseable_sql_defers_to_read_only_transaction(): + is_write, reason = read_only_guard.scan_for_writes("SELECT * FROM foo WHERE x = 'unterminated") + assert is_write is False + assert "unparseable" in reason + + +def test_scan_for_writes_sqlglot_missing_defers_to_read_only_transaction(monkeypatch): + monkeypatch.setattr(read_only_guard, "_SQLGLOT_AVAILABLE", False) + is_write, reason = read_only_guard.scan_for_writes("INSERT INTO foo VALUES (1)") + assert is_write is False + assert "sqlglot not installed" in reason 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..9d6b3aa4debe9 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,111 @@ 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.arun = 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, + "read_only": 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.arun.return_value = rows + mock_hook.descriptions = descriptions + trigger = self._make_trigger(fetch_results=True) + 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. + assert mock_hook.arun.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, "aget_hook", new=mock.AsyncMock(return_value=mock_hook)): + events = run_trigger(trigger) + + assert mock_hook.arun.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.arun.side_effect = Exception("DB error") + trigger = self._make_trigger(fetch_results=True) + with mock.patch.object(trigger, "aget_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 be6519ddade24..620f9adadaffe 100644 --- a/providers/postgres/src/airflow/providers/postgres/hooks/postgres.py +++ b/providers/postgres/src/airflow/providers/postgres/hooks/postgres.py @@ -243,8 +243,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): @@ -269,16 +269,47 @@ def get_conn(self) -> CompatConnection: if arg_name not in self.ignored_extra_options: conn_args[arg_name] = arg_val - raw_cursor = conn.extra_dejson.get("cursor") + return conn_args, conn + def get_conn(self) -> CompatConnection: + """Establish a connection to a postgres database.""" + conn_args, conn = self._build_conn_args() + raw_cursor = conn.extra_dejson.get("cursor") if raw_cursor: key, value = self._get_cursor_config(raw_cursor) conn_args[key] = value - self.conn = self._create_connection(conn_args) - return self.conn + async def aget_conn(self) -> Any: + """Establish an async connection to a postgres database.""" + if not USE_PSYCOPG3: + raise NotImplementedError("Async connections for PostgresHook require psycopg3.") + from psycopg import AsyncConnection + + conn_args, conn = self._build_conn_args() + + 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 + + async def _aenter_read_only(self, conn) -> None: + """Put the async connection's next transaction into read-only mode.""" + await conn.set_read_only(True) + @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..8becfad3575d3 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_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.aget_conn() + + @pytest.mark.asyncio + 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.aget_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_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.aget_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_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.aget_conn() + + @pytest.mark.asyncio + 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() + 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.aget_conn() + + mock_connection.add_notice_handler.assert_called_once_with(hook._notice_handler) + + @pytest.mark.asyncio + 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.aget_conn() + + mock_connection.add_notice_handler.assert_not_called() + + @pytest.mark.asyncio + 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.aget_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_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.aget_conn() + + @pytest.mark.asyncio + 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.aget_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_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.aget_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 8e3cda63ef3d7..bf3d2f553526e 100644 --- a/providers/snowflake/src/airflow/providers/snowflake/operators/snowflake.py +++ b/providers/snowflake/src/airflow/providers/snowflake/operators/snowflake.py @@ -437,7 +437,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())