kaxil commented on code in PR #75967:
URL: https://github.com/apache/airflow/pull/75967#discussion_r4501825210
##########
providers/vertica/src/airflow/providers/vertica/hooks/vertica.py:
##########
@@ -205,8 +234,20 @@ def run(
"""
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:
+ 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
+ if handler is fetch_all_handler:
Review Comment:
`vertica_fetch_all_handler` is now the same as
`_with_result_set_drain(fetch_all_handler)`, so could this branch go and every
handler get wrapped, keeping `vertica_fetch_all_handler` only as a public name?
The class docstring that says the hook uses a customized `fetch_all_handler` is
also out of date now that every handler is drained.
##########
providers/vertica/src/airflow/providers/vertica/hooks/vertica.py:
##########
@@ -29,13 +30,47 @@
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.
+
+ ``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.
+ """
+ description = cursor.description
+ while cursor.nextset():
Review Comment:
`nextset()` doesn't raise for an error that arrives after a later statement
has started returning rows. `flush_to_end_of_result()` skips every message up
to `ReadyForQuery`, `ErrorResponse` included, and vertica-python's own
`test_nextset_with_error_2` asserts exactly that for `SELECT 1; SELECT a/b FROM
t; SELECT 2` against a real server. The old `vertica_fetch_all_handler` caught
it because it called `fetchone()` after each `nextset()`, so with the default
handler that SQL now goes green where it used to fail, which is the case this
PR sets out to close. Could the drain read each set through `fetchone()` before
moving on (`while cursor.fetchone() is not None: pass` inside the `nextset()`
loop)? That also avoids a hang on vertica-python 1.3.0 to 1.3.4, still allowed
by the `>=1.3.0` floor, where `flush_to_end_of_result()` doesn't stop at
`ReadyForQuery`. A test where `nextset()` returns `True` and the next
`fetchone()` raises would pin it (the `MagicMock` cursors in the new tests
would need `fetchone` to return `None`).
##########
providers/vertica/tests/unit/vertica/hooks/test_vertica_sql.py:
##########
@@ -149,10 +149,11 @@ 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
+ # Each statement executed separately produces a single result set.
Review Comment:
This doesn't hold for the "Multiple queries, split_statements=False" case,
which sends both statements in one `execute()` and gets two result sets on real
Vertica. Could that case mock two sets, so the parametrized test goes through
the drain too?
--
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]
