Custom fetch all handler for vertica to not miss errors

airflow/providers/vertica/hooks/vertica.py

@@ -17,9 +17,24 @@
 # under the License.
 from __future__ import annotations
 
+from typing import Any, Callable, Iterable, Mapping, overload
+
 from vertica_python import connect
 
-from airflow.providers.common.sql.hooks.sql import DbApiHook
+from airflow.providers.common.sql.hooks.sql import DbApiHook, fetch_all_handler
+
+
+def vertica_fetch_all_handler(cursor) -> list[tuple] | None:
+    """Replace the default DbApiHook fetch_all_handler ."""
+    to_return = 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()
+    return to_return
 
 
 class VerticaHook(DbApiHook):
@@ -99,3 +114,40 @@ def get_conn(self) -> connect:
 
         conn = connect(**conn_config)
         return conn
+
+    @overload
+    def run(
+        self,
+        sql: str | Iterable[str],
+        autocommit: bool = ...,
+        parameters: Iterable | Mapping[str, Any] | None = ...,
+        handler: None = ...,
+        split_statements: bool = ...,
+        return_last: bool = ...,
+    ) -> None:
+        ...
+
+    @overload
+    def run(
+        self,
+        sql: str | Iterable[str],
+        autocommit: bool = ...,
+        parameters: Iterable | Mapping[str, Any] | None = ...,
+        handler: Callable[[Any], Any] = ...,
+        split_statements: bool = ...,
+        return_last: bool = ...,
+    ) -> Any | list[Any]:
+        ...
+
+    def run(
+        self,
+        sql: str | Iterable[str],
+        autocommit: bool = False,
+        parameters: Iterable | Mapping | None = None,
+        handler: Callable[[Any], Any] | None = None,
+        split_statements: bool = False,
+        return_last: bool = True,
+    ) -> Any | list[Any] | None:
+        if handler == fetch_all_handler:
+            handler = vertica_fetch_all_handler
+        return DbApiHook.run(self, sql, autocommit, parameters, handler, split_statements, return_last)

tests/providers/vertica/hooks/test_vertica.py

@@ -127,6 +127,7 @@ def test_get_conn_extra_parameters_cast(self, mock_connect):
 class TestVerticaHook:
     def setup_method(self):
         self.cur = mock.MagicMock(rowcount=0)
+        self.cur.nextset.side_effect = [None]
         self.conn = mock.MagicMock()
         self.conn.cursor.return_value = self.cur
         conn = self.conn