From 18618368721678d78934251ff8243705013458f0 Mon Sep 17 00:00:00 2001 From: Sutou Kouhei Date: Thu, 27 Mar 2025 11:24:15 +0900 Subject: [PATCH v40 3/6] Add support for implementing custom COPY handler as extension * TO: Add CopyToStateData::opaque that can be used to keep data for custom COPY TO handler implementation * TO: Export CopySendEndOfRow() to send end of row data as CopyToStateFlush() * FROM: Add CopyFromStateData::opaque that can be used to keep data for custom COPY FROM handler implementation * FROM: Export CopyGetData() to get the next data as CopyFromStateGetData() * FROM: Add CopyFromSkipErrorRow() for "ON_ERROR stop" and "LOG_VERBOSITY verbose" COPY FROM extensions must call CopyFromSkipErrorRow() when CopyFromOneRow callback reports an error by errsave(). CopyFromSkipErrorRow() handles "ON_ERROR stop" and "LOG_VERBOSITY verbose" cases. --- src/backend/commands/copyfromparse.c | 93 ++++++++++++------- src/backend/commands/copyto.c | 12 +++ src/include/commands/copyapi.h | 6 ++ src/include/commands/copyfrom_internal.h | 3 + src/include/commands/copyto_internal.h | 3 + .../expected/test_copy_format.out | 50 ++++++++++ .../test_copy_format/sql/test_copy_format.sql | 35 +++++++ .../test_copy_format/test_copy_format.c | 80 +++++++++++++++- 8 files changed, 245 insertions(+), 37 deletions(-) diff --git a/src/backend/commands/copyfromparse.c b/src/backend/commands/copyfromparse.c index 9f7171d1478..de68b53b000 100644 --- a/src/backend/commands/copyfromparse.c +++ b/src/backend/commands/copyfromparse.c @@ -739,6 +739,17 @@ CopyReadBinaryData(CopyFromState cstate, char *dest, int nbytes) return copied_bytes; } +/* + * Export CopyGetData() for extensions. We want to keep CopyGetData() as a + * static function for optimization. CopyGetData() calls in this file may be + * optimized by a compiler. + */ +int +CopyFromStateGetData(CopyFromState cstate, void *dest, int minread, int maxread) +{ + return CopyGetData(cstate, dest, minread, maxread); +} + /* * This function is exposed for use by extensions that read raw fields in the * next line. See NextCopyFromRawFieldsInternal() for details. @@ -927,6 +938,51 @@ CopyFromCSVOneRow(CopyFromState cstate, ExprContext *econtext, Datum *values, return CopyFromTextLikeOneRow(cstate, econtext, values, nulls, true); } +/* + * Call this when you report an error by errsave() in your CopyFromOneRow + * callback. This handles "ON_ERROR stop" and "LOG_VERBOSITY verbose" cases + * for you. + */ +void +CopyFromSkipErrorRow(CopyFromState cstate) +{ + Assert(cstate->opts.on_error != COPY_ON_ERROR_STOP); + + cstate->num_errors++; + + if (cstate->opts.log_verbosity == COPY_LOG_VERBOSITY_VERBOSE) + { + /* + * Since we emit line number and column info in the below notice + * message, we suppress error context information other than the + * relation name. + */ + Assert(!cstate->relname_only); + cstate->relname_only = true; + + if (cstate->cur_attval) + { + char *attval; + + attval = CopyLimitPrintoutLength(cstate->cur_attval); + ereport(NOTICE, + errmsg("skipping row due to data type incompatibility at line %" PRIu64 " for column \"%s\": \"%s\"", + cstate->cur_lineno, + cstate->cur_attname, + attval)); + pfree(attval); + } + else + ereport(NOTICE, + errmsg("skipping row due to data type incompatibility at line %" PRIu64 " for column \"%s\": null input", + cstate->cur_lineno, + cstate->cur_attname)); + + /* reset relname_only */ + cstate->relname_only = false; + } +} + /* * Workhorse for CopyFromTextOneRow() and CopyFromCSVOneRow(). * @@ -1033,42 +1089,7 @@ CopyFromTextLikeOneRow(CopyFromState cstate, ExprContext *econtext, (Node *) cstate->escontext, &values[m])) { - Assert(cstate->opts.on_error != COPY_ON_ERROR_STOP); - - cstate->num_errors++; - - if (cstate->opts.log_verbosity == COPY_LOG_VERBOSITY_VERBOSE) - { - /* - * Since we emit line number and column info in the below - * notice message, we suppress error context information other - * than the relation name. - */ - Assert(!cstate->relname_only); - cstate->relname_only = true; - - if (cstate->cur_attval) - { - char *attval; - - attval = CopyLimitPrintoutLength(cstate->cur_attval); - ereport(NOTICE, - errmsg("skipping row due to data type incompatibility at line %" PRIu64 " for column \"%s\": \"%s\"", - cstate->cur_lineno, - cstate->cur_attname, - attval)); - pfree(attval); - } - else - ereport(NOTICE, - errmsg("skipping row due to data type incompatibility at line %" PRIu64 " for column \"%s\": null input", - cstate->cur_lineno, - cstate->cur_attname)); - - /* reset relname_only */ - cstate->relname_only = false; - } - + CopyFromSkipErrorRow(cstate); return true; } diff --git a/src/backend/commands/copyto.c b/src/backend/commands/copyto.c index 265b847e255..d6fcfdfb9b1 100644 --- a/src/backend/commands/copyto.c +++ b/src/backend/commands/copyto.c @@ -454,6 +454,18 @@ CopySendEndOfRow(CopyToState cstate) resetStringInfo(fe_msgbuf); } +/* + * Export CopySendEndOfRow() for extensions. We want to keep + * CopySendEndOfRow() as a static function for + * optimization. CopySendEndOfRow() calls in this file may be optimized by a + * compiler. + */ +void +CopyToStateFlush(CopyToState cstate) +{ + CopySendEndOfRow(cstate); +} + /* * Wrapper function of CopySendEndOfRow for text and CSV formats. Sends the * line termination and do common appropriate things for the end of row. diff --git a/src/include/commands/copyapi.h b/src/include/commands/copyapi.h index 53ad3337f86..500ece7d5bb 100644 --- a/src/include/commands/copyapi.h +++ b/src/include/commands/copyapi.h @@ -56,6 +56,8 @@ typedef struct CopyToRoutine void (*CopyToEnd) (CopyToState cstate); } CopyToRoutine; +extern void CopyToStateFlush(CopyToState cstate); + /* * API structure for a COPY FROM format implementation. Note this must be * allocated in a server-lifetime manner, typically as a static const struct. @@ -106,4 +108,8 @@ typedef struct CopyFromRoutine void (*CopyFromEnd) (CopyFromState cstate); } CopyFromRoutine; +extern int CopyFromStateGetData(CopyFromState cstate, void *dest, int minread, int maxread); + +extern void CopyFromSkipErrorRow(CopyFromState cstate); + #endif /* COPYAPI_H */ diff --git a/src/include/commands/copyfrom_internal.h b/src/include/commands/copyfrom_internal.h index 24157e11a73..f9e27152313 100644 --- a/src/include/commands/copyfrom_internal.h +++ b/src/include/commands/copyfrom_internal.h @@ -181,6 +181,9 @@ typedef struct CopyFromStateData #define RAW_BUF_BYTES(cstate) ((cstate)->raw_buf_len - (cstate)->raw_buf_index) uint64 bytes_processed; /* number of bytes processed so far */ + + /* For custom format implementation */ + void *opaque; /* private space */ } CopyFromStateData; extern void ReceiveCopyBegin(CopyFromState cstate); diff --git a/src/include/commands/copyto_internal.h b/src/include/commands/copyto_internal.h index da796131988..3bd9d702bf0 100644 --- a/src/include/commands/copyto_internal.h +++ b/src/include/commands/copyto_internal.h @@ -78,6 +78,9 @@ typedef struct CopyToStateData FmgrInfo *out_functions; /* lookup info for output functions */ MemoryContext rowcontext; /* per-row evaluation context */ uint64 bytes_processed; /* number of bytes processed so far */ + + /* For custom format implementation */ + void *opaque; /* private space */ } CopyToStateData; #endif /* COPYTO_INTERNAL_H */ diff --git a/src/test/modules/test_copy_format/expected/test_copy_format.out b/src/test/modules/test_copy_format/expected/test_copy_format.out index 3916b766615..47a875f0ab1 100644 --- a/src/test/modules/test_copy_format/expected/test_copy_format.out +++ b/src/test/modules/test_copy_format/expected/test_copy_format.out @@ -4,6 +4,8 @@ INSERT INTO copy_data VALUES (1, 2, 3), (12, 34, 56), (123, 456, 789); -- schema. CREATE EXTENSION test_copy_format; -- We can find a custom COPY handler without schema. +-- 987 is accepted. +-- 654 is a hard error because ON_ERROR is stop by default. COPY copy_data FROM stdin WITH (FORMAT 'test_copy_format'); NOTICE: test_copy_format: is_from=true NOTICE: CopyFromInFunc: attribute: smallint @@ -11,7 +13,50 @@ NOTICE: CopyFromInFunc: attribute: integer NOTICE: CopyFromInFunc: attribute: bigint NOTICE: CopyFromStart: the number of attributes: 3 NOTICE: CopyFromOneRow +NOTICE: CopyFromOneRow +ERROR: invalid value: "6" +CONTEXT: COPY copy_data, line 2, column a: "6" +-- 987 is accepted. +-- 654 is a soft error because ON_ERROR is ignore. +COPY copy_data FROM stdin WITH (FORMAT 'test_copy_format', ON_ERROR ignore); +NOTICE: test_copy_format: is_from=true +NOTICE: CopyFromInFunc: attribute: smallint +NOTICE: CopyFromInFunc: attribute: integer +NOTICE: CopyFromInFunc: attribute: bigint +NOTICE: CopyFromStart: the number of attributes: 3 +NOTICE: CopyFromOneRow +NOTICE: CopyFromOneRow +NOTICE: CopyFromOneRow +NOTICE: 1 row was skipped due to data type incompatibility NOTICE: CopyFromEnd +-- 987 is accepted. +-- 654 is a soft error because ON_ERROR is ignore. +COPY copy_data FROM stdin WITH (FORMAT 'test_copy_format', ON_ERROR ignore, LOG_VERBOSITY verbose); +NOTICE: test_copy_format: is_from=true +NOTICE: CopyFromInFunc: attribute: smallint +NOTICE: CopyFromInFunc: attribute: integer +NOTICE: CopyFromInFunc: attribute: bigint +NOTICE: CopyFromStart: the number of attributes: 3 +NOTICE: CopyFromOneRow +NOTICE: CopyFromOneRow +NOTICE: skipping row due to data type incompatibility at line 2 for column "a": "6" +NOTICE: CopyFromOneRow +NOTICE: 1 row was skipped due to data type incompatibility +NOTICE: CopyFromEnd +-- 987 is accepted. +-- 654 is a soft error because ON_ERROR is ignore. +-- 321 is a hard error. +COPY copy_data FROM stdin WITH (FORMAT 'test_copy_format', ON_ERROR ignore); +NOTICE: test_copy_format: is_from=true +NOTICE: CopyFromInFunc: attribute: smallint +NOTICE: CopyFromInFunc: attribute: integer +NOTICE: CopyFromInFunc: attribute: bigint +NOTICE: CopyFromStart: the number of attributes: 3 +NOTICE: CopyFromOneRow +NOTICE: CopyFromOneRow +NOTICE: CopyFromOneRow +ERROR: too much lines: 3 +CONTEXT: COPY copy_data, line 3 COPY copy_data TO stdout WITH (FORMAT 'test_copy_format'); NOTICE: test_copy_format: is_from=false NOTICE: CopyToOutFunc: attribute: smallint @@ -21,7 +66,12 @@ NOTICE: CopyToStart: the number of attributes: 3 NOTICE: CopyToOneRow: the number of valid values: 3 NOTICE: CopyToOneRow: the number of valid values: 3 NOTICE: CopyToOneRow: the number of valid values: 3 +NOTICE: CopyToOneRow: the number of valid values: 3 +NOTICE: CopyToOneRow: the number of valid values: 3 NOTICE: CopyToEnd +-- Reset data. +TRUNCATE copy_data; +INSERT INTO copy_data VALUES (1, 2, 3), (12, 34, 56), (123, 456, 789); DROP EXTENSION test_copy_format; -- Install custom COPY handlers to a schema that isn't included in -- search_path. diff --git a/src/test/modules/test_copy_format/sql/test_copy_format.sql b/src/test/modules/test_copy_format/sql/test_copy_format.sql index b262794f878..c7beb2fb8ae 100644 --- a/src/test/modules/test_copy_format/sql/test_copy_format.sql +++ b/src/test/modules/test_copy_format/sql/test_copy_format.sql @@ -4,10 +4,45 @@ INSERT INTO copy_data VALUES (1, 2, 3), (12, 34, 56), (123, 456, 789); -- No WITH SCHEMA. It installs custom COPY handlers to the current -- schema. CREATE EXTENSION test_copy_format; + -- We can find a custom COPY handler without schema. + +-- 987 is accepted. +-- 654 is a hard error because ON_ERROR is stop by default. COPY copy_data FROM stdin WITH (FORMAT 'test_copy_format'); +987 +654 \. + +-- 987 is accepted. +-- 654 is a soft error because ON_ERROR is ignore. +COPY copy_data FROM stdin WITH (FORMAT 'test_copy_format', ON_ERROR ignore); +987 +654 +\. + +-- 987 is accepted. +-- 654 is a soft error because ON_ERROR is ignore. +COPY copy_data FROM stdin WITH (FORMAT 'test_copy_format', ON_ERROR ignore, LOG_VERBOSITY verbose); +987 +654 +\. + +-- 987 is accepted. +-- 654 is a soft error because ON_ERROR is ignore. +-- 321 is a hard error. +COPY copy_data FROM stdin WITH (FORMAT 'test_copy_format', ON_ERROR ignore); +987 +654 +321 +\. + COPY copy_data TO stdout WITH (FORMAT 'test_copy_format'); + +-- Reset data. +TRUNCATE copy_data; +INSERT INTO copy_data VALUES (1, 2, 3), (12, 34, 56), (123, 456, 789); + DROP EXTENSION test_copy_format; diff --git a/src/test/modules/test_copy_format/test_copy_format.c b/src/test/modules/test_copy_format/test_copy_format.c index 1d754201336..34ec693a7ec 100644 --- a/src/test/modules/test_copy_format/test_copy_format.c +++ b/src/test/modules/test_copy_format/test_copy_format.c @@ -14,6 +14,7 @@ #include "postgres.h" #include "commands/copyapi.h" +#include "commands/copyfrom_internal.h" #include "commands/defrem.h" #include "utils/builtins.h" @@ -35,8 +36,85 @@ TestCopyFromStart(CopyFromState cstate, TupleDesc tupDesc) static bool TestCopyFromOneRow(CopyFromState cstate, ExprContext *econtext, Datum *values, bool *nulls) { + int n_attributes = list_length(cstate->attnumlist); + char *line; + int line_size = n_attributes + 1; /* +1 is for new line */ + int read_bytes; + ereport(NOTICE, (errmsg("CopyFromOneRow"))); - return false; + + cstate->cur_lineno++; + line = palloc(line_size); + read_bytes = CopyFromStateGetData(cstate, line, line_size, line_size); + if (read_bytes == 0) + return false; + if (read_bytes != line_size) + ereport(ERROR, + (errcode(ERRCODE_BAD_COPY_FILE_FORMAT), + errmsg("one line must be %d bytes: %d", + line_size, read_bytes))); + + if (cstate->cur_lineno == 1) + { + /* Success */ + TupleDesc tupDesc = RelationGetDescr(cstate->rel); + ListCell *cur; + int i = 0; + + foreach(cur, cstate->attnumlist) + { + int attnum = lfirst_int(cur); + int m = attnum - 1; + Form_pg_attribute att = TupleDescAttr(tupDesc, m); + + if (att->atttypid == INT2OID) + { + values[i] = Int16GetDatum(line[i] - '0'); + } + else if (att->atttypid == INT4OID) + { + values[i] = Int32GetDatum(line[i] - '0'); + } + else if (att->atttypid == INT8OID) + { + values[i] = Int64GetDatum(line[i] - '0'); + } + nulls[i] = false; + i++; + } + } + else if (cstate->cur_lineno == 2) + { + /* Soft error */ + TupleDesc tupDesc = RelationGetDescr(cstate->rel); + int attnum = lfirst_int(list_head(cstate->attnumlist)); + int m = attnum - 1; + Form_pg_attribute att = TupleDescAttr(tupDesc, m); + char value[2]; + + cstate->cur_attname = NameStr(att->attname); + value[0] = line[0]; + value[1] = '\0'; + cstate->cur_attval = value; + errsave((Node *) cstate->escontext, + ( + errcode(ERRCODE_INVALID_TEXT_REPRESENTATION), + errmsg("invalid value: \"%c\"", line[0]))); + CopyFromSkipErrorRow(cstate); + cstate->cur_attname = NULL; + cstate->cur_attval = NULL; + return true; + } + else + { + /* Hard error */ + ereport(ERROR, + (errcode(ERRCODE_BAD_COPY_FILE_FORMAT), + errmsg("too much lines: %llu", + (unsigned long long) cstate->cur_lineno))); + } + + return true; } static void -- 2.47.2