Yanislava commented on code in PR #75967:
URL: https://github.com/apache/airflow/pull/75967#discussion_r4414206816
##########
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:
Good catch, thanks! Fixed in the latest commit: `_drain_result_sets` now
saves `cursor.description` before draining and restores it afterwards, so
`hook.descriptions` matches the rows returned by the handler (this also covers
`vertica_fetch_all_handler`). While at it, I noticed that with `handler=None`
the hook started collecting descriptions, which `DbApiHook.run` doesn't do
without a handler, so it's reset to `[]` there as before. Added tests for both.
--
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]
