From 4c8de5e3b971292b51dcc95c9ef375ce3f20edb1 Mon Sep 17 00:00:00 2001 From: Tomasz Kalinowski Date: Sat, 26 Sep 2026 15:57:47 -0400 Subject: [PATCH] Execute multi-statement pragmas after preceding statements --- .../pyconnection/pyconnection.hpp | 4 +- src/pyconnection.cpp | 67 ++++++++----------- tests/fast/test_multi_statement.py | 11 +++ 3 files changed, 41 insertions(+), 41 deletions(-) diff --git a/src/include/duckdb_python/pyconnection/pyconnection.hpp b/src/include/duckdb_python/pyconnection/pyconnection.hpp index 67e0fd8c..aacf3e45 100644 --- a/src/include/duckdb_python/pyconnection/pyconnection.hpp +++ b/src/include/duckdb_python/pyconnection/pyconnection.hpp @@ -250,7 +250,7 @@ struct DuckDBPyConnection : public std::enable_shared_from_this ExecuteMany(const nb::object &query, nb::object params = nb::list()); - void ExecuteImmediately(vector> statements); + void ExecuteImmediately(unique_ptr statement); unique_ptr PrepareQuery(unique_ptr statement); unique_ptr ExecuteInternal(PreparedStatement &prep, nb::object params = nb::list()); //! Binds the parameters and submits the statement. The handle is returned undriven. @@ -371,7 +371,7 @@ struct DuckDBPyConnection : public std::enable_shared_from_this> GetStatements(const nb::object &query); + unique_ptr GetLastStatement(const nb::object &query); static void DetectEnvironment(); }; diff --git a/src/pyconnection.cpp b/src/pyconnection.cpp index 034f4a6b..0289c52c 100644 --- a/src/pyconnection.cpp +++ b/src/pyconnection.cpp @@ -11,6 +11,7 @@ #include "duckdb/main/db_instance_cache.hpp" #include "duckdb/main/extension_helper.hpp" #include "duckdb/main/prepared_statement.hpp" +#include "duckdb/main/statement_iterator.hpp" #include "duckdb/main/relation/read_csv_relation.hpp" #include "duckdb/main/relation/read_json_relation.hpp" #include "duckdb/main/relation/value_relation.hpp" @@ -514,17 +515,12 @@ std::shared_ptr DuckDBPyConnection::ExecuteMany(const nb::ob params_p = nb::list(); } - auto statements = GetStatements(query); - if (statements.empty()) { + auto last_statement = GetLastStatement(query); + if (!last_statement) { // TODO: should we throw? return nullptr; } - - auto last_statement = std::move(statements.back()); - statements.pop_back(); - // First immediately execute any preceding statements (if any) // FIXME: DBAPI says to not accept an 'executemany' call with multiple statements - ExecuteImmediately(std::move(statements)); auto prep = PrepareQuery(std::move(last_statement)); @@ -704,18 +700,26 @@ unique_ptr DuckDBPyConnection::PrepareAndSubmitInternal(unique_ptr< return res; } -vector> DuckDBPyConnection::GetStatements(const nb::object &query) { +unique_ptr DuckDBPyConnection::GetLastStatement(const nb::object &query) { if (nb::isinstance(query)) { auto &statement_obj = nb::cast(query); - vector> result; - result.push_back(statement_obj.GetStatement()); - return result; + return statement_obj.GetStatement(); } if (nb::isinstance(query)) { auto &connection = con.GetConnection(); auto sql_query = nb::cast(nb::str(query)); - auto statements = connection.ExtractStatements(sql_query); - return std::move(statements); + auto iterator = connection.context->IterateStatements(sql_query); + while (iterator.Peek()) { + auto statement = iterator.GetStatement(); + if (!statement) { + continue; + } + if (!iterator.HasMore()) { + return statement; + } + ExecuteImmediately(std::move(statement)); + } + return nullptr; } throw InvalidInputException("Please provide either a DuckDBPyStatement or a string representing the query"); } @@ -729,17 +733,12 @@ std::shared_ptr DuckDBPyConnection::Execute(const nb::object ConnectionLockGuard conn_lock(*this); con.SetResult(nullptr); - auto statements = GetStatements(query); - if (statements.empty()) { + auto last_statement = GetLastStatement(query); + if (!last_statement) { // TODO: should we throw? return nullptr; } - - auto last_statement = std::move(statements.back()); - statements.pop_back(); - // First immediately execute any preceding statements (if any) // FIXME: SQLites implementation says to not accept an 'execute' call with multiple statements - ExecuteImmediately(std::move(statements)); auto res = PrepareAndSubmitInternal(std::move(last_statement), std::move(params)); @@ -1608,22 +1607,17 @@ std::unique_ptr DuckDBPyConnection::ReadCSV(const nb::object & return CreateRelation(read_csv_p->Alias(read_csv.alias)); } -void DuckDBPyConnection::ExecuteImmediately(vector> statements) { +void DuckDBPyConnection::ExecuteImmediately(unique_ptr statement) { auto &connection = con.GetConnection(); D_ASSERT(duckdb::PyUtil::GilCheck()); nb::gil_scoped_release release; - if (statements.empty()) { - return; - } - for (auto &stmt : statements) { - if (!stmt->named_param_map.empty()) { - throw NotImplementedException( - "Prepared parameters are only supported for the last statement, please split your query up into " - "separate 'execute' calls if you want to use prepared parameters"); - } - auto res = connection.Submit(std::move(stmt)); - CompleteQuery(*res); + if (!statement->named_param_map.empty()) { + throw NotImplementedException( + "Prepared parameters are only supported for the last statement, please split your query up into " + "separate 'execute' calls if you want to use prepared parameters"); } + auto res = connection.Submit(std::move(statement)); + CompleteQuery(*res); } std::unique_ptr DuckDBPyConnection::RunQuery(const nb::object &query, string alias, @@ -1633,17 +1627,12 @@ std::unique_ptr DuckDBPyConnection::RunQuery(const nb::object alias = "unnamed_relation_" + StringUtil::GenerateRandomName(16); } - auto statements = GetStatements(query); - if (statements.empty()) { + auto last_statement = GetLastStatement(query); + if (!last_statement) { // TODO: should we throw? return nullptr; } - auto last_statement = std::move(statements.back()); - statements.pop_back(); - // First immediately execute any preceding statements (if any) - ExecuteImmediately(std::move(statements)); - // Attempt to create a Relation for lazy execution if possible shared_ptr relation; bool has_params = !nb::none().is(params) && nb::len(params) > 0; diff --git a/tests/fast/test_multi_statement.py b/tests/fast/test_multi_statement.py index 76ac0b4b..4a3a844a 100644 --- a/tests/fast/test_multi_statement.py +++ b/tests/fast/test_multi_statement.py @@ -6,6 +6,17 @@ class TestMultiStatement: + def test_pragma_expansion_follows_previous_statements(self, tmp_path): + con = duckdb.connect(":memory:") + export_location = tmp_path / "export" + con.execute("CREATE TABLE integers AS SELECT 42 AS i") + + con.execute( + f"EXPORT DATABASE '{export_location}'; DROP TABLE integers; PRAGMA import_database('{export_location}')" + ) + + assert con.execute("SELECT i FROM integers").fetchall() == [(42,)] + def test_multi_statement(self, duckdb_cursor): con = duckdb.connect(":memory:")