diff --git a/src/backend/distributed/executor/multi_utility.c b/src/backend/distributed/executor/multi_utility.c index 158ffccf8..3863fbe59 100644 --- a/src/backend/distributed/executor/multi_utility.c +++ b/src/backend/distributed/executor/multi_utility.c @@ -62,6 +62,7 @@ #include "foreign/foreign.h" #include "lib/stringinfo.h" #include "nodes/bitmapset.h" +#include "nodes/memnodes.h" #include "nodes/nodes.h" #include "nodes/params.h" #include "nodes/parsenodes.h" @@ -311,12 +312,17 @@ multi_ProcessUtility(PlannedStmt *pstmt, if (IsA(parsetree, CopyStmt)) { - /* copy parse tree since we might scribble on it to fix the schema name */ - parsetree = copyObject(parsetree); + MemoryContext planContext = GetMemoryChunkContext(parsetree); + MemoryContext previousContext; + parsetree = copyObject(parsetree); parsetree = ProcessCopyStmt((CopyStmt *) parsetree, completionTag, &commandMustRunAsOwner); + previousContext = MemoryContextSwitchTo(planContext); + parsetree = copyObject(parsetree); + MemoryContextSwitchTo(previousContext); + if (parsetree == NULL) { return; diff --git a/src/test/regress/expected/multi_prepare_plsql.out b/src/test/regress/expected/multi_prepare_plsql.out index 037e5436f..1d86f035a 100644 --- a/src/test/regress/expected/multi_prepare_plsql.out +++ b/src/test/regress/expected/multi_prepare_plsql.out @@ -1233,9 +1233,31 @@ SELECT copy_in_plpgsql(); (1 row) +-- test prepared COPY on a non-distributed table +CREATE TABLE local_ddl (x int); +CREATE OR REPLACE FUNCTION local_copy_in_plpgsql() +RETURNS VOID AS +$BODY$ +BEGIN + COPY local_ddl (x) FROM PROGRAM 'echo 1' WITH CSV; +END; +$BODY$ LANGUAGE plpgsql; +SELECT local_copy_in_plpgsql(); + local_copy_in_plpgsql +----------------------- + +(1 row) + +SELECT local_copy_in_plpgsql(); + local_copy_in_plpgsql +----------------------- + +(1 row) + DROP FUNCTION ddl_in_plpgsql(); DROP FUNCTION copy_in_plpgsql(); DROP TABLE prepare_ddl; +DROP TABLE local_ddl; -- clean-up functions DROP FUNCTION plpgsql_test_1(); DROP FUNCTION plpgsql_test_2(); diff --git a/src/test/regress/sql/multi_prepare_plsql.sql b/src/test/regress/sql/multi_prepare_plsql.sql index f4cd15489..e047c157c 100644 --- a/src/test/regress/sql/multi_prepare_plsql.sql +++ b/src/test/regress/sql/multi_prepare_plsql.sql @@ -563,9 +563,24 @@ $BODY$ LANGUAGE plpgsql; SELECT copy_in_plpgsql(); SELECT copy_in_plpgsql(); +-- test prepared COPY on a non-distributed table +CREATE TABLE local_ddl (x int); + +CREATE OR REPLACE FUNCTION local_copy_in_plpgsql() +RETURNS VOID AS +$BODY$ +BEGIN + COPY local_ddl (x) FROM PROGRAM 'echo 1' WITH CSV; +END; +$BODY$ LANGUAGE plpgsql; + +SELECT local_copy_in_plpgsql(); +SELECT local_copy_in_plpgsql(); + DROP FUNCTION ddl_in_plpgsql(); DROP FUNCTION copy_in_plpgsql(); DROP TABLE prepare_ddl; +DROP TABLE local_ddl; -- clean-up functions DROP FUNCTION plpgsql_test_1();