From c94c78e7d21aef3ef57951dff36a79975ba229c6 Mon Sep 17 00:00:00 2001
From: Cheng Chen <cheng@mooncakelabs.com>
Date: Tue, 17 Dec 2024 22:03:02 +0000
Subject: [PATCH] Allow writes to columnstore and Postgres tables in the same
 transaction

Remove the lockdowns introduced in https://github.com/duckdb/pg_duckdb/pull/433 since the same restriction doesn't apply to us. In our case, Postgres remains the source of truth for columnstore tables, and DuckDB serves purely as the execution engine for reading and writing to those tables.
---
 src/pgduckdb/pgduckdb_hooks.cpp | 20 ++++++++++----------
 src/pgduckdb/pgduckdb_xact.cpp  | 22 +++++++++++-----------
 test/expected/transaction.out   | 19 +++++++++++++++++++
 test/sql/transaction.sql        | 12 ++++++++++++
 4 files changed, 52 insertions(+), 21 deletions(-)
 create mode 100644 test/expected/transaction.out
 create mode 100644 test/sql/transaction.sql

diff --git a/src/pgduckdb/pgduckdb_hooks.cpp b/src/pgduckdb/pgduckdb_hooks.cpp
index d4a32ab..15afe52 100644
--- a/src/pgduckdb/pgduckdb_hooks.cpp
+++ b/src/pgduckdb/pgduckdb_hooks.cpp
@@ -163,12 +163,12 @@ IsAllowedStatement(Query *query, bool throw_error = false) {
 			elog(elevel, "DuckDB does not support modififying Postgres tables");
 			return false;
 		}
-		if (pgduckdb::pg::IsInTransactionBlock(true)) {
-			if (pgduckdb::pg::DidWalWrites()) {
-				elog(elevel, "Writing to DuckDB and Postgres tables in the same transaction block is not supported");
-				return false;
-			}
-		}
+		// if (pgduckdb::pg::IsInTransactionBlock(true)) {
+		// 	if (pgduckdb::pg::DidWalWrites()) {
+		// 		elog(elevel, "Writing to DuckDB and Postgres tables in the same transaction block is not supported");
+		// 		return false;
+		// 	}
+		// }
 	}
 
 	/*
@@ -242,10 +242,10 @@ DuckdbPlannerHook_Cpp(Query *parse, const char *query_string, int cursor_options
 			}
 			/* If we can't create a plan, we'll fall back to Postgres */
 		}
-		if (parse->commandType != CMD_SELECT && pgduckdb::ddb::DidWrites() &&
-		    pgduckdb::pg::IsInTransactionBlock(true)) {
-			elog(ERROR, "Writing to DuckDB and Postgres tables in the same transaction block is not supported");
-		}
+		// if (parse->commandType != CMD_SELECT && pgduckdb::ddb::DidWrites() &&
+		//     pgduckdb::pg::IsInTransactionBlock(true)) {
+		// 	elog(ERROR, "Writing to DuckDB and Postgres tables in the same transaction block is not supported");
+		// }
 	}
 
 	/*
diff --git a/src/pgduckdb/pgduckdb_xact.cpp b/src/pgduckdb/pgduckdb_xact.cpp
index ddabec1..ab63cc4 100644
--- a/src/pgduckdb/pgduckdb_xact.cpp
+++ b/src/pgduckdb/pgduckdb_xact.cpp
@@ -41,7 +41,7 @@ PreventInTransactionBlock(const char *statement_type) {
  * IMPORTANT: This function should only be called at trasaction commit. At
  * other points in the transaction lifecycle its return value is not reliable.
  */
-static bool
+[[maybe_unused]] static bool
 DidWritesAtTransactionEnd() {
 	return pg::DidWalWrites() || pg::GetCurrentCommandId() > duckdb_command_id + 1;
 }
@@ -87,10 +87,10 @@ ClaimCurrentCommandId() {
 		return;
 	}
 
-	if (new_command_id != duckdb_command_id + 1) {
-		throw duckdb::NotImplementedException(
-		    "Writing to DuckDB and Postgres tables in the same transaction block is not supported");
-	}
+	// if (new_command_id != duckdb_command_id + 1) {
+	// 	throw duckdb::NotImplementedException(
+	// 	    "Writing to DuckDB and Postgres tables in the same transaction block is not supported");
+	// }
 
 	duckdb_command_id = new_command_id;
 }
@@ -157,12 +157,12 @@ DuckdbXactCallback_Cpp(XactEvent event) {
 	switch (event) {
 	case XACT_EVENT_PRE_COMMIT:
 	case XACT_EVENT_PARALLEL_PRE_COMMIT:
-		if (pg::IsInTransactionBlock(top_level_statement)) {
-			if (pg::DidWritesAtTransactionEnd() && ddb::DidWrites(context)) {
-				throw duckdb::NotImplementedException(
-				    "Writing to DuckDB and Postgres tables in the same transaction block is not supported");
-			}
-		}
+		// if (pg::IsInTransactionBlock(top_level_statement)) {
+		// 	if (pg::DidWritesAtTransactionEnd() && ddb::DidWrites(context)) {
+		// 		throw duckdb::NotImplementedException(
+		// 		    "Writing to DuckDB and Postgres tables in the same transaction block is not supported");
+		// 	}
+		// }
 		top_level_statement = true;
 		duckdb_command_id = -1;
 		// Commit the DuckDB transaction too
diff --git a/test/expected/transaction.out b/test/expected/transaction.out
new file mode 100644
index 0000000..dbdc3b6
--- /dev/null
+++ b/test/expected/transaction.out
@@ -0,0 +1,19 @@
+CREATE TABLE s (a int);
+CREATE TABLE t (b int) USING columnstore;
+BEGIN;
+INSERT INTO s VALUES (1);
+INSERT INTO t VALUES (2);
+COMMIT;
+SELECT * FROM s;
+ a 
+---
+ 1
+(1 row)
+
+SELECT * FROM t;
+ b 
+---
+ 2
+(1 row)
+
+DROP TABLE s, t;
diff --git a/test/sql/transaction.sql b/test/sql/transaction.sql
new file mode 100644
index 0000000..a303fa0
--- /dev/null
+++ b/test/sql/transaction.sql
@@ -0,0 +1,12 @@
+CREATE TABLE s (a int);
+CREATE TABLE t (b int) USING columnstore;
+
+BEGIN;
+INSERT INTO s VALUES (1);
+INSERT INTO t VALUES (2);
+COMMIT;
+
+SELECT * FROM s;
+SELECT * FROM t;
+
+DROP TABLE s, t;