diff --git a/src/backend/distributed/commands/multi_copy.c b/src/backend/distributed/commands/multi_copy.c index f769d69b3..344c538c0 100644 --- a/src/backend/distributed/commands/multi_copy.c +++ b/src/backend/distributed/commands/multi_copy.c @@ -264,6 +264,7 @@ static void EndPlacementStateCopyCommand(CopyPlacementState *placementState, static void UnclaimCopyConnections(List *connectionStateList); static void ShutdownCopyConnectionState(CopyConnectionState *connectionState, CitusCopyDestReceiver *copyDest); +static SelectStmt * CitusCopySelect(CopyStmt *copyStatement); static void CitusCopyTo(CopyStmt *copyStatement, char *completionTag); static int64 ForwardCopyDataFromConnection(CopyOutState copyOutState, MultiConnection *connection); @@ -2760,22 +2761,9 @@ ProcessCopyStmt(CopyStmt *copyStatement, char *completionTag, const char *queryS { /* * COPY table TO PROGRAM / file is handled by wrapping the table - * in a SELECT * FROM table and going through the result COPY logic. + * in a SELECT and going through the resulting COPY logic. */ - ColumnRef *allColumns = makeNode(ColumnRef); - SelectStmt *selectStmt = makeNode(SelectStmt); - ResTarget *selectTarget = makeNode(ResTarget); - - allColumns->fields = list_make1(makeNode(A_Star)); - allColumns->location = -1; - - selectTarget->name = NULL; - selectTarget->indirection = NIL; - selectTarget->val = (Node *) allColumns; - selectTarget->location = -1; - - selectStmt->targetList = list_make1(selectTarget); - selectStmt->fromClause = list_make1(copyObject(copyStatement->relation)); + SelectStmt *selectStmt = CitusCopySelect(copyStatement); /* replace original statement */ copyStatement = copyObject(copyStatement); @@ -2837,6 +2825,53 @@ ProcessCopyStmt(CopyStmt *copyStatement, char *completionTag, const char *queryS } +/* + * CitusCopySelect generates a SelectStmt such that table may be replaced in + * "COPY table FROM" for an equivalent result. + */ +static SelectStmt * +CitusCopySelect(CopyStmt *copyStatement) +{ + SelectStmt *selectStmt = makeNode(SelectStmt); + selectStmt->fromClause = list_make1(copyObject(copyStatement->relation)); + + Relation distributedRelation = heap_openrv(copyStatement->relation, AccessShareLock); + TupleDesc tupleDescriptor = RelationGetDescr(distributedRelation); + List *targetList = NIL; + + for (int i = 0; i < tupleDescriptor->natts; i++) + { + Form_pg_attribute attr = &tupleDescriptor->attrs[i]; + + if (attr->attisdropped +#if PG_VERSION_NUM >= PG_VERSION_12 + || attr->attgenerated +#endif + ) + { + continue; + } + + ColumnRef *column = makeNode(ColumnRef); + column->fields = list_make1(makeString(pstrdup(attr->attname.data))); + column->location = -1; + + ResTarget *selectTarget = makeNode(ResTarget); + selectTarget->name = NULL; + selectTarget->indirection = NIL; + selectTarget->val = (Node *) column; + selectTarget->location = -1; + + targetList = lappend(targetList, selectTarget); + } + + heap_close(distributedRelation, NoLock); + + selectStmt->targetList = targetList; + return selectStmt; +} + + /* * CitusCopyTo runs a COPY .. TO STDOUT command on each shard to do a full * table dump. @@ -3061,6 +3096,10 @@ CopyGetAttnums(TupleDesc tupDesc, Relation rel, List *attnamelist) { if (TupleDescAttr(tupDesc, i)->attisdropped) continue; +#if PG_VERSION_NUM >= PG_VERSION_12 + if (TupleDescAttr(tupDesc, i)->attgenerated) + continue; +#endif attnums = lappend_int(attnums, i + 1); } } @@ -3085,6 +3124,14 @@ CopyGetAttnums(TupleDesc tupDesc, Relation rel, List *attnamelist) continue; if (namestrcmp(&(att->attname), name) == 0) { +#if PG_VERSION_NUM >= PG_VERSION_12 + if (att->attgenerated) + ereport(ERROR, + (errcode(ERRCODE_INVALID_COLUMN_REFERENCE), + errmsg("column \"%s\" is a generated column", + name), + errdetail("Generated columns cannot be used in COPY."))); +#endif attnum = att->attnum; break; } diff --git a/src/test/regress/expected/pg12.out b/src/test/regress/expected/pg12.out index 704c6d4a2..eeec7a4e5 100644 --- a/src/test/regress/expected/pg12.out +++ b/src/test/regress/expected/pg12.out @@ -47,6 +47,7 @@ HINT: To remove the local data, run: SELECT truncate_local_data_after_distribut select create_distributed_table('gen2', 'val2'); ERROR: cannot distribute relation: gen2 DETAIL: Distribution column must not use GENERATED ALWAYS AS (...) STORED. +copy gen1 to :'temp_dir''pg12_copy_test_generated'; insert into gen1 (id, val1) values (2,4),(4,6),(6,2),(8,2); insert into gen2 (id, val1) values (2,4),(4,6),(6,2),(8,2); select * from gen1 order by 1,2,3; @@ -75,6 +76,17 @@ select * from gen2 order by 1,2,3; 8 | 2 | 4 (8 rows) +truncate gen1; +copy gen1 from :'temp_dir''pg12_copy_test_generated'; +select * from gen1 order by 1,2,3; + id | val2 | val1 +--------------------------------------------------------------------- + 1 | 6 | 4 + 3 | 8 | 6 + 5 | 4 | 2 + 7 | 4 | 2 +(4 rows) + -- Test new VACUUM/ANALYZE options analyze (skip_locked) gen1; vacuum (skip_locked) gen1; diff --git a/src/test/regress/sql/pg12.sql b/src/test/regress/sql/pg12.sql index ba7d475f8..4782bb613 100644 --- a/src/test/regress/sql/pg12.sql +++ b/src/test/regress/sql/pg12.sql @@ -43,12 +43,18 @@ insert into gen2 (id, val1) values (1,4),(3,6),(5,2),(7,2); select create_distributed_table('gen1', 'id'); select create_distributed_table('gen2', 'val2'); +copy gen1 to :'temp_dir''pg12_copy_test_generated'; + insert into gen1 (id, val1) values (2,4),(4,6),(6,2),(8,2); insert into gen2 (id, val1) values (2,4),(4,6),(6,2),(8,2); select * from gen1 order by 1,2,3; select * from gen2 order by 1,2,3; +truncate gen1; +copy gen1 from :'temp_dir''pg12_copy_test_generated'; +select * from gen1 order by 1,2,3; + -- Test new VACUUM/ANALYZE options analyze (skip_locked) gen1; vacuum (skip_locked) gen1;