Merge pull request #310 from daniellietz/add-transaction-param-execute

Mytherin · web-flow · commit 9e9e30c8a4a2 · 2025-05-07T23:39:46.000+02:00
Add `use_transaction` parameter for `postgres_query()` and `postgres_execute()`
diff --git a/src/include/postgres_scanner.hpp b/src/include/postgres_scanner.hpp
@@ -39,6 +39,7 @@ struct PostgresBindData : public FunctionData {
 	bool can_use_main_thread = true;
 	bool read_only = true;
 	bool emit_ctid = false;
+	bool use_transaction = true;
 	idx_t max_threads = 1;
 
 public:
diff --git a/src/include/storage/postgres_transaction.hpp b/src/include/storage/postgres_transaction.hpp
@@ -28,9 +28,12 @@ class PostgresTransaction : public Transaction {
 	void Commit();
 	void Rollback();
 
+	PostgresConnection &GetConnectionWithoutTransaction();
 	PostgresConnection &GetConnection();
+
 	string GetDSN();
 	unique_ptr<PostgresResult> Query(const string &query);
+	unique_ptr<PostgresResult> QueryWithoutTransaction(const string &query);
 	vector<unique_ptr<PostgresResult>> ExecuteQueries(const string &queries);
 	static PostgresTransaction &Get(ClientContext &context, Catalog &catalog);
 
diff --git a/src/postgres_execute.cpp b/src/postgres_execute.cpp
@@ -10,13 +10,14 @@
 namespace duckdb {
 
 struct PGExecuteBindData : public TableFunctionData {
-	explicit PGExecuteBindData(PostgresCatalog &pg_catalog, string query_p)
-	    : pg_catalog(pg_catalog), query(std::move(query_p)) {
+	explicit PGExecuteBindData(PostgresCatalog &pg_catalog, string query_p, bool use_transaction)
+	    : pg_catalog(pg_catalog), query(std::move(query_p)), use_transaction(use_transaction) {
 	}
 
 	bool finished = false;
 	PostgresCatalog &pg_catalog;
 	string query;
+	bool use_transaction = true;
 };
 
 static duckdb::unique_ptr<FunctionData> PGExecuteBind(ClientContext &context, TableFunctionBindInput &input,
@@ -36,7 +37,15 @@ static duckdb::unique_ptr<FunctionData> PGExecuteBind(ClientContext &context, Ta
 		throw BinderException("Attached database \"%s\" does not refer to a Postgres database", db_name);
 	}
 	auto &pg_catalog = catalog.Cast<PostgresCatalog>();
-	return make_uniq<PGExecuteBindData>(pg_catalog, input.inputs[1].GetValue<string>());
+
+	bool use_transaction = true;
+	for (auto &kv : input.named_parameters) {
+		if (kv.first == "use_transaction") {
+			use_transaction = BooleanValue::Get(kv.second);
+		}
+	}
+
+	return make_uniq<PGExecuteBindData>(pg_catalog, input.inputs[1].GetValue<string>(), use_transaction);
 }
 
 static void PGExecuteFunction(ClientContext &context, TableFunctionInput &data_p, DataChunk &output) {
@@ -45,13 +54,19 @@ static void PGExecuteFunction(ClientContext &context, TableFunctionInput &data_p
 		return;
 	}
 	auto &transaction = Transaction::Get(context, data.pg_catalog).Cast<PostgresTransaction>();
-	transaction.Query(data.query);
+	if (data.use_transaction) {
+		transaction.Query(data.query);
+	} else {
+		transaction.QueryWithoutTransaction(data.query);
+	}
+
 	data.finished = true;
 }
 
 PostgresExecuteFunction::PostgresExecuteFunction()
     : TableFunction("postgres_execute", {LogicalType::VARCHAR, LogicalType::VARCHAR}, PGExecuteFunction,
                     PGExecuteBind) {
+	    named_parameters["use_transaction"] = LogicalType::BOOLEAN;
 }
 
 } // namespace duckdb
diff --git a/src/postgres_query.cpp b/src/postgres_query.cpp
@@ -38,7 +38,16 @@ static unique_ptr<FunctionData> PGQueryBind(ClientContext &context, TableFunctio
 		StringUtil::RTrim(sql);
 	}
 
-	auto &con = transaction.GetConnection();
+	bool use_transaction = true;
+	for (auto &kv : input.named_parameters) {
+		if (kv.first == "use_transaction") {
+			use_transaction = BooleanValue::Get(kv.second);
+		}
+	}
+	result->use_transaction = use_transaction;
+
+	auto &con = use_transaction ? transaction.GetConnection() : transaction.GetConnectionWithoutTransaction();
+
 	auto conn = con.GetConn();
 	// prepare execution of the query to figure out the result types and names
 	auto prepared = PQprepare(conn, "", sql.c_str(), 0, nullptr);
@@ -87,6 +96,7 @@ static unique_ptr<FunctionData> PGQueryBind(ClientContext &context, TableFunctio
 
 PostgresQueryFunction::PostgresQueryFunction()
     : TableFunction("postgres_query", {LogicalType::VARCHAR, LogicalType::VARCHAR}, nullptr, PGQueryBind) {
+	named_parameters["use_transaction"] = LogicalType::BOOLEAN;
 	PostgresScanFunction scan_function;
 	init_global = scan_function.init_global;
 	init_local = scan_function.init_local;
diff --git a/src/postgres_scanner.cpp b/src/postgres_scanner.cpp
@@ -251,15 +251,14 @@ static void PostgresInitInternal(ClientContext &context, const PostgresBindData
 	}
 	if (bind_data->table_name.empty()) {
 		D_ASSERT(!bind_data->sql.empty());
-		lstate.sql = StringUtil::Format(
-		    R"(COPY (SELECT %s FROM (%s) AS __unnamed_subquery %s) TO STDOUT (FORMAT "binary");)",
-		    col_names, bind_data->sql, filter);
+		lstate.sql =
+		    StringUtil::Format(R"(COPY (SELECT %s FROM (%s) AS __unnamed_subquery %s) TO STDOUT (FORMAT "binary");)",
+		                       col_names, bind_data->sql, filter);
 
 	} else {
-		lstate.sql = StringUtil::Format(
-		    R"(COPY (SELECT %s FROM %s.%s %s) TO STDOUT (FORMAT "binary");)",
-		    col_names, KeywordHelper::WriteQuoted(bind_data->schema_name, '"'),
-		    KeywordHelper::WriteQuoted(bind_data->table_name, '"'), filter);
+		lstate.sql = StringUtil::Format(R"(COPY (SELECT %s FROM %s.%s %s) TO STDOUT (FORMAT "binary");)", col_names,
+		                                KeywordHelper::WriteQuoted(bind_data->schema_name, '"'),
+		                                KeywordHelper::WriteQuoted(bind_data->table_name, '"'), filter);
 	}
 	lstate.exec = false;
 	lstate.done = false;
@@ -291,11 +290,14 @@ static unique_ptr<GlobalTableFunctionState> PostgresInitGlobalState(ClientContex
 	auto pg_catalog = bind_data.GetCatalog();
 	if (pg_catalog) {
 		auto &transaction = Transaction::Get(context, *pg_catalog).Cast<PostgresTransaction>();
-		auto &con = transaction.GetConnection();
+		auto &con =
+		    bind_data.use_transaction ? transaction.GetConnection() : transaction.GetConnectionWithoutTransaction();
 		result->SetConnection(con.GetConnection());
 	} else {
 		auto con = PostgresConnection::Open(bind_data.dsn);
-		PostgresScanConnect(con, string());
+		if (bind_data.use_transaction) {
+			PostgresScanConnect(con, string());
+		}
 		result->SetConnection(std::move(con));
 	}
 	if (bind_data.requires_materialization) {
diff --git a/src/storage/postgres_transaction.cpp b/src/storage/postgres_transaction.cpp
@@ -39,6 +39,16 @@ static string GetBeginTransactionQuery(AccessMode access_mode) {
 	return result;
 }
 
+PostgresConnection &PostgresTransaction::GetConnectionWithoutTransaction() {
+	if (transaction_state == PostgresTransactionState::TRANSACTION_STARTED) {
+		throw std::runtime_error("Execution without a Transaction is not possible if a Transaction already started");
+	}
+	if (access_mode == AccessMode::READ_ONLY) {
+		throw std::runtime_error("Execution without a Transaction is not possible in Read Only Mode");
+	}
+	return connection.GetConnection();
+}
+
 PostgresConnection &PostgresTransaction::GetConnection() {
 	auto &con = GetConnectionRaw();
 	if (transaction_state == PostgresTransactionState::TRANSACTION_NOT_YET_STARTED) {
@@ -68,6 +78,17 @@ unique_ptr<PostgresResult> PostgresTransaction::Query(const string &query) {
 	return con.Query(query);
 }
 
+unique_ptr<PostgresResult> PostgresTransaction::QueryWithoutTransaction(const string &query) {
+	auto &con = GetConnectionRaw();
+	if (transaction_state == PostgresTransactionState::TRANSACTION_STARTED) {
+		throw std::runtime_error("Execution without a Transaction is not possible if a Transaction already started");
+	}
+	if (access_mode == AccessMode::READ_ONLY) {
+		throw std::runtime_error("Execution without a Transaction is not possible in Read Only Mode");
+	}
+	return con.Query(query);
+}
+
 vector<unique_ptr<PostgresResult>> PostgresTransaction::ExecuteQueries(const string &queries) {
 	auto &con = GetConnectionRaw();
 	if (transaction_state == PostgresTransactionState::TRANSACTION_NOT_YET_STARTED) {
diff --git a/test/sql/storage/postgres_execute_use_transaction.test b/test/sql/storage/postgres_execute_use_transaction.test
@@ -0,0 +1,59 @@
+# name: test/sql/storage/postgres_execute_use_transaction.test
+# description: Test use_transaction flag in postgres_execute
+# group: [storage]
+
+require postgres_scanner
+
+require-env POSTGRES_TEST_DATABASE_AVAILABLE
+
+statement ok
+PRAGMA enable_verification
+
+statement ok
+ATTACH 'dbname=postgresscanner' AS s (TYPE POSTGRES)
+
+statement error
+CALL postgres_execute('s', 'VACUUM')
+----
+Invalid Error: Failed to execute query "BEGIN TRANSACTION ISOLATION LEVEL REPEATABLE READ;
+VACUUM": ERROR:  VACUUM cannot run inside a transaction block
+
+statement error
+CALL postgres_execute('s', 'VACUUM', use_transaction=true)
+----
+Invalid Error: Failed to execute query "BEGIN TRANSACTION ISOLATION LEVEL REPEATABLE READ;
+VACUUM": ERROR:  VACUUM cannot run inside a transaction block
+
+statement error
+CALL postgres_execute('s', 'VACUUM', use_transaction=true)
+----
+Invalid Error: Failed to execute query "BEGIN TRANSACTION ISOLATION LEVEL REPEATABLE READ;
+VACUUM": ERROR:  VACUUM cannot run inside a transaction block
+
+statement ok
+CALL postgres_execute('s', 'VACUUM', use_transaction=false)
+
+statement ok
+BEGIN;
+
+statement ok
+CALL postgres_execute('s', 'SELECT 1')
+
+statement error
+CALL postgres_execute('s', 'VACUUM', use_transaction=false)
+----
+Invalid Error: Execution without a Transaction is not possible if a Transaction already started
+
+statement ok
+ROLLBACK
+
+statement ok
+CALL postgres_execute('s', 'VACUUM', use_transaction=false)
+
+statement ok
+ATTACH 'dbname=postgresscanner' AS s2 (TYPE POSTGRES, READ_ONLY)
+
+statement error
+CALL postgres_execute('s2', 'VACUUM', use_transaction=false)
+----
+Invalid Error: Execution without a Transaction is not possible in Read Only Mode
diff --git a/test/sql/storage/postgres_query_use_transaction.test b/test/sql/storage/postgres_query_use_transaction.test
@@ -0,0 +1,57 @@
+# name: test/sql/storage/postgres_query_use_transaction.test
+# description: Test use_transaction flag in postgres_query
+# group: [storage]
+
+require postgres_scanner
+
+require-env POSTGRES_TEST_DATABASE_AVAILABLE
+
+statement ok
+PRAGMA enable_verification
+
+statement ok
+ATTACH 'dbname=postgresscanner' AS s (TYPE POSTGRES)
+
+query I
+CALL postgres_query('s', 'SELECT 1')
+----
+1
+
+query I
+CALL postgres_query('s', 'SELECT 1', use_transaction=true)
+----
+1
+
+query I
+CALL postgres_query('s', 'SELECT 1', use_transaction=false)
+----
+1
+
+statement ok
+BEGIN;
+
+query I
+CALL postgres_query('s', 'SELECT 1')
+----
+1
+
+statement error
+CALL postgres_query('s', 'SELECT 1', use_transaction=false)
+----
+Invalid Error: Execution without a Transaction is not possible if a Transaction already started
+
+statement ok
+ROLLBACK
+
+query I
+CALL postgres_query('s', 'SELECT 1', use_transaction=false)
+----
+1
+
+statement ok
+ATTACH 'dbname=postgresscanner' AS s2 (TYPE POSTGRES, READ_ONLY)
+
+statement error
+CALL postgres_query('s2', 'SELECT 1', use_transaction=false)
+----
+Invalid Error: Execution without a Transaction is not possible in Read Only Mode
