Don't parse with lexer query string

Manually parse OPTIONS part of COPY query. Check settings and allow only combinations that should work with DuckDB execution.
duckdb · Oct 11, 2024 · 90b297c · 90b297c
1 parent a08020f
commit 90b297c
Show file tree

Hide file tree

Showing 3 changed files with 331 additions and 43 deletions.
diff --git a/include/pgduckdb/utility/copy.hpp b/include/pgduckdb/utility/copy.hpp
@@ -5,4 +5,5 @@ extern "C" {
 #include "nodes/plannodes.h"
 }
 
-bool DuckdbCopy(PlannedStmt *pstmt, const char *query_string, struct QueryEnvironment *query_env, uint64 *processed);
+bool DuckdbCopy(PlannedStmt *pstmt, const char *query_string, struct QueryEnvironment *query_env, uint64 *processed,
+                bool *is_copy_to_cloud);
diff --git a/src/pgduckdb_hooks.cpp b/src/pgduckdb_hooks.cpp
@@ -194,14 +194,19 @@ static void
 DuckdbUtilityHook(PlannedStmt *pstmt, const char *query_string, bool read_only_tree, ProcessUtilityContext context,
                   ParamListInfo params, struct QueryEnvironment *query_env, DestReceiver *dest, QueryCompletion *qc) {
 	Node *parsetree = pstmt->utilityStmt;
+	bool is_copy_to_cloud = false;
 	if (duckdb_execution && pgduckdb::IsExtensionRegistered() && IsA(parsetree, CopyStmt)) {
 		uint64 processed;
-		if (DuckdbCopy(pstmt, query_string, query_env, &processed)) {
+		if (DuckdbCopy(pstmt, query_string, query_env, &processed, &is_copy_to_cloud)) {
 			if (qc) {
 				SetQueryCompletion(qc, CMDTAG_COPY, processed);
 			}
 			return;
 		}
+		/* We have COPY .. TO remote so PG can't handle it, raise ERROR */
+		if (is_copy_to_cloud) {
+			elog(ERROR, "DuckDB COPY execution failed.");
+		}
 	}
 
 	if (pgduckdb::IsExtensionRegistered()) {