Create deferrable SqlExecuteQueryOperator compatible with Postgres

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

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.
 

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

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,

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),