Fixing the behaviours of SQL Hooks and Operators finally

airflow/providers/common/sql/CHANGELOG.rst

@@ -32,7 +32,9 @@ This release fixes a few errors that were introduced in common.sql operator whil
 * ``_process_output`` method in ``SQLExecuteQueryOperator`` has now consistent semantics and typing, it
   can also modify the returned (and stored in XCom) values in the operators that derive from the
   ``SQLExecuteQueryOperator``).
-* last description of the cursor whether to return scalar values are now stored in DBApiHook
+* descriptions of all returned results are stored as descriptions property in the DBApiHook
+* last description of the cursor whether to return single query results values are now exposed in
+  DBApiHook via last_description property.
 
 Lack of consistency in the operator caused ``1.3.0`` to be yanked - the ``1.3.0`` should not be used - if
 you have ``1.3.0`` installed, upgrade to ``1.3.1``.

airflow/providers/common/sql/hooks/sql.py

@@ -30,6 +30,33 @@
 from airflow.version import version
 
 
+def return_single_query_results(sql: str | Iterable[str], return_last: bool, split_statements: bool):
+    """
+    Determines when results of single query only should be returned.
+
+    For compatibility reasons, the behaviour of the DBAPIHook is somewhat confusing.
+    In cases, when multiple queries are run, the return values will be an iterable (list) of results
+    - one for each query. However, in certain cases, when single query is run - the results will be just
+    the results of that single query without wrapping the results in a list.
+
+    The cases when single query results are returned without wrapping them in a list are when:
+
+    a) sql is string and last_statement is True (regardless what split_statement value is)
+    b) sql is string and split_statement is False
+
+    In all other cases, the results are wrapped in a list, even if there is only one statement to process:
+
+    a) always when sql is an iterable of string statements (regardless what last_statement value is)
+    b) when sql is string, split_statement is True and last_statement is False
+
+    :param sql: sql to run (either string or list of strings)
+    :param return_last: whether last statement output should only be returned
+    :param split_statements: whether to split string statements.
+    :return: True if the hook should return single query results
+    """
+    return isinstance(sql, str) and (return_last or not split_statements)
+
+
 def fetch_all_handler(cursor) -> list[tuple] | None:
     """Handler for DbApiHook.run() to return results"""
     if cursor.description is not None:
