Hold one connection claim from submission to settling in execute, executemany and the DB-API Arrow fetch

src/engine/connection.cpp

@@ -65,12 +65,16 @@ bool CallbackScope::Inside(const ClientContext &context) {
 	return false;
 }
 
+void ThrowBusy() {
+	throw ResourceInUseException(BUSY_ERROR);
+}
+
 ConnectionClaim::ConnectionClaim(ConnectionState &state_p, const ClientContext &context) : state(state_p) {
 	if (CallbackScope::Inside(context)) {
 		throw ResourceInUseException(CALLBACK_ERROR);
 	}
 	if (state.busy.exchange(true)) {
-		throw ResourceInUseException(BUSY_ERROR);
+		ThrowBusy();
 	}
 	if (Exiting()) {
 		state.admitted_during_exit = OnExitThread();
@@ -89,7 +93,7 @@ void ConnectionState::ThrowIfBusy(const ClientContext &context) const {
 		throw ResourceInUseException(CALLBACK_ERROR);
 	}
 	if (busy) {
-		throw ResourceInUseException(BUSY_ERROR);
+		ThrowBusy();
 	}
 }
 
@@ -340,14 +344,9 @@ ConnectionClaim Session::Claim() {
 	return ConnectionClaim(*state, *context);
 }
 
-shared_ptr<Result> Session::Submit(unique_ptr<SQLStatement> statement, identifier_map_t<BoundParameterData> &values,
-                                   const Format &format) {
-	auto claim = Claim();
-	return SubmitClaimed(std::move(statement), values, format);
-}
-
 shared_ptr<Result> Session::SubmitClaimed(unique_ptr<SQLStatement> statement,
                                           identifier_map_t<BoundParameterData> &values, const Format &format) {
+	D_ASSERT(state->Claimed());
 	auto written_type = statement->type;
 	registry->CountStatement(written_type);
 	auto handle = context->Submit(std::move(statement), values, QueryParameters(format.EngineFormat()));
@@ -357,9 +356,9 @@ shared_ptr<Result> Session::SubmitClaimed(unique_ptr<SQLStatement> statement,
 	return result;
 }
 
-shared_ptr<Result> Session::Submit(Prepared &prepared, identifier_map_t<BoundParameterData> &values,
-                                   const Format &format) {
-	ConnectionClaim claim(*state, *context);
+shared_ptr<Result> Session::SubmitClaimed(Prepared &prepared, identifier_map_t<BoundParameterData> &values,
+                                          const Format &format) {
+	D_ASSERT(state->Claimed());
 	registry->CountStatement(prepared.statement->GetStatementType());
 	unique_ptr<QueryResult> handle;
 	{
@@ -372,8 +371,9 @@ shared_ptr<Result> Session::Submit(Prepared &prepared, identifier_map_t<BoundPar
 	return result;
 }
 
-shared_ptr<Result> Session::Submit(unique_ptr<ColumnDataCollection> collection, const vector<Identifier> &names,
-                                   const Format &format) {
+shared_ptr<Result> Session::SubmitClaimed(unique_ptr<ColumnDataCollection> collection, const vector<Identifier> &names,
+                                          const Format &format) {
+	D_ASSERT(state->Claimed());
 	// The binder rejects duplicate column names; the caller keeps the originals
 	auto deduplicated_names = names;
 	QueryResult::DeduplicateColumns(deduplicated_names);
@@ -386,7 +386,7 @@ shared_ptr<Result> Session::Submit(unique_ptr<ColumnDataCollection> collection,
 	auto select = make_uniq<SelectStatement>();
 	select->node = std::move(select_node);
 	identifier_map_t<BoundParameterData> no_values;
-	return Submit(std::move(select), no_values, format);
+	return SubmitClaimed(std::move(select), no_values, format);
 }
 
 StatementSignature Session::Bind(unique_ptr<SQLStatement> statement, ResolvedTables &tables) {
@@ -557,6 +557,10 @@ shared_ptr<Result> StatementGroup::Choose(shared_ptr<Result> last) {
 
 void StatementGroup::RefuseParameters() {
 	auto claim = session.Claim();
+	RefuseParametersClaimed();
+}
+
+void StatementGroup::RefuseParametersClaimed() {
 	Fail(ErrorData(InvalidInputException("Parameters were given, but the last statement of the query takes none")));
 }
 

src/engine/result.cpp

@@ -218,12 +218,6 @@ shared_ptr<ClientContext> Result::TryGetContext() const {
 	return context;
 }
 
-void Result::Retain(const InterruptCheck &check) {
-	DUCKDB_PY_ASSERT_GIL_RELEASED();
-	auto claim = ClaimWhileRunning();
-	RetainClaimed(check);
-}
-
 void Result::CloseAfterFailure(std::exception_ptr error) noexcept {
 	DUCKDB_PY_ASSERT_GIL_RELEASED();
 	{
@@ -271,6 +265,7 @@ void Result::RetainClaimed(const InterruptCheck &check) {
 	if (state != State::PENDING) {
 		return;
 	}
+	D_ASSERT(connection_state->Claimed());
 	try {
 		Complete(*handle, check);
 		CaptureMetadata();
@@ -281,17 +276,12 @@ void Result::RetainClaimed(const InterruptCheck &check) {
 	state = State::RETAINED;
 }
 
-void Result::Settle(const InterruptCheck &check) {
-	DUCKDB_PY_ASSERT_GIL_RELEASED();
-	auto claim = ClaimWhileRunning();
-	SettleClaimed(check);
-}
-
 void Result::SettleClaimed(const InterruptCheck &check) {
 	DUCKDB_PY_ASSERT_GIL_RELEASED();
 	if (state != State::PENDING) {
 		return;
 	}
+	D_ASSERT(connection_state->Claimed());
 	if (CompletesBeforeReturning()) {
 		RetainClaimed(check);
 		return;
@@ -375,10 +365,16 @@ unique_ptr<ArrowArrayWrapper> Result::FetchArray(const InterruptCheck &check) {
 
 unique_ptr<ColumnDataCollection> Result::TakeChunks(const InterruptCheck &check) {
 	DUCKDB_PY_ASSERT_GIL_RELEASED();
-	D_ASSERT(!format.IsArrow());
 	ThrowIfFailed();
 	// Claimed only while the query can still run engine work: rows already in the client need no connection
 	auto claim = ClaimWhileRunning();
+	return TakeChunksClaimed(check);
+}
+
+unique_ptr<ColumnDataCollection> Result::TakeChunksClaimed(const InterruptCheck &check) {
+	DUCKDB_PY_ASSERT_GIL_RELEASED();
+	D_ASSERT(!format.IsArrow());
+	ThrowIfFailed();
 	if (state == State::PENDING) {
 		RetainClaimed(check);
 	}
@@ -448,9 +444,13 @@ void Result::CopyArrowSchema(ArrowSchema &out) const {
 }
 
 void Result::BuildArrowSchema(const vector<string> &schema_names, ArrowSchema &out) const {
+	ConnectionClaim claim(*connection_state, *TryGetContext());
+	BuildArrowSchemaClaimed(schema_names, out);
+}
+
+void Result::BuildArrowSchemaClaimed(const vector<string> &schema_names, ArrowSchema &out) const {
 	auto context = TryGetContext();
 	auto properties = client_properties;
-	ConnectionClaim claim(*connection_state, *context);
 	context->RunFunctionInTransaction(
 	    [&]() { ArrowConverter::ToArrowSchema(&out, Types(), schema_names, properties); });
 }

src/include/duckdb_python/engine/connection.hpp

@@ -86,6 +86,10 @@ class ConnectionState : public ClientContextState {
 	bool AdmittedDuringExit() const {
 		return admitted_during_exit;
 	}
+	//! Whether a call holds the connection's claim right now
+	bool Claimed() const {
+		return busy;
+	}
 
 private:
 	friend class ConnectionClaim;
@@ -168,6 +172,9 @@ class SourceRegistry : public ObjectCacheEntry {
 	std::deque<string> released_order;
 };
 
+//! The refusal of a call that finds the connection taken by another call
+[[noreturn]] void ThrowBusy();
+
 //! Holds the connection for one stretch of engine work. A second claim, and any claim from a callback of the
 //! connection's queries, is refused instead of waiting: the work it would wait for may be waiting for it.
 class ConnectionClaim {
@@ -205,15 +212,16 @@ class Session {
 		return *context;
 	}
 	Prepared Prepare(unique_ptr<SQLStatement> statement);
-	shared_ptr<Result> Submit(unique_ptr<SQLStatement> statement, identifier_map_t<BoundParameterData> &values,
-	                          const Format &format);
-	//! Submit for a caller that already holds the connection's claim
+	//! Submits a statement whose result the caller reads later, for a caller that holds the connection's claim;
+	//! the caller keeps its claim until the result is settled or retained, so no other call finds the statement
+	//! pending. Run and RunClaimed below submit and complete a statement in one go.
 	shared_ptr<Result> SubmitClaimed(unique_ptr<SQLStatement> statement, identifier_map_t<BoundParameterData> &values,
 	                                 const Format &format);
-	shared_ptr<Result> Submit(Prepared &prepared, identifier_map_t<BoundParameterData> &values, const Format &format);
+	shared_ptr<Result> SubmitClaimed(Prepared &prepared, identifier_map_t<BoundParameterData> &values,
+	                                 const Format &format);
 	//! Scans a collection, for example to read rows a result already holds in another format
-	shared_ptr<Result> Submit(unique_ptr<ColumnDataCollection> collection, const vector<Identifier> &names,
-	                          const Format &format);
+	shared_ptr<Result> SubmitClaimed(unique_ptr<ColumnDataCollection> collection, const vector<Identifier> &names,
+	                                 const Format &format);
 	//! Binds a statement without running it, keeping the tables its replacement scans resolved in tables
 	StatementSignature Bind(unique_ptr<SQLStatement> statement, ResolvedTables &tables);
 	//! Binds a statement with values for its parameters, as the engine rebinds a prepared statement for the values it
@@ -267,6 +275,8 @@ class StatementGroup {
 	shared_ptr<Result> Choose(shared_ptr<Result> last);
 	//! Throws for parameters given to a group whose text after the statements that ran held no statement
 	[[noreturn]] void RefuseParameters();
+	//! RefuseParameters for a caller that already holds the connection's claim
+	[[noreturn]] void RefuseParametersClaimed();
 
 private:
 	bool HasMore();

src/include/duckdb_python/engine/result.hpp

@@ -52,9 +52,10 @@ class Format {
 class ConnectionState;
 class ConnectionClaim;
 
-//! A submitted statement. Retention is settled by the first consumer: a fetch opens a stream, Retain and the
-//! Take calls keep every row. Every call that drives the query claims the connection, takes the client context
-//! lock and may wait for running tasks, so callers must not hold a lock those tasks need.
+//! A submitted statement. Retention is settled by the first consumer: a fetch opens a stream, RetainClaimed and
+//! the Take calls keep every row. Every call that drives the query runs under the connection's claim, its own or
+//! the caller's, takes the client context lock and may wait for running tasks, so callers must not hold a lock
+//! those tasks need.
 class Result : public enable_shared_from_this<Result> {
 public:
 	//! written_type is the submitted statement's own type, since the handle of a prepared execution reports EXECUTE
@@ -105,14 +106,11 @@ class Result : public enable_shared_from_this<Result> {
 	shared_ptr<ClientContext> TryGetContext() const;
 	//! The database whose buffers hold the rows, null once it is gone. Weak for the same reason as the context.
 	shared_ptr<DatabaseInstance> TryGetDatabase() const;
-	//! Keeps every row: drives the query to its end unless it already ended. A failure ends the query.
-	void Retain(const InterruptCheck &check);
-	//! Retain for a caller that already holds the connection's claim. A failure ends the query.
+	//! Keeps every row: drives the query to its end unless it already ended, for a caller that holds the
+	//! connection's claim. A failure ends the query.
 	void RetainClaimed(const InterruptCheck &check);
-	//! Runs the statement until its first unit is ready or it ends, without deciding how its rows are read,
-	//! and throws the error it ended with. A failure ends the query.
-	void Settle(const InterruptCheck &check);
-	//! Settle for a caller that already holds the connection's claim. A failure ends the query.
+	//! Runs the statement until its first unit is ready or it ends, without deciding how its rows are read, and
+	//! throws the error it ended with, for a caller that holds the connection's claim. A failure ends the query.
 	void SettleClaimed(const InterruptCheck &check);
 	//! Rows of a retained result. Materializes it first.
 	idx_t RowCount(const InterruptCheck &check);
@@ -122,12 +120,16 @@ class Result : public enable_shared_from_this<Result> {
 	unique_ptr<ArrowArrayWrapper> FetchArray(const InterruptCheck &check);
 	//! Every remaining row, the result is empty afterwards
 	unique_ptr<ColumnDataCollection> TakeChunks(const InterruptCheck &check);
+	//! TakeChunks for a caller that already holds the connection's claim
+	unique_ptr<ColumnDataCollection> TakeChunksClaimed(const InterruptCheck &check);
 	//! Every remaining array, the result is empty afterwards
 	vector<unique_ptr<ArrowArrayWrapper>> TakeArrays(const InterruptCheck &check);
 	//! A copy of the result's Arrow schema, owned by the caller. Arrow results only.
 	void CopyArrowSchema(ArrowSchema &out) const;
 	//! An Arrow schema for the result's types under the given names, built in a transaction on its connection
 	void BuildArrowSchema(const vector<string> &names, ArrowSchema &out) const;
+	//! BuildArrowSchema for a caller that already holds the connection's claim
+	void BuildArrowSchemaClaimed(const vector<string> &names, ArrowSchema &out) const;
 	//! Renders the rows of a retained result
 	string ToBox(BoxRendererContext &context, const BoxRendererConfig &config);
 

src/include/duckdb_python/pyconnection/pyconnection.hpp

@@ -7,6 +7,8 @@
 //===----------------------------------------------------------------------===//
 
 #pragma once
+
+#include <thread>
 #include "duckdb.hpp"
 #include "duckdb_python/nb/casters.hpp"
 #include "duckdb_python/import_cache/python_import_cache.hpp"
@@ -88,8 +90,9 @@ class DbapiResult {
 	void Supersede();
 	//! Every row the result still holds, under names, which leaves it read out as a consuming fetch does; null
 	//! when the rows were used up already. The rows are gone even when the take fails, since it ran the
-	//! statement to its end or ended it.
-	unique_ptr<ColumnDataCollection> TakeRows(vector<Identifier> &names);
+	//! statement to its end or ended it. The caller holds the connection's claim, so nothing can refuse the take
+	//! once it has begun.
+	unique_ptr<ColumnDataCollection> TakeRowsClaimed(vector<Identifier> &names);
 	shared_ptr<engine::Result> EngineResult() const {
 		return result ? result->EngineResult() : nullptr;
 	}
@@ -266,38 +269,77 @@ struct DuckDBPyConnection : public std::enable_shared_from_this<DuckDBPyConnecti
 	};
 
 public:
-	// RAII guard for the connection mutex (see py_connection_lock below). Constructing
-	// one releases the GIL while waiting for the mutex and reacquires it before
-	// returning, so callers always come out of the constructor with the GIL held
-	// and the mutex locked. The mutex is released when the guard goes out of scope.
-	// Holding the GIL while blocked on this mutex would deadlock against a thread
-	// that holds the mutex and is mid-way through a GIL-releasing native call —
-	// see duckdb-python#435.
+	//! What a call does about other calls on the connection mutex
+	enum class IfHeld : uint8_t {
+		//! Waits for the mutex, with the GIL released: the holder may be inside a native call that needs the GIL back
+		WAIT,
+		//! For the calls that submit a statement. One of them is in flight at a time: a second one arriving is refused
+		//! as busy at once, before any wait, since the first may be in the stretch of its call outside any claim
+		//! (parsing, superseding the old result, converting its parameters, installing the result), where a caller
+		//! that passed the busy check would otherwise wait for the whole call and then replace its result. Behind
+		//! any other holder of the mutex it waits as every call does
+		REFUSE
+	};
+
+	//! Holds the connection mutex (py_connection_lock below) for one call; the GIL is held on return. The mutex is
+	//! recursive, so a call from the thread that holds it enters once it passed the busy check; a submitting call
+	//! from that thread also needs the mark to be its own, which it is unless another thread's submitting call
+	//! arrived first and waits for the mutex
 	class ConnectionLockGuard {
 	public:
-		explicit ConnectionLockGuard(DuckDBPyConnection &conn) : lock_(conn.py_connection_lock, std::defer_lock) {
+		explicit ConnectionLockGuard(DuckDBPyConnection &conn_p, IfHeld if_held = IfHeld::WAIT)
+		    : conn(conn_p), lock_(conn_p.py_connection_lock, std::defer_lock) {
 			DUCKDB_PY_ASSERT_GIL_HELD();
 			// The holder of the mutex may be running a query that is waiting for this very call
 			conn.con.ThrowIfBusy();
+			if (if_held == IfHeld::REFUSE) {
+				MarkSubmitting();
+			}
 			nb::gil_scoped_release release;
 			DUCKDB_PY_ASSERT_GIL_RELEASED();
 			lock_.lock();
 		}
+		~ConnectionLockGuard() {
+			DropSubmittingMark();
+		}
+
+		//! Ends the stretch during which other submitting calls are refused, for a call whose last claim is over
+		//! while it still holds the mutex, so the rest of the call is waited for like any other holder
+		void DropSubmittingMark() {
+			if (marks_submitting_call) {
+				conn.submitting_thread.store(std::thread::id());
+				marks_submitting_call = false;
+			}
+		}
 
 	private:
+		//! Decided on arrival, before any wait, so the loser never waits for the winner's call; a nested call of the
+		//! submitting thread itself passes, as it does the recursive mutex
+		void MarkSubmitting() {
+			std::thread::id nobody;
+			auto self = std::this_thread::get_id();
+			if (conn.submitting_thread.compare_exchange_strong(nobody, self)) {
+				marks_submitting_call = true;
+				return;
+			}
+			if (nobody != self) {
+				engine::ThrowBusy();
+			}
+		}
+
+		DuckDBPyConnection &conn;
 		std::unique_lock<std::recursive_mutex> lock_;
+		bool marks_submitting_call = false;
 	};
 
 	ConnectionGuard con;
 	Cursors cursors;
-	// Recursive so that the outer lock taken at the top of execute/fetch
-	// methods (while still holding the GIL) does not deadlock against the
-	// inner lock taken by PrepareQuery / ExecuteInternal /
-	// PrepareAndSubmitInternal (after releasing the GIL). Serialises every
-	// path that touches `con.result` so concurrent calls on a single
-	// DuckDBPyConnection cannot dereference an already-freed result — see
-	// duckdb-python#435.
+	//! Serialises every path that touches con.result, so concurrent calls on one connection cannot dereference a
+	//! freed result (duckdb-python#435). Recursive, so the lock taken at the top of execute and the fetches does not
+	//! deadlock against the one PrepareQuery and ExecuteInternal take after releasing the GIL
 	std::recursive_mutex py_connection_lock;
+	//! The thread whose call that submits a statement (IfHeld::REFUSE) is in flight, or no thread
+	std::atomic<std::thread::id> submitting_thread {std::thread::id()};
 	//! The database's store of file-like objects for reading, kept once this connection has found it
 	shared_ptr<ObjectStore> object_store;
 	case_insensitive_map_t<unique_ptr<ExternalDependency>> registered_functions;
@@ -372,12 +414,11 @@ struct DuckDBPyConnection : public std::enable_shared_from_this<DuckDBPyConnecti
 	engine::Prepared PrepareQuery(unique_ptr<SQLStatement> statement);
 	shared_ptr<engine::Result> ExecuteInternal(engine::Prepared &prep, nb::object params,
 	                                           engine::StatementGroup &statements);
-	//! Binds the parameters and submits the statement. The result is returned undriven.
-	shared_ptr<engine::Result> PrepareAndSubmitInternal(unique_ptr<SQLStatement> statement,
-	                                                    nb::object params = nb::list());
 	engine::Session EngineSession();
 
 	std::shared_ptr<DuckDBPyConnection> Execute(const nb::object &query, nb::object params = nb::list());
+	//! Makes the result the connection's DB-API result, with the GIL released; false when there is none
+	bool RegisterResult(shared_ptr<engine::Result> res);
 	std::shared_ptr<DuckDBPyConnection> ExecuteFromString(const string &query);
 
 	std::shared_ptr<DuckDBPyConnection> Append(const string &name, const PandasDataFrame &value, bool by_name);