Fix memory leaks in custom executor (#256)

include/pgduckdb/pgduckdb_utils.hpp

@@ -5,6 +5,7 @@ extern "C" {
 }
 
 #include "duckdb/common/exception.hpp"
+#include "duckdb/common/error_data.hpp"
 
 #include <vector>
 #include <string>
@@ -74,4 +75,30 @@ PostgresFunctionGuard(FuncType postgres_function, FuncArgs... args) {
 	}
 }
 
+template <typename FuncRetT, typename FuncType, typename... FuncArgs>
+FuncRetT
+DuckDBFunctionGuard(FuncType duckdb_function, const char* function_name, FuncArgs... args) {
+	const char *error_message = nullptr;
+	try {
+		return duckdb_function(args...);
+	} catch (duckdb::Exception &ex) {
+		duckdb::ErrorData edata(ex.what());
+		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());
+		} else {
+			error_message = pstrdup(ex.what());
+		}
+	}
+
+	if (error_message) {
+		elog(ERROR, "(PGDuckDB/%s) %s", function_name, error_message);
+	}
+
+	std::abort(); // Cannot reach.
+}
+
 } // namespace pgduckdb

src/pgduckdb_node.cpp

@@ -13,6 +13,7 @@ extern "C" {
 #include "pgduckdb/pgduckdb_types.hpp"
 #include "pgduckdb/pgduckdb_duckdb.hpp"
 #include "pgduckdb/pgduckdb_planner.hpp"
+#include "pgduckdb/pgduckdb_utils.hpp"
 
 /* global variables */
 CustomScanMethods duckdb_scan_scan_methods;
@@ -36,9 +37,21 @@ typedef struct DuckdbScanState {
 
 static void
 CleanupDuckdbScanState(DuckdbScanState *state) {
+	MemoryContextReset(state->css.ss.ps.ps_ExprContext->ecxt_per_tuple_memory);
+	ExecClearTuple(state->css.ss.ss_ScanTupleSlot);
+
 	state->query_results.reset();
-	delete state->prepared_statement;
-	delete state->duckdb_connection;
+	state->current_data_chunk.reset();
+
+	if (state->prepared_statement) {
+		delete state->prepared_statement;
+		state->prepared_statement = nullptr;
+	}
+
+	if (state->duckdb_connection) {
+		delete state->duckdb_connection;
+		state->duckdb_connection = nullptr;
+	}
 }
 
 /* static callbacks */
@@ -105,13 +118,13 @@ ExecuteQuery(DuckdbScanState *state) {
 		} else {
 			std::ostringstream oss;
 			oss << "parameter '" << i << "' has an invalid type (" << pg_param->ptype << ") during query execution";
-			elog(ERROR, oss.str().c_str());
+			throw duckdb::Exception(duckdb::ExceptionType::EXECUTOR, oss.str().c_str());
 		}
 	}
 
 	auto pending = prepared.PendingQuery(duckdb_params, true);
 	if (pending->HasError()) {
-		elog(ERROR, "DuckDB execute returned an error: %s", pending->GetError().c_str());
+		return pending->ThrowError();
 	}
 
 	duckdb::PendingExecutionResult execution_result;
@@ -124,16 +137,14 @@ ExecuteQuery(DuckdbScanState *state) {
 			// Wait for all tasks to terminate
 			executor.CancelTasks();
 			// Delete the scan state
-			CleanupDuckdbScanState(state);
 			// Process the interrupt on the Postgres side
 			ProcessInterrupts();
-			elog(ERROR, "Query cancelled");
+			throw duckdb::Exception(duckdb::ExceptionType::EXECUTOR, "Query cancelled");
 		}
 	} while (!duckdb::PendingQueryResult::IsResultReady(execution_result));
 
 	if (execution_result == duckdb::PendingExecutionResult::EXECUTION_ERROR) {
-		CleanupDuckdbScanState(state);
-		elog(ERROR, "(PGDuckDB/ExecuteQuery) %s", pending->GetError().c_str());
+		return pending->ThrowError();
 	}
 
 	query_results = pending->Execute();
@@ -149,7 +160,7 @@ Duckdb_ExecCustomScan(CustomScanState *node) {
 
 	bool already_executed = duckdb_scan_state->is_executed;
 	if (!already_executed) {
-		ExecuteQuery(duckdb_scan_state);
+		pgduckdb::DuckDBFunctionGuard<void>(ExecuteQuery, "ExecuteQuery", duckdb_scan_state);
 	}
 
 	if (duckdb_scan_state->fetch_next) {
@@ -209,7 +220,8 @@ Duckdb_ReScanCustomScan(CustomScanState *node) {
 void
 Duckdb_ExplainCustomScan(CustomScanState *node, List *ancestors, ExplainState *es) {
 	DuckdbScanState *duckdb_scan_state = (DuckdbScanState *)node;
-	ExecuteQuery(duckdb_scan_state);
+	pgduckdb::DuckDBFunctionGuard<void>(ExecuteQuery, "ExecuteQuery", duckdb_scan_state);
+
 	auto chunk = duckdb_scan_state->query_results->Fetch();
 	if (!chunk || chunk->size() == 0) {
 		return;

src/pgduckdb_planner.cpp

@@ -108,7 +108,8 @@ CreatePlan(Query *query) {
 PlannedStmt *
 DuckdbPlanNode(Query *parse, int cursor_options) {
 	/* We need to check can we DuckDB create plan */
-	Plan *duckdb_plan = (Plan *)castNode(CustomScan, CreatePlan(parse));
+	Plan *plan = pgduckdb::DuckDBFunctionGuard<Plan*>(CreatePlan, "CreatePlan", parse);
+	Plan *duckdb_plan = (Plan *)castNode(CustomScan, plan);
 
 	if (!duckdb_plan) {
 		return nullptr;