From 53cdd9defcee3c6f0f2e2743d885cec3d16f6b9f Mon Sep 17 00:00:00 2001 From: Sutou Kouhei Date: Mon, 25 Nov 2024 14:11:55 +0900 Subject: [PATCH v35 4/7] Add support for adding custom COPY FROM format This uses the same handler for COPY TO and COPY FROM but uses different routine. This uses CopyToRoutine for COPY TO and CopyFromRoutine for COPY FROM. PostgreSQL calls a COPY TO/FROM handler with "is_from" argument. It's true for COPY FROM and false for COPY TO: copy_handler(true) returns CopyToRoutine copy_handler(false) returns CopyFromRoutine This also add a test module for custom COPY FROM handler. --- src/backend/commands/copy.c | 13 +++---- src/backend/commands/copyfrom.c | 36 +++++++++++++---- src/include/catalog/pg_type.dat | 2 +- src/include/commands/copyapi.h | 2 + .../expected/test_copy_format.out | 10 +++-- .../test_copy_format/sql/test_copy_format.sql | 1 + .../test_copy_format/test_copy_format.c | 39 ++++++++++++++++++- 7 files changed, 82 insertions(+), 21 deletions(-) diff --git a/src/backend/commands/copy.c b/src/backend/commands/copy.c index 8d94bc313eb..b4417bb6819 100644 --- a/src/backend/commands/copy.c +++ b/src/backend/commands/copy.c @@ -483,8 +483,8 @@ defGetCopyLogVerbosityChoice(DefElem *def, ParseState *pstate) * This function checks whether the option value is a built-in format such as * "text" and "csv" or not. If the option value isn't a built-in format, this * function finds a COPY format handler that returns a CopyToRoutine (for - * is_from == false). If no COPY format handler is found, this function - * reports an error. + * is_from == false) or CopyFromRountine (for is_from == true). If no COPY + * format handler is found, this function reports an error. */ static void ProcessCopyOptionFormat(ParseState *pstate, @@ -518,12 +518,9 @@ ProcessCopyOptionFormat(ParseState *pstate, } /* custom format */ - if (!is_from) - { - funcargtypes[0] = INTERNALOID; - handlerOid = LookupFuncName(list_make1(makeString(format)), 1, - funcargtypes, true); - } + funcargtypes[0] = INTERNALOID; + handlerOid = LookupFuncName(list_make1(makeString(format)), 1, + funcargtypes, true); if (!OidIsValid(handlerOid)) ereport(ERROR, (errcode(ERRCODE_INVALID_PARAMETER_VALUE), diff --git a/src/backend/commands/copyfrom.c b/src/backend/commands/copyfrom.c index 198cee2bc48..114ea969dfa 100644 --- a/src/backend/commands/copyfrom.c +++ b/src/backend/commands/copyfrom.c @@ -129,6 +129,7 @@ static void CopyFromBinaryEnd(CopyFromState cstate); /* text format */ static const CopyFromRoutine CopyFromRoutineText = { + .type = T_CopyFromRoutine, .CopyFromInFunc = CopyFromTextLikeInFunc, .CopyFromStart = CopyFromTextLikeStart, .CopyFromOneRow = CopyFromTextOneRow, @@ -137,6 +138,7 @@ static const CopyFromRoutine CopyFromRoutineText = { /* CSV format */ static const CopyFromRoutine CopyFromRoutineCSV = { + .type = T_CopyFromRoutine, .CopyFromInFunc = CopyFromTextLikeInFunc, .CopyFromStart = CopyFromTextLikeStart, .CopyFromOneRow = CopyFromCSVOneRow, @@ -145,6 +147,7 @@ static const CopyFromRoutine CopyFromRoutineCSV = { /* binary format */ static const CopyFromRoutine CopyFromRoutineBinary = { + .type = T_CopyFromRoutine, .CopyFromInFunc = CopyFromBinaryInFunc, .CopyFromStart = CopyFromBinaryStart, .CopyFromOneRow = CopyFromBinaryOneRow, @@ -153,15 +156,32 @@ static const CopyFromRoutine CopyFromRoutineBinary = { /* Return a COPY FROM routine for the given options */ static const CopyFromRoutine * -CopyFromGetRoutine(CopyFormatOptions opts) +CopyFromGetRoutine(CopyFormatOptions *opts) { - if (opts.csv_mode) - return &CopyFromRoutineCSV; - else if (opts.binary) - return &CopyFromRoutineBinary; + if (OidIsValid(opts->handler)) + { + Datum datum; + Node *routine; - /* default is text */ - return &CopyFromRoutineText; + datum = OidFunctionCall1(opts->handler, BoolGetDatum(true)); + routine = (Node *) DatumGetPointer(datum); + if (routine == NULL || !IsA(routine, CopyFromRoutine)) + ereport( + ERROR, + (errcode( + ERRCODE_INVALID_PARAMETER_VALUE), + errmsg("COPY handler function " + "%u did not return " + "CopyFromRoutine struct", + opts->handler))); + return castNode(CopyFromRoutine, routine); + } + else if (opts->csv_mode) + return &CopyFromRoutineCSV; + else if (opts->binary) + return &CopyFromRoutineBinary; + else + return &CopyFromRoutineText; } /* Implementation of the start callback for text and CSV formats */ @@ -1574,7 +1594,7 @@ BeginCopyFrom(ParseState *pstate, ProcessCopyOptions(pstate, &cstate->opts, true /* is_from */ , options); /* Set the format routine */ - cstate->routine = CopyFromGetRoutine(cstate->opts); + cstate->routine = CopyFromGetRoutine(&cstate->opts); /* Process the target relation */ cstate->rel = rel; diff --git a/src/include/catalog/pg_type.dat b/src/include/catalog/pg_type.dat index 340e0cd0a8d..63b7d65f982 100644 --- a/src/include/catalog/pg_type.dat +++ b/src/include/catalog/pg_type.dat @@ -634,7 +634,7 @@ typoutput => 'tsm_handler_out', typreceive => '-', typsend => '-', typalign => 'i' }, { oid => '8752', - descr => 'pseudo-type for the result of a copy to method function', + descr => 'pseudo-type for the result of a copy to/from method function', typname => 'copy_handler', typlen => '4', typbyval => 't', typtype => 'p', typcategory => 'P', typinput => 'copy_handler_in', typoutput => 'copy_handler_out', typreceive => '-', typsend => '-', diff --git a/src/include/commands/copyapi.h b/src/include/commands/copyapi.h index 5c5ea6592e3..895c105d8d8 100644 --- a/src/include/commands/copyapi.h +++ b/src/include/commands/copyapi.h @@ -64,6 +64,8 @@ extern void CopyToStateFlush(CopyToState cstate); */ typedef struct CopyFromRoutine { + NodeTag type; + /* * Set input function information. This callback is called once at the * beginning of COPY FROM. 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 adfe7d1572a..016893e7026 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 @@ -2,9 +2,13 @@ CREATE EXTENSION test_copy_format; CREATE TABLE public.test (a smallint, b integer, c bigint); INSERT INTO public.test VALUES (1, 2, 3), (12, 34, 56), (123, 456, 789); COPY public.test FROM stdin WITH (FORMAT 'test_copy_format'); -ERROR: COPY format "test_copy_format" not recognized -LINE 1: COPY public.test FROM stdin WITH (FORMAT 'test_copy_format')... - ^ +NOTICE: test_copy_format: is_from=true +NOTICE: CopyFromInFunc: atttypid=21 +NOTICE: CopyFromInFunc: atttypid=23 +NOTICE: CopyFromInFunc: atttypid=20 +NOTICE: CopyFromStart: natts=3 +NOTICE: CopyFromOneRow +NOTICE: CopyFromEnd COPY public.test TO stdout WITH (FORMAT 'test_copy_format'); NOTICE: test_copy_format: is_from=false NOTICE: CopyToOutFunc: atttypid=21 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 810b3d8cedc..0dfdfa00080 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 @@ -2,4 +2,5 @@ CREATE EXTENSION test_copy_format; CREATE TABLE public.test (a smallint, b integer, c bigint); INSERT INTO public.test VALUES (1, 2, 3), (12, 34, 56), (123, 456, 789); COPY public.test FROM stdin WITH (FORMAT 'test_copy_format'); +\. COPY public.test TO stdout WITH (FORMAT '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 b42d472d851..abafc668463 100644 --- a/src/test/modules/test_copy_format/test_copy_format.c +++ b/src/test/modules/test_copy_format/test_copy_format.c @@ -18,6 +18,40 @@ PG_MODULE_MAGIC; +static void +CopyFromInFunc(CopyFromState cstate, Oid atttypid, + FmgrInfo *finfo, Oid *typioparam) +{ + ereport(NOTICE, (errmsg("CopyFromInFunc: atttypid=%d", atttypid))); +} + +static void +CopyFromStart(CopyFromState cstate, TupleDesc tupDesc) +{ + ereport(NOTICE, (errmsg("CopyFromStart: natts=%d", tupDesc->natts))); +} + +static bool +CopyFromOneRow(CopyFromState cstate, ExprContext *econtext, Datum *values, bool *nulls) +{ + ereport(NOTICE, (errmsg("CopyFromOneRow"))); + return false; +} + +static void +CopyFromEnd(CopyFromState cstate) +{ + ereport(NOTICE, (errmsg("CopyFromEnd"))); +} + +static const CopyFromRoutine CopyFromRoutineTestCopyFormat = { + .type = T_CopyFromRoutine, + .CopyFromInFunc = CopyFromInFunc, + .CopyFromStart = CopyFromStart, + .CopyFromOneRow = CopyFromOneRow, + .CopyFromEnd = CopyFromEnd, +}; + static void CopyToOutFunc(CopyToState cstate, Oid atttypid, FmgrInfo *finfo) { @@ -59,5 +93,8 @@ test_copy_format(PG_FUNCTION_ARGS) ereport(NOTICE, (errmsg("test_copy_format: is_from=%s", is_from ? "true" : "false"))); - PG_RETURN_POINTER(&CopyToRoutineTestCopyFormat); + if (is_from) + PG_RETURN_POINTER(&CopyFromRoutineTestCopyFormat); + else + PG_RETURN_POINTER(&CopyToRoutineTestCopyFormat); } -- 2.47.2