End a result whose fetch fails to convert

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<ConnectionClaim> claim;
+	try {
+		claim = ClaimWhileRunning();
+	} catch (...) { // NOLINT
+	}
+	CloseAfterFailure(std::move(error));
+}
+
 void Result::ThrowIfFailed() const {
 	if (failed) {
 		std::rethrow_exception(failure);

src/include/duckdb_python/engine/result.hpp

@@ -138,6 +138,8 @@ class Result : public enable_shared_from_this<Result> {
 	static void Register(const shared_ptr<Result> &result);
 	static vector<shared_ptr<Result>> 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();
 

src/include/duckdb_python/pyresult.hpp

@@ -116,6 +116,9 @@ struct DuckDBPyResult {
 	unique_ptr<DataChunk> FetchNext();
 	unique_ptr<DataChunk> FetchNextRaw();
 	unique_ptr<DataChunk> 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<NumpyResultConversion> InitializeNumpyConversion(bool pandas = false);
 
 	template <typename T>

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 <class NUMPY_T, bool PANDAS>
 	static NUMPY_T NullValue(bool &set_mask) {

src/pyresult.cpp

@@ -136,6 +136,27 @@ unique_ptr<ColumnDataCollection> 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<nb::tuple> DuckDBPyResult::Fetchone() {
 	if (Empty()) {
 		throw InvalidInputException("result closed");
@@ -154,14 +175,19 @@ Optional<nb::tuple> 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<DataChunk> chunk;
-		{
-			DUCKDB_PY_ASSERT_GIL_HELD();
-			nb::gil_scoped_release release;
-			chunk = FetchNextRaw();
+		for (; count_vec < vectors_per_chunk; count_vec++) {
+			unique_ptr<DataChunk> 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() {

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,)]