Return real row counts from cursor.rowcount for CHANGED_ROWS statements

src/include/duckdb_python/pyconnection/pyconnection.hpp

@@ -321,7 +321,7 @@ struct DuckDBPyConnection : public std::enable_shared_from_this<DuckDBPyConnecti
 
 	Optional<nb::list> GetDescription();
 
-	int GetRowcount();
+	int64_t GetRowcount();
 
 	// these should be functions on the result but well
 	Optional<nb::tuple> FetchOne();

src/include/duckdb_python/pyrelation.hpp

@@ -30,6 +30,15 @@ struct DuckDBPyRelation {
 
 	nb::list Description();
 
+	//! Cached at construction time (see the DuckDBPyResult constructor overload) so that it survives
+	//! the Fetch*() methods below, several of which null out `result` once fully consumed. Note this
+	//! is only ever non-default for relations built directly from a DuckDBPyResult (the connection's
+	//! execute()/executemany() path) - relations built from a lazy Relation (rel != nullptr, executed
+	//! via ExecuteOrThrow) are always SELECT-shaped and so never produce a CHANGED_ROWS value anyway.
+	int64_t GetRowcount() const {
+		return row_changes;
+	}
+
 	void Close();
 
 	std::unique_ptr<DuckDBPyRelation> GetAttribute(const string &name);
@@ -300,6 +309,8 @@ struct DuckDBPyRelation {
 	vector<string> names;
 	std::shared_ptr<DuckDBPyResult> result;
 	std::string rendered_result;
+	//! Cached row-changed count - see GetRowcount().
+	int64_t row_changes = -1;
 };
 
 } // namespace duckdb

src/include/duckdb_python/pyresult.hpp

