diff --git a/src/engine/result.cpp b/src/engine/result.cpp index 83de0d8e..3f293b47 100644 --- a/src/engine/result.cpp +++ b/src/engine/result.cpp @@ -238,6 +238,17 @@ void Result::CloseAfterFailure(std::exception_ptr error) noexcept { } } +void Result::CloseWithError(std::exception_ptr error) { + // The claim makes Python code that the close runs refuse to reenter the connection, and a busy or closed + // connection ends the query without it, as Close does + unique_ptr claim; + try { + claim = ClaimWhileRunning(); + } catch (...) { // NOLINT + } + CloseAfterFailure(std::move(error)); +} + void Result::ThrowIfFailed() const { if (failed) { std::rethrow_exception(failure); diff --git a/src/include/duckdb_python/engine/result.hpp b/src/include/duckdb_python/engine/result.hpp index f75e8665..feb5cadd 100644 --- a/src/include/duckdb_python/engine/result.hpp +++ b/src/include/duckdb_python/engine/result.hpp @@ -138,6 +138,8 @@ class Result : public enable_shared_from_this { static void Register(const shared_ptr &result); static vector> Live(); static void AfterForkInChild(); + //! Close, keeping the error for any later use of the result to report + void CloseWithError(std::exception_ptr error); //! The rows a completed DML statement changed, or -1 int64_t ChangedRows(); diff --git a/src/include/duckdb_python/pyresult.hpp b/src/include/duckdb_python/pyresult.hpp index 0c4978fe..abad6e35 100644 --- a/src/include/duckdb_python/pyresult.hpp +++ b/src/include/duckdb_python/pyresult.hpp @@ -116,6 +116,9 @@ struct DuckDBPyResult { unique_ptr FetchNext(); unique_ptr FetchNextRaw(); unique_ptr TakeBufferedRows(); + //! Called while handling an error from reading rows: ends the result, so the connection is free and every later + //! fetch raises the error again + void EndWithCurrentError(); std::unique_ptr InitializeNumpyConversion(bool pandas = false); template diff --git a/src/numpy/array_wrapper.cpp b/src/numpy/array_wrapper.cpp index e4fe8562..97808625 100644 --- a/src/numpy/array_wrapper.cpp +++ b/src/numpy/array_wrapper.cpp @@ -9,6 +9,7 @@ #include "duckdb_python/python_objects.hpp" #include "duckdb_python/pyconnection/pyconnection.hpp" #include "duckdb/common/types/uuid.hpp" +#include "duckdb/main/error_manager.hpp" #include "duckdb/common/vector/list_vector.hpp" #include "duckdb/common/vector/array_vector.hpp" @@ -167,7 +168,15 @@ struct StringConvert { (void)append_data; auto data = const_data_ptr_cast(val.GetData()); auto len = val.GetSize(); - return PyUnicode_FromStringAndSize(const_char_ptr_cast(data), len); + auto str = PyUnicode_FromStringAndSize(const_char_ptr_cast(data), len); + if (!str) { + if (!PyErr_ExceptionMatches(PyExc_UnicodeDecodeError)) { + throw nb::python_error(); + } + PyErr_Clear(); + throw ErrorManager::InvalidUnicodeError(val.GetString(), "value construction"); + } + return str; } template static NUMPY_T NullValue(bool &set_mask) { diff --git a/src/pyresult.cpp b/src/pyresult.cpp index 41456c40..56e46119 100644 --- a/src/pyresult.cpp +++ b/src/pyresult.cpp @@ -136,6 +136,27 @@ unique_ptr DuckDBPyResult::TakeRows(const engine::Interrup return rows; } +void DuckDBPyResult::EndWithCurrentError() { + std::exception_ptr error; + try { + throw; + } catch (nb::python_error &python_error) { + // Raising a python_error consumes it, so a later fetch raises its text instead + error = std::make_exception_ptr(InvalidInputException(python_error.what())); + } catch (...) { + error = std::current_exception(); + } + current_chunk.reset(); + chunk_offset = 0; + try { + // Ending the query waits for its running tasks, and a task inside a Python UDF cannot finish until the GIL + // is free + duckdb::PyUtil::GilReleaseIfHeld release; + result->CloseWithError(std::move(error)); + } catch (...) { // NOLINT: the conversion error the caller rethrows is the one to report + } +} + Optional DuckDBPyResult::Fetchone() { if (Empty()) { throw InvalidInputException("result closed"); @@ -154,14 +175,19 @@ Optional DuckDBPyResult::Fetchone() { auto &types = GetTypes(); auto &client_properties = GetClientProperties(); duckdb::PyUtil::TupleBuilder row(types.size()); - for (idx_t col_idx = 0; col_idx < types.size(); col_idx++) { - auto &mask = FlatVector::Validity(current_chunk->data[col_idx]); - if (!mask.RowIsValid(chunk_offset)) { - row.append(nb::none()); - } else { - auto val = current_chunk->data[col_idx].GetValue(chunk_offset); - row.append(PythonObject::FromValue(val, types[col_idx], client_properties)); + try { + for (idx_t col_idx = 0; col_idx < types.size(); col_idx++) { + auto &mask = FlatVector::Validity(current_chunk->data[col_idx]); + if (!mask.RowIsValid(chunk_offset)) { + row.append(nb::none()); + } else { + auto val = current_chunk->data[col_idx].GetValue(chunk_offset); + row.append(PythonObject::FromValue(val, types[col_idx], client_properties)); + } } + } catch (...) { + EndWithCurrentError(); + throw; } chunk_offset++; return row.take(); @@ -267,37 +293,42 @@ nb::dict DuckDBPyResult::FetchNumpyInternal(bool chunked, idx_t vectors_per_chun } idx_t count_vec = 0; - if (vectors_per_chunk > 0) { - auto buffered = TakeBufferedRows(); - if (buffered) { - conversion.Append(*buffered); - count_vec++; + try { + if (vectors_per_chunk > 0) { + auto buffered = TakeBufferedRows(); + if (buffered) { + conversion.Append(*buffered); + count_vec++; + } } - } - for (; count_vec < vectors_per_chunk; count_vec++) { - unique_ptr chunk; - { - DUCKDB_PY_ASSERT_GIL_HELD(); - nb::gil_scoped_release release; - chunk = FetchNextRaw(); + for (; count_vec < vectors_per_chunk; count_vec++) { + unique_ptr chunk; + { + DUCKDB_PY_ASSERT_GIL_HELD(); + nb::gil_scoped_release release; + chunk = FetchNextRaw(); + } + if (!chunk || chunk->size() == 0) { + break; + } + conversion.Append(*chunk); } - if (!chunk || chunk->size() == 0) { - break; + InsertCategory(GetTypes(), categories); + + // now that we have materialized the result in contiguous arrays, construct the actual NumPy arrays or + // categorical types + nb::dict res; + auto names = ResultNames(); + QueryResult::DeduplicateColumns(names); + for (idx_t col_idx = 0; col_idx < names.size(); col_idx++) { + auto &name = names[col_idx]; + FillNumpy(res, col_idx, conversion, name.c_str()); } - conversion.Append(*chunk); + return res; + } catch (...) { + EndWithCurrentError(); + throw; } - InsertCategory(GetTypes(), categories); - - // now that we have materialized the result in contiguous arrays, construct the actual NumPy arrays or categorical - // types - nb::dict res; - auto names = ResultNames(); - QueryResult::DeduplicateColumns(names); - for (idx_t col_idx = 0; col_idx < names.size(); col_idx++) { - auto &name = names[col_idx]; - FillNumpy(res, col_idx, conversion, name.c_str()); - } - return res; } static void ReplaceDFColumn(PandasDataFrame &df, const char *col_name, idx_t idx, const nb::handle &new_value) { @@ -405,12 +436,24 @@ PandasDataFrame DuckDBPyResult::FrameFromNumpy(bool date_as_object, const nb::ha PandasDataFrame DuckDBPyResult::FetchDF(bool date_as_object) { auto conversion = InitializeNumpyConversion(true); - return FrameFromNumpy(date_as_object, FetchNumpyInternal(false, 1, std::move(conversion))); + auto arrays = FetchNumpyInternal(false, 1, std::move(conversion)); + try { + return FrameFromNumpy(date_as_object, arrays); + } catch (...) { + EndWithCurrentError(); + throw; + } } PandasDataFrame DuckDBPyResult::FetchDFChunk(idx_t num_of_vectors, bool date_as_object) { auto conversion = InitializeNumpyConversion(true); - return FrameFromNumpy(date_as_object, FetchNumpyInternal(true, num_of_vectors, std::move(conversion))); + auto arrays = FetchNumpyInternal(true, num_of_vectors, std::move(conversion)); + try { + return FrameFromNumpy(date_as_object, arrays); + } catch (...) { + EndWithCurrentError(); + throw; + } } nb::dict DuckDBPyResult::FetchPyTorch() { diff --git a/tests/fast/api/test_fetch_conversion_errors.py b/tests/fast/api/test_fetch_conversion_errors.py new file mode 100644 index 00000000..c9f4965b --- /dev/null +++ b/tests/fast/api/test_fetch_conversion_errors.py @@ -0,0 +1,137 @@ +"""A VARCHAR the client cannot turn into a Python str raises InvalidInputException and ends the result.""" + +import importlib.util +import re + +import pytest + +import duckdb + +pa = pytest.importorskip("pyarrow") + +needs_pandas = pytest.mark.skipif(importlib.util.find_spec("pandas") is None, reason="pandas not installed") + + +def _icu_loads(): + with duckdb.connect() as con: + try: + con.execute("LOAD icu") + except duckdb.Error: + return False + return True + + +needs_icu = pytest.mark.skipif(not _icu_loads(), reason="SET TimeZone needs the icu extension") + +ROWS = 5000 +INVALID_UTF8 = b"\xff" +INVALID_UNICODE = "Invalid unicode" + + +def strings(values): + # The Arrow scan does not check UTF-8, so the invalid bytes reach the client unchanged. + return pa.table({"s": pa.array(values, pa.binary()).cast(pa.string(), safe=False)}) + + +def strings_with_one_invalid_row(row): + values = [b"ok"] * ROWS + values[row] = INVALID_UTF8 + return strings(values) + + +@pytest.fixture +def con(): + connection = duckdb.connect() + yield connection + connection.close() + + +def fetch_one_at_a_time(con): + while con.fetchone() is not None: + pass + + +def fetch_many_at_a_time(con): + while con.fetchmany(1000): + pass + + +FETCHES = [ + pytest.param(fetch_one_at_a_time, id="fetchone"), + pytest.param(fetch_many_at_a_time, id="fetchmany"), + pytest.param(lambda con: con.fetchall(), id="fetchall"), + pytest.param(lambda con: con.df(), id="df", marks=needs_pandas), + pytest.param(lambda con: con.fetchnumpy(), id="fetchnumpy"), +] +INVALID_ROWS = [pytest.param(0, id="first_vector"), pytest.param(3000, id="later_vector")] + + +class TestAConversionErrorEndsTheResult: + @pytest.mark.parametrize("row", INVALID_ROWS) + @pytest.mark.parametrize("fetch", FETCHES) + def test_the_next_statement_runs(self, con, fetch, row): + con.register("strings", strings_with_one_invalid_row(row)) + con.execute("SELECT s FROM strings") + with pytest.raises(duckdb.InvalidInputException, match=INVALID_UNICODE): + fetch(con) + assert con.sql("SELECT 42").fetchall() == [(42,)] + + @pytest.mark.parametrize("fetch", FETCHES) + def test_a_later_fetch_raises_the_same_error(self, con, fetch): + con.register("strings", strings_with_one_invalid_row(3000)) + con.execute("SELECT s FROM strings") + with pytest.raises(duckdb.InvalidInputException, match=INVALID_UNICODE) as first: + fetch(con) + for later in (con.fetchall, con.fetchone): + with pytest.raises(duckdb.InvalidInputException) as again: + later() + assert str(again.value) == str(first.value) + assert con.sql("SELECT 42").fetchall() == [(42,)] + + @needs_icu + def test_a_python_error_is_raised_again_as_its_text(self, con): + # The instant fits a datetime in UTC, but its wall clock in Tokyo is past datetime.max. + con.execute("SET TimeZone = 'Asia/Tokyo'") + con.execute("SELECT TIMESTAMPTZ '9999-12-31 23:59:59+00'") + with pytest.raises(OverflowError) as first: + con.fetchone() + with pytest.raises(duckdb.InvalidInputException, match=re.escape(str(first.value))): + con.fetchall() + assert con.sql("SELECT 42").fetchall() == [(42,)] + + @needs_icu + @needs_pandas + def test_a_python_error_in_a_data_frame_chunk_ends_the_result(self, con): + # The engine accepts the old Java three letter zone names, which the IANA database pandas reads lacks. + con.execute("SET TimeZone = 'ACT'") + con.execute("SELECT TIMESTAMPTZ '2024-01-01 00:00:00+00' + INTERVAL (i) SECOND FROM range(10000) t(i)") + with pytest.raises(KeyError) as first: + con.fetch_df_chunk(1) + with pytest.raises(duckdb.InvalidInputException, match=re.escape(str(first.value))): + con.fetchall() + assert con.sql("SELECT 42").fetchall() == [(42,)] + + def test_the_transaction_stays_usable(self, con): + con.register("strings", strings_with_one_invalid_row(3000)) + con.execute("BEGIN") + con.execute("CREATE TABLE kept (i INTEGER)") + con.execute("SELECT s FROM strings") + with pytest.raises(duckdb.InvalidInputException, match=INVALID_UNICODE): + con.fetchall() + con.sql("INSERT INTO kept VALUES (1)") + con.execute("COMMIT") + assert con.sql("SELECT i FROM kept").fetchall() == [(1,)] + + +class TestNumpyConversionOfInvalidUtf8: + @pytest.mark.parametrize( + "query", + [pytest.param("SELECT s FROM strings", id="varchar"), pytest.param("SELECT [s] AS s FROM strings", id="list")], + ) + @pytest.mark.parametrize("method", [pytest.param("fetchnumpy"), pytest.param("df", marks=needs_pandas)]) + @pytest.mark.parametrize("source", ["execute", "sql"]) + def test_raises_invalid_input_like_fetchall(self, con, source, method, query): + con.register("strings", strings([b"ok", INVALID_UTF8, None])) + with pytest.raises(duckdb.InvalidInputException, match=INVALID_UNICODE): + getattr(getattr(con, source)(query), method)() + assert con.sql("SELECT 42").fetchall() == [(42,)]