@@ -114,8 +141,7 @@ def __init__(self, *args, schema: str | None = None, log_sql: bool = True, **kwa
         # Hook deriving from the DBApiHook to still have access to the field in its constructor
         self.__schema = schema
         self.log_sql = log_sql
-        self.scalar_return_last = False
-        self.last_description: Sequence[Sequence] | None = None
+        self.descriptions: list[Sequence[Sequence] | None] = []
 
     def get_conn(self):
         """Returns a connection object"""
@@ -222,6 +248,12 @@ def split_sql_string(sql: str) -> list[str]:
         statements: list[str] = list(filter(None, splits))
         return statements
 
+    @property
+    def last_description(self) -> Sequence[Sequence] | None:
+        if not self.descriptions:
+            return None
+        return self.descriptions[-1]
+
     def run(
         self,
         sql: str | Iterable[str],
@@ -234,7 +266,42 @@ def run(
         """
         Runs a command or a list of commands. Pass a list of sql
         statements to the sql parameter to get them to execute
-        sequentially
+        sequentially.
+
+        The method will return either single query results (typically list of rows) or list of those results
+        where each element in the list are results of one of the queries (typically list of list of rows :D)
+
+        For compatibility reasons, the behaviour of the DBAPIHook is somewhat confusing.
+        In cases, when multiple queries are run, the return values will be an iterable (list) of results
+        - one for each query. However, in certain cases, when single query is run - the results will be just
+        the results of that query without wrapping the results in a list.
+
+        The cases when single query results are returned without wrapping them in a list are when:
+
+        a) sql is string and last_statement is True (regardless what split_statement value is)
+        b) sql is string and split_statement is False
+
+        In all other cases, the results are wrapped in a list, even if there is only one statement to process:
+
+        a) always when sql is an iterable of string statements (regardless what last_statement value is)
+        b) when sql is string, split_statement is True and last_statement is False
+
+
+        In any of those cases, however you can access the following properties of the Hook after running it:
+
+          * descriptions - has an array of cursor descriptions - each statement executed contain the list
+            of descriptions executed. If ``return_last`` is used, this is always a one-element array
+          * last_description - description of the last statement executed
+
+        Note that return value from the hook will ONLY be actually returned when handler is provided. Setting
+        the ``handler`` to None, results in this method returning None.
+
+        Handler is a way to process the rows from cursor (Iterator) into a value that is suitable to be
+        returned to XCom and generally fit in memory. As an optimization, handler is usually not executed
+        by the SQLExecuteQuery operator if `do_xcom_push` is not specified.
+
+        You can use pre-defined handles (`fetch_all_handler``, ''fetch_one_handler``) or implement your
+        own handler.
 
         :param sql: the sql statement to be executed (str) or a list of
             sql statements to execute
@@ -246,40 +313,48 @@ def run(
         :param return_last: Whether to return result for only last statement or for all after split
         :return: return only result of the ALL SQL expressions if handler was provided.
         """
-        self.scalar_return_last = isinstance(sql, str) and return_last
+        self.descriptions = []
+
         if isinstance(sql, str):
             if split_statements:
-                sql = self.split_sql_string(sql)
+                sql_list: Iterable[str] = self.split_sql_string(sql)
             else:
-                sql = [sql]
+                sql_list = [sql] if sql.strip() else []
+        else:
+            sql_list = sql
 
-        if sql:
-            self.log.debug("Executing following statements against DB: %s", 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
         with closing(self.get_conn()) as conn:
             if self.supports_autocommit:
                 self.set_autocommit(conn, autocommit)
 
             with closing(conn.cursor()) as cur:
                 results = []
-                for sql_statement in sql:
+                for sql_statement in sql_list:
                     self._run_command(cur, sql_statement, parameters)
 
                     if handler is not None:
                         result = handler(cur)
-                        results.append(result)
-                self.last_description = cur.description
+                        if 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 autocommit was set to False or db does not support autocommit, we do a manual commit.
             if not self.get_autocommit(conn):
                 conn.commit()
 
         if handler is None:
             return None
-        elif self.scalar_return_last:
-            return results[-1]
+        if return_single_query_results(sql, return_last, split_statements):
+            self.descriptions = [_last_description]
+            return _last_result
         else:
             return results
 

airflow/providers/common/sql/operators/sql.py

@@ -19,14 +19,13 @@
 
 import ast
 import re
-from typing import TYPE_CHECKING, Any, Callable, Iterable, Mapping, NoReturn, Sequence, SupportsAbs, overload
+from typing import TYPE_CHECKING, Any, Callable, Iterable, Mapping, NoReturn, Sequence, SupportsAbs
 
 from airflow.compat.functools import cached_property
 from airflow.exceptions import AirflowException, AirflowFailException
 from airflow.hooks.base import BaseHook
 from airflow.models import BaseOperator, SkipMixin
-from airflow.providers.common.sql.hooks.sql import DbApiHook, fetch_all_handler
-from airflow.typing_compat import Literal
+from airflow.providers.common.sql.hooks.sql import DbApiHook, fetch_all_handler, return_single_query_results
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context
@@ -190,11 +189,17 @@ class SQLExecuteQueryOperator(BaseSQLOperator):
     Executes SQL code in a specific database
     :param sql: the SQL code or string pointing to a template file to be executed (templated).
     File must have a '.sql' extensions.
+
+    When implementing a specific Operator, you can also implement `_process_output` method in the
+    hook to perform additional processing of values returned by the DB Hook of yours. For example, you
+    can join description retrieved from the cursors of your statements with returned values, or save
+    the output of your operator to a file.
+
     :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).
     :param split_statements: (optional) if split single SQL string into statements (default: False).
-    :param return_last: (optional) if return the result of only last statement (default: True).
+    :param return_last: (optional) return the result of only last statement (default: True).
 
     .. seealso::
         For more information on how to use this operator, take a look at the guide:
@@ -225,54 +230,42 @@ def __init__(
         self.split_statements = split_statements
         self.return_last = return_last
 
-    @overload
-    def _process_output(
-        self, results: Any, description: Sequence[Sequence] | None, scalar_results: Literal[True]
-    ) -> Any:
-        pass
-
-    @overload
-    def _process_output(
-        self, results: list[Any], description: Sequence[Sequence] | None, scalar_results: Literal[False]
-    ) -> Any:
-        pass
-
-    def _process_output(
-        self, results: Any | list[Any], description: Sequence[Sequence] | None, scalar_results: bool
-    ) -> Any:
+    def _process_output(self, results: list[Any], descriptions: list[Sequence[Sequence] | None]) -> list[Any]:
         """
-        Can be overridden by the subclass in case some extra processing is needed.
+        Processes output before it is returned by the operator.
+
+        It can be overridden by the subclass in case some extra processing is needed. Note that unlike
+        DBApiHook return values returned - the results passed and returned by ``_process_output`` should
+        always be lists of results - each element of the list is a result from a single SQL statement
+        (typically this will be list of Rows). You have to make sure that this is the same for returned
+        values = there should be one element in the list for each statement executed by the hook..
+
         The "process_output" method can override the returned output - augmenting or processing the
         output as needed - the output returned will be returned as execute return value and if
-        do_xcom_push is set to True, it will be set as XCom returned
+        do_xcom_push is set to True, it will be set as XCom returned.
 
         :param results: results in the form of list of rows.
-        :param description: as returned by ``cur.description`` in the Python DBAPI
-        :param scalar_results: True if result is single scalar value rather than list of rows
+        :param descriptions: list of descriptions returned by ``cur.description`` in the Python DBAPI
         """
         return results
 
     def execute(self, context):
         self.log.info("Executing: %s", self.sql)
         hook = self.get_db_hook()
-        if self.do_xcom_push:
-            output = hook.run(
-                sql=self.sql,
-                autocommit=self.autocommit,
-                parameters=self.parameters,
-                handler=self.handler,
-                split_statements=self.split_statements,
-                return_last=self.return_last,
-            )
-        else:
-            output = hook.run(
-                sql=self.sql,
-                autocommit=self.autocommit,
-                parameters=self.parameters,
-                split_statements=self.split_statements,
-            )
-
-        return self._process_output(output, hook.last_description, hook.scalar_return_last)
+        output = hook.run(
+            sql=self.sql,
+            autocommit=self.autocommit,
+            parameters=self.parameters,
+            handler=self.handler if self.do_xcom_push else None,
+            split_statements=self.split_statements,
+            return_last=self.return_last,
+        )
+        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]
+        return self._process_output(output, hook.descriptions)
 
     def prepare_template(self) -> None:
         """Parse template file for attribute parameters."""
@@ -284,6 +277,7 @@ class SQLColumnCheckOperator(BaseSQLOperator):
     """
     Performs one or more of the templated checks in the column_checks dictionary.
     Checks are performed on a per-column basis specified by the column_mapping.
+
     Each check can take one or more of the following options:
     - equal_to: an exact value to equal, cannot be used with other comparison options
     - greater_than: value that result should be strictly greater than

airflow/providers/databricks/hooks/databricks_sql.py

@@ -24,7 +24,7 @@
 from databricks.sql.client import Connection  # type: ignore[attr-defined]
 
 from airflow.exceptions import AirflowException
-from airflow.providers.common.sql.hooks.sql import DbApiHook
+from airflow.providers.common.sql.hooks.sql import DbApiHook, return_single_query_results
 from airflow.providers.databricks.hooks.databricks_base import BaseDatabricksHook
 
 LIST_SQL_ENDPOINTS_ENDPOINT = ("GET", "api/2.0/sql/endpoints")
@@ -151,49 +151,57 @@ def run(
         """
         Runs a command or a list of commands. Pass a list of sql
         statements to the sql parameter to get them to execute
-        sequentially
+        sequentially.
 
         :param sql: the sql statement to be executed (str) or a list of
             sql statements to execute
         :param autocommit: What to set the connection's autocommit setting to
-            before executing the query.
+            before executing the query. Note that currently there is no commit functionality
+            in Databricks SQL so this flag has no effect.
+
         :param parameters: The parameters to render the SQL query with.
         :param handler: The result handler which is called with the result of each statement.
         :param split_statements: Whether to split a single SQL string into statements and run separately
         :param return_last: Whether to return result for only last statement or for all after split
-        :return: return only result of the LAST SQL expression if handler was provided.
+        :return: return only result of the LAST SQL expression if handler was provided unless return_last
+            is set to False.
         """
-        self.scalar_return_last = isinstance(sql, str) and return_last
+        self.descriptions = []
         if isinstance(sql, str):
             if split_statements:
-                sql = self.split_sql_string(sql)
+                sql_list = [self.strip_sql_string(s) for s in self.split_sql_string(sql)]
             else:
-                sql = [self.strip_sql_string(sql)]
+                sql_list = [self.strip_sql_string(sql)]
+        else:
+            sql_list = [self.strip_sql_string(s) for s in sql]
 
-        if sql:
-            self.log.debug("Executing following statements against Databricks DB: %s", list(sql))
+        if sql_list:
+            self.log.debug("Executing following statements against Databricks DB: %s", sql_list)
         else:
             raise ValueError("List of SQL statements is empty")
 
         results = []
-        for sql_statement in sql:
+        for sql_statement in sql_list:
             # when using AAD tokens, it could expire if previous query run longer than token lifetime
             with closing(self.get_conn()) as conn:
                 self.set_autocommit(conn, autocommit)
 
                 with closing(conn.cursor()) as cur:
                     self._run_command(cur, sql_statement, parameters)
-
                     if handler is not None:
                         result = handler(cur)
-                        results.append(result)
-                    self.last_description = cur.description
+                        if return_single_query_results(sql, return_last, split_statements):
+                            results = [result]
+                            self.descriptions = [cur.description]
+                        else:
+                            results.append(result)
+                            self.descriptions.append(cur.description)
 
             self._sql_conn = None
 
         if handler is None:
             return None
-        elif self.scalar_return_last:
+        if return_single_query_results(sql, return_last, split_statements):
             return results[-1]
         else:
             return results