@@ -59,6 +59,12 @@ struct DuckDBPyResult {
 
 	ClientProperties GetClientProperties();
 
+	//! Number of rows changed by the last CHANGED_ROWS-returning statement (INSERT/UPDATE/DELETE/...).
+	//! Returns -1 when not applicable/unknown, as permitted by the DB-API 2.0 spec for 'rowcount'.
+	int64_t GetRowcount() const {
+		return row_changes;
+	}
+
 private:
 	void FillNumpy(nb::dict &res, idx_t col_idx, NumpyResultConversion &conversion, const char *name);
 
@@ -84,6 +90,12 @@ struct DuckDBPyResult {
 	duckdb::pyarrow::Table MaterializedResultToArrowTable(const ArrowSchema &arrow_schema, idx_t rows_per_batch);
 	ArrowArrayStream FetchArrowArrayStream(idx_t rows_per_batch);
 
+	//! Computes the CHANGED_ROWS value (if any) up front, before any Fetch call has had a chance to
+	//! consume it, so that later fetch calls do not affect what GetRowcount() reports. Materializes a
+	//! streaming result if necessary; this is cheap since CHANGED_ROWS results are always exactly one
+	//! already-computed row.
+	int64_t ComputeRowChanges();
+
 private:
 	idx_t chunk_offset = 0;
 
@@ -97,6 +109,8 @@ struct DuckDBPyResult {
 	// Holds the categorical type of Categorical/ENUM types
 	unordered_map<idx_t, nb::object> categories_type;
 	bool result_closed = false;
+	//! Cached by ComputeRowChanges() at construction time - see GetRowcount().
+	int64_t row_changes = -1;
 };
 
 } // namespace duckdb

src/pyconnection.cpp

@@ -1916,8 +1916,13 @@ Optional<nb::list> DuckDBPyConnection::GetDescription() {
 	return result.Description();
 }
 
-int DuckDBPyConnection::GetRowcount() {
-	return -1;
+int64_t DuckDBPyConnection::GetRowcount() {
+	ConnectionLockGuard conn_lock(*this);
+	if (!con.HasResult()) {
+		return -1;
+	}
+	auto &result = con.GetResult();
+	return result.GetRowcount();
 }
 
 void DuckDBPyConnection::Close() {

src/pyrelation.cpp

@@ -71,6 +71,7 @@ DuckDBPyRelation::DuckDBPyRelation(std::shared_ptr<DuckDBPyResult> result_p)
 	this->executed = true;
 	this->types = result->GetTypes();
 	this->names = result->GetNames();
+	this->row_changes = result->GetRowcount();
 }
 
 std::unique_ptr<DuckDBPyRelation> DuckDBPyRelation::ProjectFromExpression(const string &expression) {

src/pyresult.cpp

@@ -33,6 +33,9 @@ DuckDBPyResult::DuckDBPyResult(unique_ptr<QueryResult> result_p) : result(std::m
 	if (!result) {
 		throw InternalException("PyResult created without a result object");
 	}
+	// Must happen before any Fetch call (from this object or a caller holding a reference to the
+	// same underlying QueryResult) has a chance to consume the single row we're reading here.
+	row_changes = ComputeRowChanges();
 }
 
 const vector<Identifier> &DuckDBPyResult::ResultNames() const {
@@ -72,6 +75,43 @@ const vector<LogicalType> &DuckDBPyResult::GetTypes() {
 	return result->GetTypes();
 }
 
+int64_t DuckDBPyResult::ComputeRowChanges() {
+	if (!result || result->HasError()) {
+		return -1;
+	}
+	if (result->properties.return_type != StatementReturnType::CHANGED_ROWS) {
+		// The row count of a SELECT (or a statement that returns nothing) is not known without fully
+		// consuming the result - report -1, as permitted by the DB-API 2.0 spec for 'rowcount'.
+		return -1;
+	}
+	if (result->type == QueryResultType::STREAM_RESULT) {
+		// CHANGED_ROWS statements always produce a single already-computed row, so materializing here
+		// does not trigger any additional query execution - it just changes the in-memory representation.
+		// Still release the GIL around it though: Materialize() drives the same native Fetch() machinery
+		// as any other result consumption, and other fetch paths (e.g. Fetchone()) release the GIL around it.
+		unique_ptr<MaterializedQueryResult> materialized;
+		{
+			D_ASSERT(duckdb::PyUtil::GilCheck());
+			nb::gil_scoped_release release;
+			auto &stream_result = result->Cast<StreamQueryResult>();
+			materialized = stream_result.Materialize();
+		}
+		if (!materialized || materialized->HasError()) {
+			return -1;
+		}
+		result = std::move(materialized);
+	}
+	if (result->type != QueryResultType::MATERIALIZED_RESULT) {
+		// e.g. an Arrow result - can't peek at the value without disturbing it.
+		return -1;
+	}
+	auto &materialized_result = result->Cast<MaterializedQueryResult>();
+	if (materialized_result.RowCount() != 1 || materialized_result.ColumnCount() != 1) {
+		return -1;
+	}
+	return materialized_result.GetValue(0, 0).GetValue<int64_t>();
+}
+
 unique_ptr<DataChunk> DuckDBPyResult::FetchChunk() {
 	if (!result) {
 		throw InternalException("FetchChunk called without a result object");

tests/fast/api/test_dbapi10.py

@@ -56,3 +56,104 @@ def test_none_description(self, duckdb_empty_cursor):
 class TestCursorRowcount:
     def test_rowcount(self, duckdb_cursor):
         assert duckdb_cursor.rowcount == -1
+
+    def test_rowcount_no_query_yet(self, duckdb_cursor):
+        assert duckdb_cursor.rowcount == -1
+
+    def test_rowcount_insert(self, duckdb_cursor):
+        duckdb_cursor.execute("CREATE TABLE t (i INTEGER)")
+        duckdb_cursor.execute("INSERT INTO t VALUES (1), (2), (3)")
+        assert duckdb_cursor.rowcount == 3
+
+    def test_rowcount_insert_select(self, duckdb_cursor):
+        duckdb_cursor.execute("CREATE TABLE t (i INTEGER)")
+        duckdb_cursor.execute("INSERT INTO t SELECT * FROM range(100)")
+        assert duckdb_cursor.rowcount == 100
+
+    def test_rowcount_update(self, duckdb_cursor):
+        duckdb_cursor.execute("CREATE TABLE t AS SELECT * FROM range(10) t(i)")
+        duckdb_cursor.execute("UPDATE t SET i = i + 1 WHERE i < 4")
+        assert duckdb_cursor.rowcount == 4
+
+    def test_rowcount_delete(self, duckdb_cursor):
+        duckdb_cursor.execute("CREATE TABLE t AS SELECT * FROM range(10) t(i)")
+        duckdb_cursor.execute("DELETE FROM t WHERE i < 3")
+        assert duckdb_cursor.rowcount == 3
+
+    def test_rowcount_create_table_as(self, duckdb_cursor):
+        duckdb_cursor.execute("CREATE TABLE t AS SELECT * FROM range(42) t(i)")
+        assert duckdb_cursor.rowcount == 42
+
+    def test_rowcount_select_is_unknown(self, duckdb_cursor):
+        # Matches DB-API 2.0: rowcount is -1 when it can't be determined without consuming the result.
+        duckdb_cursor.execute("SELECT * FROM range(10)")
+        assert duckdb_cursor.rowcount == -1
+        duckdb_cursor.fetchall()
+        assert duckdb_cursor.rowcount == -1
+
+    def test_rowcount_reset_by_next_execute(self, duckdb_cursor):
+        duckdb_cursor.execute("CREATE TABLE t (i INTEGER)")
+        duckdb_cursor.execute("INSERT INTO t VALUES (1), (2)")
+        assert duckdb_cursor.rowcount == 2
+        duckdb_cursor.execute("SELECT * FROM t")
+        assert duckdb_cursor.rowcount == -1
+
+    def test_rowcount_does_not_disturb_fetch(self, duckdb_cursor):
+        # Accessing rowcount must not consume the single-row result that fetchone() also reads.
+        duckdb_cursor.execute("CREATE TABLE t (i INTEGER)")
+        duckdb_cursor.execute("INSERT INTO t VALUES (1), (2), (3)")
+        assert duckdb_cursor.rowcount == 3
+        assert duckdb_cursor.fetchone() == (3,)
+
+    def test_rowcount_ddl_is_unknown(self, duckdb_cursor):
+        duckdb_cursor.execute("CREATE TABLE t (i INTEGER)")
+        assert duckdb_cursor.rowcount == -1
+
+    def test_rowcount_after_fetchall(self, duckdb_cursor):
+        # Regression test: rowcount must survive full consumption of the result via fetchall(),
+        # not just a bare execute() with no fetch at all.
+        duckdb_cursor.execute("CREATE TABLE t (i INTEGER)")
+        duckdb_cursor.execute("INSERT INTO t VALUES (1), (2), (3), (4)")
+        assert duckdb_cursor.fetchall() == [(4,)]
+        assert duckdb_cursor.rowcount == 4
+
+    def test_rowcount_after_fetchone(self, duckdb_cursor):
+        duckdb_cursor.execute("CREATE TABLE t (i INTEGER)")
+        duckdb_cursor.execute("INSERT INTO t VALUES (1), (2), (3), (4)")
+        assert duckdb_cursor.fetchone() == (4,)
+        assert duckdb_cursor.rowcount == 4
+
+    def test_rowcount_after_fetchmany(self, duckdb_cursor):
+        duckdb_cursor.execute("CREATE TABLE t (i INTEGER)")
+        duckdb_cursor.execute("INSERT INTO t VALUES (1), (2), (3), (4)")
+        assert duckdb_cursor.fetchmany(1) == [(4,)]
+        assert duckdb_cursor.rowcount == 4
+
+    def test_rowcount_after_fetchdf(self, duckdb_cursor):
+        duckdb_cursor.execute("CREATE TABLE t (i INTEGER)")
+        duckdb_cursor.execute("INSERT INTO t VALUES (1), (2), (3), (4)")
+        duckdb_cursor.fetchdf()
+        assert duckdb_cursor.rowcount == 4
+
+    def test_rowcount_after_fetchnumpy(self, duckdb_cursor):
+        duckdb_cursor.execute("CREATE TABLE t (i INTEGER)")
+        duckdb_cursor.execute("INSERT INTO t VALUES (1), (2), (3), (4)")
+        duckdb_cursor.fetchnumpy()
+        assert duckdb_cursor.rowcount == 4
+
+    def test_rowcount_after_to_arrow_table(self, duckdb_cursor):
+        pytest.importorskip("pyarrow")
+        duckdb_cursor.execute("CREATE TABLE t (i INTEGER)")
+        duckdb_cursor.execute("INSERT INTO t VALUES (1), (2), (3), (4)")
+        duckdb_cursor.to_arrow_table()
+        assert duckdb_cursor.rowcount == 4
+
+    def test_rowcount_executemany_reflects_last_statement_only(self, duckdb_cursor):
+        # Documents current (limited) behavior: executemany() only keeps the QueryResult of the last
+        # parameter set it executes, so rowcount reflects that one statement rather than the total
+        # number of rows affected across all parameter sets. Not the DB-API-idiomatic answer, but
+        # pinned here so a future change to this behavior is a deliberate, visible decision.
+        duckdb_cursor.execute("CREATE TABLE t (i INTEGER)")
+        duckdb_cursor.executemany("INSERT INTO t VALUES (?)", [(1,), (2,), (3,)])
+        assert duckdb_cursor.table("t").fetchall() == [(1,), (2,), (3,)]
+        assert duckdb_cursor.rowcount == 1