This is an automated email from the ASF dual-hosted git repository.
potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 6c364f7e29b Raise errors from all statements of multi-statement SQL in
VerticaHook (#75967)
6c364f7e29b is described below
commit 6c364f7e29b6a5efacec95b9e7598129512257e9
Author: Ianislava Chalbysheva <[email protected]>
AuthorDate: Mon Nov 9 22:14:09 2026 +0700
Raise errors from all statements of multi-statement SQL in VerticaHook
(#75967)
* Raise errors from all statements of multi-statement SQL in VerticaHook
vertica-python only raises the error of a statement other than the first
one in a multi-statement query once its result set is reached with
cursor.nextset(). VerticaHook only drained the remaining result sets when
the handler was fetch_all_handler and the first statement returned rows.
With no handler (e.g. SQLExecuteQueryOperator with do_xcom_push=False),
fetch_one_handler, a custom handler, or a first statement without result
description (DDL), a failing statement was silently ignored, the
transaction was committed and the task succeeded.
Drain all result sets after every statement, whatever the handler is.
Co-Authored-By: Claude Opus 5.5 <[email protected]>
* Keep description of the consumed result set after draining
nextset() overwrites cursor.description with the columns of each next
result set, while DbApiHook.run reads it after the handler returns, so
hook.descriptions described the last result set instead of the rows
returned by the handler. Restore the description after draining, and keep
hook.descriptions empty when no handler is given, as before.
Co-Authored-By: Claude Opus 5.5 <[email protected]>
* Read every result set with fetchone() when draining
nextset() skips the rest of the current result set, including an error
that arrives after the set has started returning rows (e.g.
"SELECT 1; SELECT a/b FROM t" failing on the second row of t), so such an
error was still swallowed; vertica-python's own test_nextset_with_error_2
asserts that behaviour. On vertica-python 1.3.0-1.3.4 skipping does not
stop at ReadyForQuery and can hang. Read each result set with fetchone(),
which raises on the error, before moving to the next one.
Also wrap every handler the same way (vertica_fetch_all_handler is kept
as a public name only), update the class docstring, and mock two result
sets for the "split_statements=False" case of test_vertica_run_queries so
it goes through the drain.
Co-Authored-By: Claude Opus 5.5 <[email protected]>
---------
Co-authored-by: Claude Opus 5.5 <[email protected]>
---
.../src/airflow/providers/vertica/hooks/vertica.py | 81 ++++++++++----
.../tests/unit/vertica/hooks/test_vertica.py | 122 ++++++++++++++++++++-
.../tests/unit/vertica/hooks/test_vertica_sql.py | 19 +++-
3 files changed, 194 insertions(+), 28 deletions(-)
diff --git a/providers/vertica/src/airflow/providers/vertica/hooks/vertica.py
b/providers/vertica/src/airflow/providers/vertica/hooks/vertica.py
index 84634ff26bf..1913555bd32 100644
--- a/providers/vertica/src/airflow/providers/vertica/hooks/vertica.py
+++ b/providers/vertica/src/airflow/providers/vertica/hooks/vertica.py
@@ -16,6 +16,7 @@
# under the License.
from __future__ import annotations
+import functools
from collections.abc import Callable, Iterable, Mapping
from typing import TYPE_CHECKING, Any, overload
@@ -29,29 +30,54 @@ if TYPE_CHECKING:
from sqlalchemy.engine import URL
-def vertica_fetch_all_handler(cursor) -> list[tuple] | None:
+def _drain_result_sets(cursor) -> None:
+ """
+ Consume all remaining rows and result sets of the cursor.
+
+ vertica-python streams results: an error raised by a statement other than
the first one in a
+ multi-statement query is only reported once the client reads up to it.
``nextset()`` alone is
+ not enough: it skips the rest of the current result set, including an
error that arrives after
+ the set has started returning rows (e.g. ``SELECT a/b FROM t`` failing on
its second row), and
+ on vertica-python < 1.3.5 skipping can hang instead. So every result set
is read row by row
+ with ``fetchone()``, which raises on such an error, before moving to the
next one.
+
+ ``nextset()`` also overwrites ``cursor.description`` with the columns of
each next result set,
+ while ``DbApiHook.run`` reads it after the handler returns. The
description of the result set
+ the handler consumed is restored, so that ``hook.descriptions`` matches
the returned rows.
+
+ See https://github.com/apache/airflow/issues/32993 and
+ https://github.com/vertica/vertica-python/issues/197.
"""
- Replace the default DbApiHook fetch_all_handler in order to fix this issue
https://github.com/apache/airflow/issues/32993.
+ description = cursor.description
+ while True:
+ while cursor.fetchone() is not None:
+ pass
+ if not cursor.nextset():
+ break
+ cursor.description = description
+
- Returned value will not change after the initial call of
fetch_all_handler, all the remaining code is here
- only to make vertica client throws error.
- With Vertica, if you run the following sql (with split_statements set to
false):
+def _with_result_set_drain(handler: Callable[[Any], Any]) -> Callable[[Any],
Any]:
+ """Wrap a cursor handler so that all result sets are consumed after the
handler has run."""
- INSERT INTO MyTable (Key, Label) values (1, 'test 1');
- INSERT INTO MyTable (Key, Label) values (1, 'test 2');
- INSERT INTO MyTable (Key, Label) values (3, 'test 3');
+ @functools.wraps(handler)
+ def wrapper(cursor):
+ result = handler(cursor)
+ _drain_result_sets(cursor)
+ return result
- each insert will have its own result set and if you don't try to fetch
data of those result sets
- you won't detect error on the second insert.
+ return wrapper
+
+
+def vertica_fetch_all_handler(cursor) -> list[tuple] | None:
+ """
+ Return all rows of the current result set, then consume the remaining
result sets.
+
+ Kept for backward compatibility: ``VerticaHook.run`` now drains the result
sets for every
+ handler, so passing ``fetch_all_handler`` is enough. See
https://github.com/apache/airflow/issues/32993.
"""
result = fetch_all_handler(cursor)
- # loop on all statement result sets to get errors
- if cursor.description is not None:
- while cursor.nextset():
- if cursor.description is not None:
- row = cursor.fetchone()
- while row:
- row = cursor.fetchone()
+ _drain_result_sets(cursor)
return result
@@ -59,7 +85,8 @@ class VerticaHook(DbApiHook):
"""
Interact with Vertica.
- This hook use a customized version of default fetch_all_handler named
vertica_fetch_all_handler.
+ ``run`` consumes all rows and result sets after each executed statement,
whatever the handler is,
+ so that an error raised by any statement of a multi-statement query is not
silently ignored.
"""
conn_name_attr = "vertica_conn_id"
@@ -205,8 +232,18 @@ class VerticaHook(DbApiHook):
"""
Overwrite the common sql run.
- Will automatically replace fetch_all_handler by
vertica_fetch_all_handler.
+ vertica-python only reports an error raised by a statement once its
result set is reached, so
+ every result set is consumed after each executed statement, whatever
the handler is (including
+ when no handler is passed). Otherwise a failing statement in a
multi-statement query could be
+ silently ignored.
"""
- if handler == fetch_all_handler:
- handler = vertica_fetch_all_handler
- return DbApiHook.run(self, sql, autocommit, parameters, handler,
split_statements, return_last)
+ if handler is None:
+ DbApiHook.run(
+ self, sql, autocommit, parameters, _drain_result_sets,
split_statements, return_last
+ )
+ # Without a handler, DbApiHook.run does not collect descriptions;
keep that behaviour.
+ self.descriptions = []
+ return None
+ return DbApiHook.run(
+ self, sql, autocommit, parameters,
_with_result_set_drain(handler), split_statements, return_last
+ )
diff --git a/providers/vertica/tests/unit/vertica/hooks/test_vertica.py
b/providers/vertica/tests/unit/vertica/hooks/test_vertica.py
index 8fc260f1d85..2e14150eb93 100644
--- a/providers/vertica/tests/unit/vertica/hooks/test_vertica.py
+++ b/providers/vertica/tests/unit/vertica/hooks/test_vertica.py
@@ -21,8 +21,11 @@ import json
from unittest import mock
from unittest.mock import patch
+import pytest
+
from airflow.models import Connection
-from airflow.providers.vertica.hooks.vertica import VerticaHook
+from airflow.providers.common.sql.hooks.handlers import fetch_all_handler,
fetch_one_handler
+from airflow.providers.vertica.hooks.vertica import VerticaHook,
vertica_fetch_all_handler
class TestVerticaHookConn:
@@ -128,6 +131,8 @@ class TestVerticaHook:
def setup_method(self):
self.cur = mock.MagicMock(rowcount=0)
self.cur.nextset.side_effect = [None]
+ # The hook drains the cursor with fetchone() until it returns None.
+ self.cur.fetchone.return_value = None
self.conn = mock.MagicMock()
self.conn.cursor.return_value = self.cur
conn = self.conn
@@ -152,7 +157,7 @@ class TestVerticaHook:
def test_get_first_record(self):
statement = "SQL"
result_sets = [("row1",), ("row2",)]
- self.cur.fetchone.return_value = result_sets[0]
+ self.cur.fetchone.side_effect = [result_sets[0], None]
assert result_sets[0] == self.db_hook.get_first(statement)
self.conn.close.assert_called_once_with()
@@ -202,6 +207,8 @@ class TestVerticaHook:
class TestVerticaHookLineage:
def setup_method(self):
self.cur = mock.MagicMock(rowcount=0)
+ self.cur.nextset.return_value = False
+ self.cur.fetchone.return_value = None
self.conn = mock.MagicMock()
self.conn.cursor.return_value = self.cur
conn = self.conn
@@ -274,3 +281,114 @@ class TestVerticaHookLineage:
assert call_kw["context"] is self.db_hook
assert call_kw["sql"] == sql
assert call_kw["sql_parameters"] == parameters
+
+
+class TestVerticaHookResultSetDrain:
+ """
+ vertica-python only raises the error of a non-first statement once the
client reads up to it.
+
+ The mocks follow the driver: ``fetchone()`` returns rows of the current
result set, then
+ ``None``, and raises if the server error arrives in the middle of the set;
``nextset()`` moves to
+ the next result set and raises if the next "set" is an error.
+ """
+
+ def setup_method(self):
+ self.cur = mock.MagicMock(rowcount=0, description=None)
+ self.cur.fetchone.return_value = None
+ self.conn = mock.MagicMock()
+ self.conn.cursor.return_value = self.cur
+ conn = self.conn
+
+ class UnitTestVerticaHook(VerticaHook):
+ conn_name_attr = "test_conn_id"
+
+ def get_conn(self):
+ return conn
+
+ self.db_hook = UnitTestVerticaHook()
+
+ @pytest.mark.parametrize(
+ "handler",
+ [None, fetch_all_handler, fetch_one_handler, lambda cur: "custom"],
+ ids=["no_handler", "fetch_all", "fetch_one", "custom"],
+ )
+ def test_error_in_later_statement_is_raised(self, handler):
+ self.cur.nextset.side_effect = [True, RuntimeError("Duplicate key
values")]
+ with pytest.raises(RuntimeError, match="Duplicate key values"):
+ self.db_hook.run("INSERT INTO t VALUES (1); INSERT INTO t VALUES
(1);", handler=handler)
+
+ @pytest.mark.parametrize(
+ "handler",
+ [None, fetch_all_handler, fetch_one_handler, lambda cur: "custom"],
+ ids=["no_handler", "fetch_all", "fetch_one", "custom"],
+ )
+ def test_error_after_rows_of_later_statement_is_raised(self, handler):
+ # SELECT 1; SELECT a/b FROM t; -- the second statement returns a row,
then fails on the next
+ # one. nextset() would skip that error, only fetchone() raises it.
+ self.cur.fetchall.return_value = [(1,)]
+ self.cur.fetchone.side_effect = [None, (4,), RuntimeError("Division by
zero")]
+ self.cur.nextset.side_effect = [True, False]
+ with pytest.raises(RuntimeError, match="Division by zero"):
+ self.db_hook.run("SELECT 1; SELECT a/b FROM t;", handler=handler)
+
+ def test_error_in_rest_of_consumed_result_set_is_raised(self):
+ # fetch_one_handler reads one row of the only statement; its next row
fails.
+ self.cur.fetchone.side_effect = [(4,), RuntimeError("Division by
zero")]
+ self.cur.nextset.return_value = False
+ with pytest.raises(RuntimeError, match="Division by zero"):
+ self.db_hook.run("SELECT a/b FROM t;", handler=fetch_one_handler)
+
+ def test_all_result_sets_are_consumed(self):
+ self.cur.nextset.side_effect = [True, True, False]
+ assert self.db_hook.run("CREATE TABLE t (a INT); INSERT INTO t VALUES
(1); SELECT 1;") is None
+ assert self.cur.nextset.call_count == 3
+
+ def test_handler_result_is_returned_after_drain(self):
+ self.cur.nextset.side_effect = [True, False]
+ self.cur.description = [("?column?", None, None, None, None, None,
None)]
+ self.cur.fetchone.side_effect = [(1,), None, (2,), None]
+ assert self.db_hook.run("SELECT 1; SELECT 2;",
handler=fetch_one_handler) == (1,)
+ assert self.cur.nextset.call_count == 2
+
+ def test_each_split_statement_is_drained(self):
+ self.cur.nextset.side_effect = [False, False]
+ self.db_hook.run("SELECT 1; SELECT 2;", split_statements=True)
+ assert self.cur.execute.call_count == 2
+ assert self.cur.nextset.call_count == 2
+
+ def test_vertica_fetch_all_handler_drains(self):
+ self.cur.fetchall.return_value = [(1,)]
+ self.cur.nextset.side_effect = [True, RuntimeError("Division by zero")]
+ with pytest.raises(RuntimeError, match="Division by zero"):
+ vertica_fetch_all_handler(self.cur)
+
+ @pytest.mark.parametrize(
+ "handler",
+ [fetch_all_handler, fetch_one_handler, lambda cur: cur.fetchone()],
+ ids=["fetch_all", "fetch_one", "custom"],
+ )
+ def test_description_of_consumed_result_set_is_kept(self, handler):
+ first_description = [("a", None, None, None, None, None, None)]
+ last_description = [
+ ("b", None, None, None, None, None, None),
+ ("c", None, None, None, None, None, None),
+ ]
+ self.cur.description = first_description
+ self.cur.fetchone.side_effect = [(1,), None, (2, 3), None]
+ self.cur.fetchall.return_value = [(1,)]
+
+ def nextset():
+ if self.cur.description is first_description:
+ self.cur.description = last_description
+ return True
+ return False
+
+ self.cur.nextset.side_effect = nextset
+ self.db_hook.run("SELECT 1 AS a; SELECT 2 AS b, 3 AS c;",
handler=handler)
+ assert self.db_hook.descriptions == [first_description]
+
+ def test_no_descriptions_without_handler(self):
+ self.cur.description = [("a", None, None, None, None, None, None)]
+ self.cur.nextset.side_effect = [False]
+ self.db_hook.run("SELECT 1 AS a;")
+ assert self.db_hook.descriptions == []
diff --git a/providers/vertica/tests/unit/vertica/hooks/test_vertica_sql.py
b/providers/vertica/tests/unit/vertica/hooks/test_vertica_sql.py
index e95adb47989..3618d55baad 100644
--- a/providers/vertica/tests/unit/vertica/hooks/test_vertica_sql.py
+++ b/providers/vertica/tests/unit/vertica/hooks/test_vertica_sql.py
@@ -77,6 +77,8 @@ def mock_cursor(mock_get_conn):
cursor.fetchall.return_value = [("1", "row1")]
cursor.description = get_cursor_descriptions(["id", "value"])
cursor.nextset.side_effect = [False]
+ # The hook drains each result set with fetchone() until it returns None.
+ cursor.fetchone.return_value = None
mock_get_conn.return_value.cursor.return_value = cursor
return cursor
@@ -149,10 +151,17 @@ def test_vertica_run_queries(
):
if split_statements:
mock_cursor.fetchall.side_effect = cursor_results
- mock_cursor.nextset.side_effect = [True] * (len(cursor_results) - 1) +
[False]
else:
mock_cursor.fetchall.return_value = cursor_results
- mock_cursor.nextset.side_effect = lambda: False
+ # Like on a real Vertica, each executed SQL string produces one result set
per statement, and
+ # the hook moves through all of them with nextset().
+ result_sets_per_execute = [
+ len([statement for statement in sql_call.split(";") if
statement.strip()])
+ for sql_call in expected_calls
+ ]
+ mock_cursor.nextset.side_effect = [
+ has_next for count in result_sets_per_execute for has_next in [True] *
(count - 1) + [False]
+ ]
result = vertica_hook.run(
sql,
@@ -163,13 +172,15 @@ def test_vertica_run_queries(
expected_mock_calls = [mock.call(sql_call) for sql_call in expected_calls]
mock_cursor.execute.assert_has_calls(expected_mock_calls)
+ assert mock_cursor.nextset.call_count == sum(result_sets_per_execute)
assert [SerializableRow(*row) for row in result] == expected_result
def test_run_with_multiple_statements(vertica_hook, mock_cursor):
mock_cursor.fetchall.side_effect = [[(1,)], [(2,)]]
-
- mock_cursor.nextset.side_effect = [True, False]
+ # Statements are split, so each execute() produces a single result set.
+ mock_cursor.nextset.side_effect = None
+ mock_cursor.nextset.return_value = False
sql = "SELECT 1; SELECT 2;"
