Improve CPP & PG Exception handlers

Author: Y--

include/pgduckdb/pg/error_data.hpp

@@ -3,8 +3,8 @@
 #include "duckdb/common/exception.hpp"
 #include "duckdb/common/error_data.hpp"
 #include "pgduckdb/pgduckdb_duckdb.hpp"
-#include "pgduckdb/pg/error_data.hpp"
 #include "pgduckdb/logger.hpp"
+#include "pgduckdb/utility/exception.hpp"
 
 #include <setjmp.h>
 
@@ -14,6 +14,7 @@ extern "C" {
 // Note: these forward-declarations could live in a header under the `pg/` folder.
 // But since they are (hopefully) only used in this file, we keep them here.
 struct ErrorContextCallback;
+struct ErrorData;
 struct MemoryContextData;
 
 typedef struct MemoryContextData *MemoryContext;
@@ -26,6 +27,7 @@ extern ErrorData *CopyErrorData();
 extern void FlushErrorState();
 extern pg_stack_base_t set_stack_base();
 extern void restore_stack_base(pg_stack_base_t base);
+extern int errdetail(const char *fmt, ...);
 }
 
 namespace pgduckdb {
@@ -83,7 +85,7 @@ struct PostgresScopedStackReset {
  */
 template <typename Func, Func func, typename... FuncArgs>
 typename std::invoke_result<Func, FuncArgs...>::type
-__PostgresFunctionGuard__(const char *func_name, FuncArgs... args) {
+__PostgresFunctionGuard__(const char *func_name, const char *file_name, int line, FuncArgs... args) {
 	std::lock_guard<std::recursive_mutex> lock(pgduckdb::GlobalProcessLock::GetLock());
 	MemoryContext ctx = CurrentMemoryContext;
 
@@ -110,20 +112,22 @@ __PostgresFunctionGuard__(const char *func_name, FuncArgs... args) {
 			FlushErrorState();
 		} else {
 			// This is a pretty bad situation - we failed to extract the error message.
-			throw duckdb::Exception(duckdb::ExceptionType::EXECUTOR, "Failed to extract Postgres error message");
+			pd_log(WARNING, "(PGGuard/%s) Exception in %s:%d: could not extract ErrorData.", func_name, file_name,
+			       line);
+			throw Exception();
 		}
 	} // PG_END_TRY
 
-	auto message = duckdb::StringUtil::Format("(PGDuckDB/%s) %s", func_name, pg::GetErrorDataMessage(edata));
-	throw duckdb::Exception(duckdb::ExceptionType::EXECUTOR, message);
+	throw pgduckdb::Exception(edata);
 }
 
 #define PostgresFunctionGuard(FUNC, ...)                                                                               \
-	pgduckdb::__PostgresFunctionGuard__<decltype(&FUNC), &FUNC>(#FUNC, ##__VA_ARGS__)
+	pgduckdb::__PostgresFunctionGuard__<decltype(&FUNC), &FUNC>(#FUNC, __FILE__, __LINE__, ##__VA_ARGS__)
 
 template <typename T, typename ReturnType, typename... FuncArgs>
 ReturnType
-__PostgresMemberGuard__(ReturnType (T::*func)(FuncArgs... args), T *instance, const char *func_name, FuncArgs... args) {
+__PostgresMemberGuard__(ReturnType (T::*func)(FuncArgs... args), T *instance, const char *func_name,
+                        const char *file_name, int line, FuncArgs... args) {
 	MemoryContext ctx = CurrentMemoryContext;
 
 	{ // Scope for PG_END_TRY
@@ -150,15 +154,17 @@ __PostgresMemberGuard__(ReturnType (T::*func)(FuncArgs... args), T *instance, co
 			FlushErrorState();
 		} else {
 			// This is a pretty bad situation - we failed to extract the error message.
-			throw duckdb::Exception(duckdb::ExceptionType::EXECUTOR, "Failed to extract Postgres error message");
+			pd_log(WARNING, "(PGGuard/%s) Exception in %s:%d: could not extract ErrorData.", func_name, file_name,
+			       line);
+			throw Exception();
 		}
 	} // PG_END_TRY
 
-	auto message = duckdb::StringUtil::Format("(PGDuckDB/%s) %s", func_name, pg::GetErrorDataMessage(edata));
-	throw duckdb::Exception(duckdb::ExceptionType::EXECUTOR, message);
+	throw pgduckdb::Exception(edata);
 }
 
-#define PostgresMemberGuard(FUNC, ...) pgduckdb::__PostgresMemberGuard__(&FUNC, this, __func__, ##__VA_ARGS__)
+#define PostgresMemberGuard(FUNC, ...)                                                                                 \
+	pgduckdb::__PostgresMemberGuard__(&FUNC, this, __func__, __FILE__, __LINE__, ##__VA_ARGS__)
 
 duckdb::unique_ptr<duckdb::QueryResult> DuckDBQueryOrThrow(duckdb::ClientContext &context, const std::string &query);
 

include/pgduckdb/pgduckdb_utils.hpp

@@ -1,6 +1,7 @@
 #pragma once
 
 #include <duckdb/common/error_data.hpp>
+#include "pgduckdb/utility/exception.hpp"
 
 extern "C" {
 #include "postgres.h"
@@ -13,16 +14,39 @@ typename std::invoke_result<Func, FuncArgs &...>::type
 __CPPFunctionGuard__(const char *func_name, const char *file_name, int line, FuncArgs &...args) {
 	const char *error_message = nullptr;
 	auto pg_es_start = PG_exception_stack;
+
+	// We may restore the PG error data from two ways:
+	// 1. The error data is stored in the exception object and the error type was preserved
+	// 2. The error pointer was stored in the extra info of the DuckDB exception object
+	ErrorData *maybe_error_data = nullptr;
 	try {
 		return func(args...);
+	} catch (pgduckdb::Exception &ex) {
+		if (ex.error_data == nullptr) {
+			error_message = "Failed to extract Postgres error message";
+		} else {
+			maybe_error_data = ex.error_data;
+		}
 	} catch (duckdb::Exception &ex) {
 		duckdb::ErrorData edata(ex.what());
-		error_message = pstrdup(edata.Message().c_str());
+		auto it = edata.ExtraInfo().find("error_data_ptr");
+		if (it != edata.ExtraInfo().end()) {
+			// Restore pointer stored in DuckDB exception
+			maybe_error_data = reinterpret_cast<ErrorData *>(std::stoull(it->second));
+		} else {
+			error_message = pstrdup(edata.Message().c_str());
+		}
 	} catch (std::exception &ex) {
 		const auto msg = ex.what();
 		if (msg[0] == '{') {
 			duckdb::ErrorData edata(ex.what());
-			error_message = pstrdup(edata.Message().c_str());
+			auto it = edata.ExtraInfo().find("error_data_ptr");
+			if (it != edata.ExtraInfo().end()) {
+				// Restore pointer stored in DuckDB exception
+				maybe_error_data = reinterpret_cast<ErrorData *>(std::stoull(it->second));
+			} else {
+				error_message = pstrdup(edata.Message().c_str());
+			}
 		} else {
 			error_message = pstrdup(ex.what());
 		}
@@ -73,11 +97,14 @@ __CPPFunctionGuard__(const char *func_name, const char *file_name, int line, Fun
 		PG_exception_stack = pg_es_start;
 	}
 
-	// Simplified version of `elog(ERROR, ...)`, with arguments inlined
-	if (errstart_cold(ERROR, TEXTDOMAIN)) {
+	if (maybe_error_data) {
+		ReThrowError(maybe_error_data);
+	} else if (errstart_cold(ERROR, TEXTDOMAIN)) {
+		// Simplified version of `elog(ERROR, ...)`, with arguments inlined
 		errmsg_internal("(PGDuckDB/%s) %s", func_name, error_message);
 		errfinish(file_name, line, func_name);
 	}
+
 	pg_unreachable();
 }
 

include/pgduckdb/utility/cpp_wrapper.hpp

@@ -0,0 +1,36 @@
+#pragma once
+
+#include "duckdb/common/exception.hpp"
+
+extern "C" {
+struct ErrorData;
+}
+
+namespace pgduckdb {
+
+/*
+Custom exception class for PGDuckDB that acts as a shell for ErrorData.
+
+Note: we have to capture the ErrorData pointer in the exception which is
+a subclass of duckdb::Exception. This is because in some cases, the exception
+is caught by DuckDB and re-thrown. So the original exception is lost.
+
+While it is not ideal to stringify the pointer, it should be valid as long as
+it is properly caught by the CPP exception handler in an upper stack.
+*/
+
+struct Exception : public duckdb::Exception {
+	Exception(ErrorData *edata = nullptr)
+	    : duckdb::Exception::Exception(duckdb::ExceptionType::INVALID, "PGDuckDB Exception",
+	                                   {{"error_data_ptr", std::to_string(reinterpret_cast<uintptr_t>(edata))}}),
+	      error_data(edata) {
+	}
+
+	ErrorData *error_data = nullptr;
+
+private:
+	Exception(const Exception &) = delete;
+	Exception &operator=(const Exception &) = delete;
+};
+
+} // namespace pgduckdb

include/pgduckdb/utility/exception.hpp

@@ -50,7 +50,7 @@ EXPLAIN SELECT count(*) FROM ta;
 -- Not allowed in a transaction
 BEGIN;
 CALL duckdb.recycle_ddb();
-ERROR:  (PGDuckDB/pgduckdb_recycle_ddb_cpp) Executor Error: (PGDuckDB/::PreventInTransactionBlock) duckdb.recycle_ddb() cannot run inside a transaction block
+ERROR:  duckdb.recycle_ddb() cannot run inside a transaction block
 END;
 -- Nor in a function
 CREATE OR REPLACE FUNCTION f() RETURNS void
@@ -64,7 +64,7 @@ END;
 $$;
 SET duckdb.force_execution = false;
 SELECT * FROM f();
-ERROR:  (PGDuckDB/pgduckdb_recycle_ddb_cpp) Executor Error: (PGDuckDB/::PreventInTransactionBlock) duckdb.recycle_ddb() cannot be executed from a function
+ERROR:  duckdb.recycle_ddb() cannot be executed from a function
 CONTEXT:  SQL statement "CALL duckdb.recycle_ddb()"
 PL/pgSQL function f() line 3 at CALL
 DROP TABLE ta;

src/pg/error_data.cpp

@@ -87,7 +87,7 @@ SELECT 1;
 (1 row)
 
 CALL duckdb.enable_motherduck();
-ERROR:  (PGDuckDB/pgduckdb_enable_motherduck_cpp) Executor Error: (PGDuckDB/::PreventInTransactionBlock) duckdb.enable_motherduck() cannot run inside a transaction block
+ERROR:  duckdb.enable_motherduck() cannot run inside a transaction block
 ROLLBACK;
 -- Use helper to enable MD: will fail since there's no token in environment
 CALL duckdb.enable_motherduck();
@@ -103,7 +103,7 @@ SELECT 1;
 (1 row)
 
 CALL duckdb.enable_motherduck();
-ERROR:  (PGDuckDB/pgduckdb_enable_motherduck_cpp) Executor Error: (PGDuckDB/::PreventInTransactionBlock) duckdb.enable_motherduck() cannot run inside a transaction block
+ERROR:  duckdb.enable_motherduck() cannot run inside a transaction block
 ROLLBACK;
 -- Now MD is enabled again
 SELECT * FROM duckdb.is_motherduck_enabled();

test/regression/expected/duckdb_recycle.out

@@ -2,4 +2,4 @@ SET duckdb.force_execution = TRUE;
 COPY (SELECT 1 INTO frak UNION SELECT 2) TO 'blob';
 ERROR:  COPY (SELECT INTO) is not supported
 COPY (SELECT 1 INTO frak UNION SELECT 2) TO 'blob.parquet';
-ERROR:  (PGDuckDB/DuckdbUtilityHook_Cpp) Executor Error: (PGDuckDB/MakeDuckdbCopyQuery) DuckDB COPY only supports SELECT statements
+ERROR:  DuckDB COPY only supports SELECT statements

test/regression/expected/foreign_data_wrapper.out

@@ -25,7 +25,7 @@ ERROR:  permission denied for table extensions
 -- permissions checked even if it happens straight from DuckDB.
 SET duckdb.force_execution = false;
 SELECT * FROM duckdb.raw_query($$ SELECT * FROM pgduckdb.public.t $$);
-ERROR:  (PGDuckDB/pgduckdb_raw_query_cpp) Executor Error: (PGDuckDB/Init) permission denied for table t
+ERROR:  permission denied for table t
 SET duckdb.force_execution = true;
 -- read_csv from the local filesystem should be disallowed
 SELECT count(r['sepal.length']) FROM read_csv('../../data/iris.csv') r;

test/regression/expected/issue_789.out

@@ -20,6 +20,7 @@ BEGIN
 END;
 $$;
 SET duckdb.force_execution = false;
+
 SELECT * FROM f();
 
 DROP TABLE ta;