diff --git a/contrib/postgres_fdw/connection.c b/contrib/postgres_fdw/connection.c index b5d4cf3dccc..f3492cd9c85 100644 --- a/contrib/postgres_fdw/connection.c +++ b/contrib/postgres_fdw/connection.c @@ -1214,6 +1214,21 @@ pgfdw_xact_callback(XactEvent event, void *arg) */ pgfdw_reject_incomplete_xact_state_change(entry); + /* Close COPY if in progress before committing */ + if (entry->state.copy_in_progress) + { + PGresult *copyres; + + if (PQputCopyEnd(entry->conn, NULL) < 0 || + PQflush(entry->conn)) + pgfdw_report_error(NULL, entry->conn, NULL); + copyres = pgfdw_get_result(entry->conn); + if (PQresultStatus(copyres) != PGRES_COMMAND_OK) + pgfdw_report_error(copyres, entry->conn, NULL); + PQclear(copyres); + entry->state.copy_in_progress = false; + } + /* Commit all remote transactions during pre-commit */ entry->changing_xact_state = true; if (entry->parallel_commit) @@ -1899,6 +1914,27 @@ pgfdw_abort_cleanup(ConnCacheEntry *entry, bool toplevel) /* Assume we might have lost track of prepared statements */ entry->have_error = true; + /* + * If a COPY FROM STDIN started by ExecForeignBatchCopy is still open on + * this connection, terminate it before we can send any other commands. We + * send a COPY end with an error message so the remote aborts it. + */ + if (entry->state.copy_in_progress) + { + if (PQputCopyEnd(entry->conn, "COPY aborted due to local error") == 1) + { + PGresult *res; + + /* + * Consume the error result from the aborted COPY. We don't care + * about the specific error since we're aborting anyway. + */ + res = PQgetResult(entry->conn); + PQclear(res); + } + entry->state.copy_in_progress = false; + } + /* * If a command has been submitted to the remote server by using an * asynchronous execution function, the command might not have yet @@ -1972,6 +2008,26 @@ pgfdw_abort_cleanup_begin(ConnCacheEntry *entry, bool toplevel, /* Assume we might have lost track of prepared statements */ entry->have_error = true; + /* + * If a COPY FROM STDIN started by ExecForeignBatchCopy is still open on + * this connection, terminate it before we can send any other commands. + */ + if (entry->state.copy_in_progress) + { + if (PQputCopyEnd(entry->conn, "COPY aborted due to local error") == 1) + { + PGresult *res; + + /* + * Consume the error result from the aborted COPY. We don't care + * about the specific error since we're aborting anyway. + */ + res = PQgetResult(entry->conn); + PQclear(res); + } + entry->state.copy_in_progress = false; + } + /* * If a command has been submitted to the remote server by using an * asynchronous execution function, the command might not have yet diff --git a/contrib/postgres_fdw/deparse.c b/contrib/postgres_fdw/deparse.c index 673b678826c..c73b3be2cd3 100644 --- a/contrib/postgres_fdw/deparse.c +++ b/contrib/postgres_fdw/deparse.c @@ -2282,6 +2282,61 @@ deparseInsertSql(StringInfo buf, RangeTblEntry *rte, withCheckOptionList, returningList, retrieved_attrs); } +/* + * Build a COPY FROM STDIN statement using the TEXT format + */ +void +deparseCopySql(StringInfo buf, Relation rel, List *target_attrs) +{ + Oid relid = RelationGetRelid(rel); + TupleDesc tupdesc = RelationGetDescr(rel); + bool first = true; + + appendStringInfo(buf, "COPY "); + deparseRelation(buf, rel); + + /* + * Emit the column list, skipping generated columns (which the remote + * server fills in itself). The opening parenthesis is emitted lazily so + * that a relation whose transmittable columns are all generated produces + * no (empty) column list at all. + */ + foreach_int(attnum, target_attrs) + { + Form_pg_attribute attr = TupleDescAttr(tupdesc, attnum - 1); + char *colname; + List *options; + ListCell *lc; + + if (attr->attgenerated) + continue; + + appendStringInfoString(buf, first ? "(" : ", "); + first = false; + + /* Use attribute name or column_name option. */ + colname = NameStr(attr->attname); + options = GetForeignColumnOptions(relid, attnum); + foreach(lc, options) + { + DefElem *def = (DefElem *) lfirst(lc); + + if (strcmp(def->defname, "column_name") == 0) + { + colname = defGetString(def); + break; + } + } + + appendStringInfoString(buf, quote_identifier(colname)); + } + if (!first) + appendStringInfoChar(buf, ')'); + + appendStringInfoString(buf, " FROM STDIN (FORMAT TEXT)"); +} + + /* * rebuild remote INSERT statement * diff --git a/contrib/postgres_fdw/expected/postgres_fdw.out b/contrib/postgres_fdw/expected/postgres_fdw.out index 517d15cf1fa..c9358d7a019 100644 --- a/contrib/postgres_fdw/expected/postgres_fdw.out +++ b/contrib/postgres_fdw/expected/postgres_fdw.out @@ -8454,6 +8454,32 @@ select * from grem1; (2 rows) delete from grem1; +-- test that fdw also use COPY FROM as a remote sql +alter foreign table grem1 options (add batch_size '10'); -- add batch size to enable COPY +set client_min_messages to 'log'; +create function insert_or_copy() returns trigger as $$ +declare query text; +begin + query := current_query(); + raise notice '%', query; +return new; +end; +$$ language plpgsql; +CREATE TRIGGER trig_row_before +BEFORE INSERT OR UPDATE OR DELETE ON gloc1 +FOR EACH ROW EXECUTE PROCEDURE insert_or_copy(); +copy grem1 from stdin; +LOG: received message via remote connection: NOTICE: COPY public.gloc1(a) FROM STDIN (FORMAT TEXT) +drop trigger trig_row_before on gloc1; +reset client_min_messages; +alter foreign table grem1 options (drop batch_size); +-- test that copy does not fail with column_name alias +create table gloc2(xxx int); +create foreign table grem2(a int) server loopback options(table_name 'gloc2'); +alter foreign table grem2 alter column a options (column_name 'xxx'); +alter foreign table grem2 options (add batch_size '10'); -- enable COPY +copy grem2 from stdin; +alter foreign table grem2 options (drop batch_size); -- test batch insert alter server loopback options (add batch_size '10'); explain (verbose, costs off) @@ -8471,16 +8497,18 @@ insert into grem1 (a) values (1), (2); select * from gloc1; a | b | c ---+---+--- + 3 | 6 | 1 | 2 | 2 | 4 | -(2 rows) +(3 rows) select * from grem1; a | b | c ---+---+--- + 3 | 6 | 9 1 | 2 | 3 2 | 4 | 6 -(2 rows) +(3 rows) delete from grem1; -- batch insert with foreign partitions. @@ -8505,6 +8533,12 @@ select count(*) from tab_batch_sharded; drop table tab_batch_local; drop table tab_batch_sharded; drop table tab_batch_sharded_p1_remote; +-- test batch insert using copy +set client_min_messages to 'debug1'; +copy grem1 from stdin; +DEBUG: postgres_fdw: COPY 10 rows into foreign table +DEBUG: postgres_fdw: COPY 2 rows into foreign table +reset client_min_messages; alter server loopback options (drop batch_size); -- =================================================================== -- test local triggers @@ -10387,6 +10421,18 @@ select * from rem2; (2 rows) delete from rem2; +-- Test invalid data types to assert that COPY protocol is finished on +-- transaction abort +alter foreign table rem2 options (add batch_size '10'); -- enable COPY +copy rem2 from stdin; +ERROR: invalid input syntax for type integer: "bar" +CONTEXT: COPY rem2, line 2, column f1: "bar" +alter foreign table rem2 options (drop batch_size); +select * from rem2; + f1 | f2 +----+---- +(0 rows) + -- Test check constraints alter table loc2 add constraint loc2_f1positive check (f1 >= 0); alter foreign table rem2 add constraint rem2_f1positive check (f1 >= 0); @@ -10552,7 +10598,8 @@ copy rem2 from stdin; copy rem2 from stdin; -- ERROR ERROR: new row for relation "loc2" violates check constraint "loc2_f1positive" DETAIL: Failing row contains (-1, xyzzy). -CONTEXT: remote SQL command: INSERT INTO public.loc2(f1, f2) VALUES ($1, $2) +CONTEXT: COPY loc2, line 1: "-1 xyzzy" +remote SQL command: COPY public.loc2(f1, f2) FROM STDIN (FORMAT TEXT) COPY rem2 select * from rem2; f1 | f2 @@ -10591,6 +10638,45 @@ select * from rem2; drop trigger trig_null on loc2; delete from rem2; +-- Test COPY FROM with column list and special characters +alter foreign table rem2 options (add batch_size '10'); -- add batch size to enable COPY +copy rem2 (f1, f2) from stdin; +select * from rem2; + f1 | f2 +----+------- + 1 | hello+ + | world +(1 row) + +delete from rem2; +-- Test COPY with NULL and special characters +copy rem2 from stdin; +select * from rem2; + f1 | f2 +----+----- + 1 | + | bar + 3 | a"b +(3 rows) + +delete from rem2; +alter foreign table rem2 options (drop batch_size); +-- Test that float numbers do not loose precision when sending to the foreign +-- server +create table f(a float); +create foreign table f_fdw(a float) server loopback options(table_name 'f'); +alter foreign table f_fdw options (add batch_size '10'); -- enable COPY +set extra_float_digits = 0; +copy f_fdw from stdin; +reset extra_float_digits; +alter foreign table f_fdw options (drop batch_size); +select * from f; + a +-------------------- + 1.0000000000000002 +(1 row) + +drop table f; -- Check with zero-column foreign table; batch insert will be disabled alter table loc2 drop column f1; alter table loc2 drop column f2; diff --git a/contrib/postgres_fdw/postgres_fdw.c b/contrib/postgres_fdw/postgres_fdw.c index 47197a733ff..7d290f1035e 100644 --- a/contrib/postgres_fdw/postgres_fdw.c +++ b/contrib/postgres_fdw/postgres_fdw.c @@ -18,6 +18,7 @@ #include "access/sysattr.h" #include "access/table.h" #include "catalog/pg_opfamily.h" +#include "commands/copy.h" #include "commands/defrem.h" #include "commands/explain_format.h" #include "commands/explain_state.h" @@ -26,6 +27,7 @@ #include "executor/instrument.h" #include "foreign/fdwapi.h" #include "funcapi.h" +#include "mb/pg_wchar.h" #include "miscadmin.h" #include "nodes/makefuncs.h" #include "nodes/nodeFuncs.h" @@ -70,6 +72,9 @@ PG_MODULE_MAGIC_EXT( /* If no remote estimates, assume a sort costs 20% extra */ #define DEFAULT_FDW_SORT_MULTIPLIER 1.2 +/* Buffer size to send COPY IN data */ +#define COPYBUFSIZ 8192 + /* * Indexes of FDW-private information stored in fdw_private lists. * @@ -459,6 +464,12 @@ static void postgresBeginForeignInsert(ModifyTableState *mtstate, ResultRelInfo *resultRelInfo); static void postgresEndForeignInsert(EState *estate, ResultRelInfo *resultRelInfo); + +static void postgresExecForeignBatchCopy(EState *estate, + ResultRelInfo *resultRelInfo, + TupleTableSlot **slots, + int numSlots); + static int postgresIsForeignRelUpdatable(Relation rel); static bool postgresPlanDirectModify(PlannerInfo *root, ModifyTable *plan, @@ -666,6 +677,9 @@ static void merge_fdw_options(PgFdwRelationInfo *fpinfo, const PgFdwRelationInfo *fpinfo_o, const PgFdwRelationInfo *fpinfo_i); static int get_batch_size_option(Relation rel); +static void convert_slot_to_copy_text(StringInfo buf, + PgFdwModifyState *fmstate, + TupleTableSlot *slot); /* @@ -698,6 +712,7 @@ postgres_fdw_handler(PG_FUNCTION_ARGS) routine->EndForeignModify = postgresEndForeignModify; routine->BeginForeignInsert = postgresBeginForeignInsert; routine->EndForeignInsert = postgresEndForeignInsert; + routine->ExecForeignBatchCopy = postgresExecForeignBatchCopy; routine->IsForeignRelUpdatable = postgresIsForeignRelUpdatable; routine->PlanDirectModify = postgresPlanDirectModify; routine->BeginDirectModify = postgresBeginDirectModify; @@ -2370,6 +2385,80 @@ postgresEndForeignModify(EState *estate, finish_foreign_modify(fmstate); } +/* + * postgresExecForeignBatchCopy + * Insert a batch of tuples into a foreign table using the COPY protocol. + * + * This is a self-contained operation: it opens a COPY ... FROM STDIN on the + * remote connection, streams all the given tuples, and closes the COPY, all + * within a single call. The connection is therefore left idle again once we + * return, which lets several foreign partitions (or local code that runs + * between batches, such as triggers) safely share the same cached connection. + */ +static void +postgresExecForeignBatchCopy(EState *estate, + ResultRelInfo *resultRelInfo, + TupleTableSlot **slots, + int numSlots) +{ + PgFdwModifyState *fmstate = (PgFdwModifyState *) resultRelInfo->ri_FdwState; + StringInfoData buf; + int nestlevel; + int i; + + elog(DEBUG1, "postgres_fdw: COPY %d rows into foreign table", numSlots); + + /* Make sure any constants are printed portably. */ + nestlevel = set_transmission_modes(); + + /* Start COPY if not already in progress on this connection */ + if (!fmstate->conn_state->copy_in_progress) + { + StringInfoData sql; + PGresult *res; + + /* Build the COPY ... FROM STDIN command for this relation. */ + initStringInfo(&sql); + deparseCopySql(&sql, fmstate->rel, fmstate->target_attrs); + + /* Start the COPY protocol. */ + if (!PQsendQuery(fmstate->conn, sql.data)) + pgfdw_report_error(NULL, fmstate->conn, sql.data); + + res = pgfdw_get_result(fmstate->conn); + if (PQresultStatus(res) != PGRES_COPY_IN) + pgfdw_report_error(res, fmstate->conn, sql.data); + PQclear(res); + + fmstate->conn_state->copy_in_progress = true; + + pfree(sql.data); + } + + /* Stream all the rows, flushing the buffer when it grows large. */ + initStringInfo(&buf); + for (i = 0; i < numSlots; i++) + { + convert_slot_to_copy_text(&buf, fmstate, slots[i]); + + if (buf.len >= COPYBUFSIZ) + { + if (PQputCopyData(fmstate->conn, buf.data, buf.len) <= 0) + pgfdw_report_error(NULL, fmstate->conn, NULL); + resetStringInfo(&buf); + } + } + if (buf.len > 0) + { + if (PQputCopyData(fmstate->conn, buf.data, buf.len) <= 0) + pgfdw_report_error(NULL, fmstate->conn, NULL); + } + + pfree(buf.data); + + reset_transmission_modes(nestlevel); +} + /* * postgresBeginForeignInsert * Begin an insert operation on a foreign table @@ -9266,3 +9355,56 @@ get_batch_size_option(Relation rel) return batch_size; } + +/* + * Write target attribute values from fmstate into buf buffer to be sent as + * COPY FROM STDIN data + */ +static void +convert_slot_to_copy_text(StringInfo buf, + PgFdwModifyState *fmstate, + TupleTableSlot *slot) +{ + TupleDesc tupdesc = RelationGetDescr(fmstate->rel); + bool first = true; + int i = 0; + + foreach_int(attnum, fmstate->target_attrs) + { + CompactAttribute *attr = TupleDescCompactAttr(tupdesc, attnum - 1); + Datum datum; + bool isnull; + + /* Ignore generated columns; they are set to DEFAULT */ + if (attr->attgenerated) + continue; + + if (!first) + appendStringInfoCharMacro(buf, '\t'); + first = false; + + datum = slot_getattr(slot, attnum, &isnull); + + if (isnull) + appendStringInfoString(buf, "\\N"); + else + { + const char *value = OutputFunctionCall(&fmstate->p_flinfo[i], + datum); + + /* + * Append a string to buf, escaping special characters for COPY + * TEXT format. + */ + CopyEscapeText(buf, + value, + '\t', + GetDatabaseEncoding(), + false, + false); + } + i++; + } + + appendStringInfoCharMacro(buf, '\n'); +} diff --git a/contrib/postgres_fdw/postgres_fdw.h b/contrib/postgres_fdw/postgres_fdw.h index da7da1c2ea9..3135ba8a7cf 100644 --- a/contrib/postgres_fdw/postgres_fdw.h +++ b/contrib/postgres_fdw/postgres_fdw.h @@ -147,6 +147,8 @@ typedef struct PgFdwRelationInfo typedef struct PgFdwConnState { AsyncRequest *pendingAreq; /* pending async request */ + bool copy_in_progress; /* true while a COPY FROM STDIN is open on + * the connection */ } PgFdwConnState; /* @@ -216,6 +218,7 @@ extern void rebuildInsertSql(StringInfo buf, Relation rel, char *orig_query, List *target_attrs, int values_end_len, int num_params, int num_rows); +extern void deparseCopySql(StringInfo buf, Relation rel, List *target_attrs); extern void deparseUpdateSql(StringInfo buf, RangeTblEntry *rte, Index rtindex, Relation rel, List *targetAttrs, diff --git a/contrib/postgres_fdw/sql/postgres_fdw.sql b/contrib/postgres_fdw/sql/postgres_fdw.sql index ec766e2b28a..193ff49806b 100644 --- a/contrib/postgres_fdw/sql/postgres_fdw.sql +++ b/contrib/postgres_fdw/sql/postgres_fdw.sql @@ -2337,6 +2337,41 @@ select * from gloc1; select * from grem1; delete from grem1; +-- test that fdw also use COPY FROM as a remote sql +alter foreign table grem1 options (add batch_size '10'); -- add batch size to enable COPY +set client_min_messages to 'log'; + +create function insert_or_copy() returns trigger as $$ +declare query text; +begin + query := current_query(); + raise notice '%', query; +return new; +end; +$$ language plpgsql; + +CREATE TRIGGER trig_row_before +BEFORE INSERT OR UPDATE OR DELETE ON gloc1 +FOR EACH ROW EXECUTE PROCEDURE insert_or_copy(); + +copy grem1 from stdin; +3 +\. + +drop trigger trig_row_before on gloc1; +reset client_min_messages; +alter foreign table grem1 options (drop batch_size); + +-- test that copy does not fail with column_name alias +create table gloc2(xxx int); +create foreign table grem2(a int) server loopback options(table_name 'gloc2'); +alter foreign table grem2 alter column a options (column_name 'xxx'); +alter foreign table grem2 options (add batch_size '10'); -- enable COPY +copy grem2 from stdin; +1 +\. +alter foreign table grem2 options (drop batch_size); + -- test batch insert alter server loopback options (add batch_size '10'); explain (verbose, costs off) @@ -2363,6 +2398,24 @@ drop table tab_batch_local; drop table tab_batch_sharded; drop table tab_batch_sharded_p1_remote; +-- test batch insert using copy +set client_min_messages to 'debug1'; +copy grem1 from stdin; +1 +2 +3 +4 +5 +6 +7 +8 +9 +10 +11 +12 +\. +reset client_min_messages; + alter server loopback options (drop batch_size); -- =================================================================== @@ -3230,6 +3283,16 @@ select * from rem2; delete from rem2; +-- Test invalid data types to assert that COPY protocol is finished on +-- transaction abort +alter foreign table rem2 options (add batch_size '10'); -- enable COPY +copy rem2 from stdin; +1 foo +bar 2 +\. +alter foreign table rem2 options (drop batch_size); +select * from rem2; + -- Test check constraints alter table loc2 add constraint loc2_f1positive check (f1 >= 0); alter foreign table rem2 add constraint rem2_f1positive check (f1 >= 0); @@ -3427,6 +3490,43 @@ drop trigger trig_null on loc2; delete from rem2; +-- Test COPY FROM with column list and special characters +alter foreign table rem2 options (add batch_size '10'); -- add batch size to enable COPY +copy rem2 (f1, f2) from stdin; +1 hello\nworld +\. +select * from rem2; + +delete from rem2; + +-- Test COPY with NULL and special characters +copy rem2 from stdin; +1 \N +\N bar +3 a"b +\. +select * from rem2; + +delete from rem2; + +alter foreign table rem2 options (drop batch_size); + +-- Test that float numbers do not loose precision when sending to the foreign +-- server +create table f(a float); +create foreign table f_fdw(a float) server loopback options(table_name 'f'); + +alter foreign table f_fdw options (add batch_size '10'); -- enable COPY +set extra_float_digits = 0; +copy f_fdw from stdin; +1.0000000000000002 +\. +reset extra_float_digits; +alter foreign table f_fdw options (drop batch_size); +select * from f; + +drop table f; + -- Check with zero-column foreign table; batch insert will be disabled alter table loc2 drop column f1; alter table loc2 drop column f2; diff --git a/doc/src/sgml/fdwhandler.sgml b/doc/src/sgml/fdwhandler.sgml index 502441fefcb..8a15cdc4b01 100644 --- a/doc/src/sgml/fdwhandler.sgml +++ b/doc/src/sgml/fdwhandler.sgml @@ -908,6 +908,55 @@ EndForeignInsert(EState *estate, +void +ExecForeignBatchCopy(EState *estate, + ResultRelInfo *rinfo, + TupleTableSlot **slots, + int numSlots); + + + Insert a batch of tuples into a foreign table during a + COPY FROM command, using an optimized bulk-load + protocol (such as the COPY protocol for remote + PostgreSQL servers) instead of inserting the + tuples one at a time. + estate is global execution state for the query. + rinfo is the ResultRelInfo struct + describing the target foreign table. + slots is an array of tuple table slots containing the + tuples to be inserted; numSlots is the number of tuples + in the array. + + + + This callback is used only when a COPY FROM is executed + directly on a foreign table, or is routed into a foreign-table partition, + and batching is enabled for the table (that is, the + GetForeignModifyBatchSize callback reports a batch + size greater than one). It is not used when the target has + AFTER ROW triggers, because the transferred rows are not + returned to feed those triggers. In all other cases the core code falls + back to ExecForeignBatchInsert or + ExecForeignInsert. + + + + Each call is expected to be self-contained: the callback should open the + bulk-load session, transfer all numSlots tuples, and + close the session before returning. This leaves the remote connection idle + between calls, which allows several foreign partitions to share one + connection safely. + + + + If the ExecForeignBatchCopy pointer is set to + NULL, COPY FROM falls back to + inserting tuples through + ExecForeignInsert/ExecForeignBatchInsert. + + + + int IsForeignRelUpdatable(Relation rel); diff --git a/doc/src/sgml/postgres-fdw.sgml b/doc/src/sgml/postgres-fdw.sgml index 87b1433aacb..0853c734ca5 100644 --- a/doc/src/sgml/postgres-fdw.sgml +++ b/doc/src/sgml/postgres-fdw.sgml @@ -467,11 +467,13 @@ OPTIONS (ADD password_required 'false'); - This option also applies when copying into foreign tables. In that case - the actual number of rows postgres_fdw copies at - once is determined in a similar way to the insert case, but it is - limited to at most 1000 due to implementation restrictions of the - COPY command. + This option also applies when copying into foreign tables. When + COPY FROM is executed on a foreign table (or routed + into a foreign-table partition) that has this option set greater than + one and has no row-level triggers, postgres_fdw uses + the COPY protocol to transfer data to the remote + server, which is significantly faster than inserting rows individually. + Otherwise the rows are inserted individually. diff --git a/src/backend/commands/copyfrom.c b/src/backend/commands/copyfrom.c index 3782db171be..63c5bd1501e 100644 --- a/src/backend/commands/copyfrom.c +++ b/src/backend/commands/copyfrom.c @@ -464,7 +464,6 @@ CopyMultiInsertBufferFlush(CopyMultiInsertInfo *miinfo, Assert(buffer->bistate == NULL); /* Ensure that the FDW supports batching and it's enabled */ - Assert(resultRelInfo->ri_FdwRoutine->ExecForeignBatchInsert); Assert(batch_size > 1); /* @@ -474,6 +473,40 @@ CopyMultiInsertBufferFlush(CopyMultiInsertInfo *miinfo, Assert(!cstate->relname_only); cstate->relname_only = true; + if (resultRelInfo->ri_FdwRoutine->ExecForeignBatchCopy != NULL && + (resultRelInfo->ri_TrigDesc == NULL || + !resultRelInfo->ri_TrigDesc->trig_insert_after_row)) + { + /* + * Send the buffered tuples to the FDW in batches of at most + * batch_size, as we do for ExecForeignBatchInsert. Each call is a + * self-contained COPY operation. + * + * COPY provides no RETURNING, so this path is only usable when + * there are no AFTER ROW triggers that would need the stored rows. + */ + while (sent < nused) + { + int size = (batch_size < nused - sent) ? batch_size : (nused - sent); + + resultRelInfo->ri_FdwRoutine->ExecForeignBatchCopy(estate, + resultRelInfo, + &slots[sent], + size); + + sent += size; + + /* Update the row counter and progress of the COPY command */ + *processed += size; + pgstat_progress_update_param(PROGRESS_COPY_TUPLES_PROCESSED, + *processed); + } + } + else + { + /* Ensure that the FDW supports batching and it's enabled */ + Assert(resultRelInfo->ri_FdwRoutine->ExecForeignBatchInsert); + while (sent < nused) { int size = (batch_size < nused - sent) ? batch_size : (nused - sent); @@ -525,6 +558,7 @@ CopyMultiInsertBufferFlush(CopyMultiInsertInfo *miinfo, pgstat_progress_update_param(PROGRESS_COPY_TUPLES_PROCESSED, *processed); } + } for (i = 0; i < nused; i++) ExecClearTuple(slots[i]); @@ -774,6 +808,34 @@ CopyMultiInsertInfoStore(CopyMultiInsertInfo *miinfo, ResultRelInfo *rri, miinfo->bufferedBytes += tuplen; } +/* + * Can this foreign result relation accept batches of tuples during COPY FROM, + * i.e. does CopyMultiInsertBufferFlush() have a usable path for it? + * + * The FDW can batch if it provides ExecForeignBatchInsert, or if it provides + * ExecForeignBatchCopy and there are no AFTER ROW triggers to feed (COPY has no + * RETURNING, so the stored rows are not available for them). These two + * callbacks are independent: an FDW may implement either or both. + * + * Callers must have already established that batching is enabled for the + * relation (ri_BatchSize > 1). + */ +static bool +CopyFromFdwCanBatch(ResultRelInfo *resultRelInfo) +{ + FdwRoutine *fdwroutine = resultRelInfo->ri_FdwRoutine; + + if (fdwroutine->ExecForeignBatchInsert != NULL) + return true; + + if (fdwroutine->ExecForeignBatchCopy != NULL && + (resultRelInfo->ri_TrigDesc == NULL || + !resultRelInfo->ri_TrigDesc->trig_insert_after_row)) + return true; + + return false; +} + /* * Copy FROM file to relation. */ @@ -951,7 +1013,8 @@ CopyFrom(CopyFromState cstate) */ if (resultRelInfo->ri_FdwRoutine != NULL && resultRelInfo->ri_FdwRoutine->GetForeignModifyBatchSize && - resultRelInfo->ri_FdwRoutine->ExecForeignBatchInsert) + (resultRelInfo->ri_FdwRoutine->ExecForeignBatchInsert || + resultRelInfo->ri_FdwRoutine->ExecForeignBatchCopy)) resultRelInfo->ri_BatchSize = resultRelInfo->ri_FdwRoutine->GetForeignModifyBatchSize(resultRelInfo); else @@ -1007,11 +1070,14 @@ CopyFrom(CopyFromState cstate) insertMethod = CIM_SINGLE; } else if (resultRelInfo->ri_FdwRoutine != NULL && - resultRelInfo->ri_BatchSize == 1) + (resultRelInfo->ri_BatchSize == 1 || + !CopyFromFdwCanBatch(resultRelInfo))) { /* * Can't support multi-inserts to a foreign table if the FDW does not - * support batching, or it's disabled for the server or foreign table. + * support batching, or it's disabled for the server or foreign table, + * or the FDW only offers COPY-based batching but the table has AFTER + * ROW triggers to feed. */ insertMethod = CIM_SINGLE; } @@ -1235,7 +1301,8 @@ CopyFrom(CopyFromState cstate) !has_before_insert_row_trig && !has_instead_insert_row_trig && (resultRelInfo->ri_FdwRoutine == NULL || - resultRelInfo->ri_BatchSize > 1); + (resultRelInfo->ri_BatchSize > 1 && + CopyFromFdwCanBatch(resultRelInfo))); /* Set the multi-insert buffer to use for this partition. */ if (leafpart_use_multi_insert) diff --git a/src/backend/commands/copyto.c b/src/backend/commands/copyto.c index 5850608a3fb..94c04baa4e7 100644 --- a/src/backend/commands/copyto.c +++ b/src/backend/commands/copyto.c @@ -571,6 +571,12 @@ SendCopyEnd(CopyToState cstate) pq_putemptymessage(PqMsg_CopyDone); } +#define CopySendCharBuf(buf, c) \ + appendStringInfoCharMacro(buf, c) + +#define CopySendDataBuf(buf, databuf, datasize) \ + appendBinaryStringInfo(buf, databuf, datasize) + /*---------- * CopySendData sends output data to the destination (file or frontend) * CopySendString does the same for null-terminated strings @@ -584,7 +590,7 @@ SendCopyEnd(CopyToState cstate) static void CopySendData(CopyToState cstate, const void *databuf, int datasize) { - appendBinaryStringInfo(cstate->fe_msgbuf, databuf, datasize); + CopySendDataBuf(cstate->fe_msgbuf, databuf, datasize); } static void @@ -596,7 +602,7 @@ CopySendString(CopyToState cstate, const char *str) static void CopySendChar(CopyToState cstate, char c) { - appendStringInfoCharMacro(cstate->fe_msgbuf, c); + CopySendCharBuf(cstate->fe_msgbuf, c); } static void @@ -1439,16 +1445,44 @@ CopyOneRowTo(CopyToState cstate, TupleTableSlot *slot) CopySendData(cstate, start, ptr - start); \ } while (0) -static void -CopyAttributeOutText(CopyToState cstate, const char *string) +/* Like above, but it works with a string buffer */ +#define DUMPSOFAR_TO_BUF() \ + do { \ + if (ptr > start) \ + CopySendDataBuf(buf, start, ptr - start); \ + } while (0) + + +/* + * Escape a string for COPY TEXT format output + * + * Escapes control characters, backslashes, and the delimiter character + * according to COPY TEXT format rules. The escaped string is appended + * to 'buf'. + * + * Parameters: + * buf - StringInfo buffer to append the escaped text to + * string - the input string to escape + * delimc - the delimiter character that must be escaped + * file_encoding - target encoding for the output + * need_transcoding - if true, convert from server encoding to file_encoding + * encoding_embeds_ascii - if true, the encoding may have ASCII bytes as + * non-first bytes of multi-byte characters + */ +void +CopyEscapeText(StringInfo buf, + const char *string, + char delimc, + int file_encoding, + bool need_transcoding, + bool encoding_embeds_ascii) { const char *ptr; const char *start; char c; - char delimc = cstate->opts.delim[0]; - if (cstate->need_transcoding) - ptr = pg_server_to_any(string, strlen(string), cstate->file_encoding); + if (need_transcoding) + ptr = pg_server_to_any(string, strlen(string), file_encoding); else ptr = string; @@ -1466,7 +1500,7 @@ CopyAttributeOutText(CopyToState cstate, const char *string) * it's worth making two copies of it to get the IS_HIGHBIT_SET() test out * of the normal safe-encoding path. */ - if (cstate->encoding_embeds_ascii) + if (encoding_embeds_ascii) { start = ptr; while ((c = *ptr) != '\0') @@ -1509,19 +1543,19 @@ CopyAttributeOutText(CopyToState cstate, const char *string) continue; /* fall to end of loop */ } /* if we get here, we need to convert the control char */ - DUMPSOFAR(); - CopySendChar(cstate, '\\'); - CopySendChar(cstate, c); + DUMPSOFAR_TO_BUF(); + CopySendCharBuf(buf, '\\'); + CopySendCharBuf(buf, c); start = ++ptr; /* do not include char in next run */ } else if (c == '\\' || c == delimc) { - DUMPSOFAR(); - CopySendChar(cstate, '\\'); + DUMPSOFAR_TO_BUF(); + CopySendCharBuf(buf, '\\'); start = ptr++; /* we include char in next run */ } else if (IS_HIGHBIT_SET(c)) - ptr += pg_encoding_mblen(cstate->file_encoding, ptr); + ptr += pg_encoding_mblen(file_encoding, ptr); else ptr++; } @@ -1569,15 +1603,15 @@ CopyAttributeOutText(CopyToState cstate, const char *string) continue; /* fall to end of loop */ } /* if we get here, we need to convert the control char */ - DUMPSOFAR(); - CopySendChar(cstate, '\\'); - CopySendChar(cstate, c); + DUMPSOFAR_TO_BUF(); + CopySendCharBuf(buf, '\\'); + CopySendCharBuf(buf, c); start = ++ptr; /* do not include char in next run */ } else if (c == '\\' || c == delimc) { - DUMPSOFAR(); - CopySendChar(cstate, '\\'); + DUMPSOFAR_TO_BUF(); + CopySendCharBuf(buf, '\\'); start = ptr++; /* we include char in next run */ } else @@ -1585,7 +1619,18 @@ CopyAttributeOutText(CopyToState cstate, const char *string) } } - DUMPSOFAR(); + DUMPSOFAR_TO_BUF(); +} + +static void +CopyAttributeOutText(CopyToState cstate, const char *string) +{ + CopyEscapeText(cstate->fe_msgbuf, + string, + cstate->opts.delim[0], + cstate->file_encoding, + cstate->need_transcoding, + cstate->encoding_embeds_ascii); } /* diff --git a/src/include/commands/copy.h b/src/include/commands/copy.h index abecfe51098..27ee58c7fdb 100644 --- a/src/include/commands/copy.h +++ b/src/include/commands/copy.h @@ -136,5 +136,11 @@ extern void EndCopyTo(CopyToState cstate); extern uint64 DoCopyTo(CopyToState cstate); extern List *CopyGetAttnums(TupleDesc tupDesc, Relation rel, List *attnamelist); +extern void CopyEscapeText(StringInfo buf, + const char *string, + char delimc, + int file_encoding, + bool need_transcoding, + bool encoding_embeds_ascii); #endif /* COPY_H */ diff --git a/src/include/foreign/fdwapi.h b/src/include/foreign/fdwapi.h index abf59a0d8ad..0439fa8b24e 100644 --- a/src/include/foreign/fdwapi.h +++ b/src/include/foreign/fdwapi.h @@ -113,6 +113,11 @@ typedef void (*BeginForeignInsert_function) (ModifyTableState *mtstate, typedef void (*EndForeignInsert_function) (EState *estate, ResultRelInfo *rinfo); +typedef void (*ExecForeignBatchCopy_function) (EState *estate, + ResultRelInfo *rinfo, + TupleTableSlot **slots, + int numSlots); + typedef int (*IsForeignRelUpdatable_function) (Relation rel); typedef bool (*PlanDirectModify_function) (PlannerInfo *root, @@ -241,6 +246,7 @@ typedef struct FdwRoutine EndForeignModify_function EndForeignModify; BeginForeignInsert_function BeginForeignInsert; EndForeignInsert_function EndForeignInsert; + ExecForeignBatchCopy_function ExecForeignBatchCopy; IsForeignRelUpdatable_function IsForeignRelUpdatable; PlanDirectModify_function PlanDirectModify; BeginDirectModify_function BeginDirectModify;