its-amann commented on code in PR #75967:
URL: https://github.com/apache/airflow/pull/75967#discussion_r4414128622
##########
providers/vertica/src/airflow/providers/vertica/hooks/vertica.py:
##########
@@ -29,13 +30,41 @@
from sqlalchemy.engine import URL
+def _drain_result_sets(cursor) -> None:
+ """
+ Advance the cursor through all remaining result sets.
+
+ 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 reaches its result
set. ``nextset()``
+ discards the remaining rows of the current set (without loading them in
memory) and raises
+ the server error if the next set is an error.
+
+ See https://github.com/apache/airflow/issues/32993 and
+ https://github.com/vertica/vertica-python/issues/197.
+ """
+ while cursor.nextset():
+ pass
+
+
+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."""
+
+ @functools.wraps(handler)
+ def wrapper(cursor):
+ result = handler(cursor)
+ _drain_result_sets(cursor)
Review Comment:
one thing here: DbApiHook.run reads cur.description after the handler
returns, and nextset() overwrites description with the next result set's
columns. so with fetch_one_handler on `SELECT 1 AS a; SELECT 2 AS b, 3 AS c`
you get (1,) back but hook.descriptions says b, c. maybe save
cursor.description before draining and set it back after?
vertica_fetch_all_handler has the same thing
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
