diff --git a/src/backend/commands/copy.c b/src/backend/commands/copy.c index cfad47b562..d3325b130d 100644 --- a/src/backend/commands/copy.c +++ b/src/backend/commands/copy.c @@ -23,6 +23,7 @@ #include "access/xact.h" #include "catalog/pg_authid.h" #include "commands/copy.h" +#include "commands/copyapi.h" #include "commands/defrem.h" #include "executor/executor.h" #include "mb/pg_wchar.h" @@ -32,6 +33,7 @@ #include "parser/parse_coerce.h" #include "parser/parse_collate.h" #include "parser/parse_expr.h" +#include "parser/parse_func.h" #include "parser/parse_relation.h" #include "rewrite/rewriteHandler.h" #include "utils/acl.h" @@ -427,6 +429,8 @@ ProcessCopyOptions(ParseState *pstate, opts_out->file_encoding = -1; + /* Text is the default format. */ + opts_out->to_ops = &CopyToRoutineText; /* Extract options from the statement node tree */ foreach(option, options) { @@ -442,9 +446,26 @@ ProcessCopyOptions(ParseState *pstate, if (strcmp(fmt, "text") == 0) /* default format */ ; else if (strcmp(fmt, "csv") == 0) + { opts_out->csv_mode = true; + opts_out->to_ops = &CopyToRoutineCSV; + } else if (strcmp(fmt, "binary") == 0) + { opts_out->binary = true; + opts_out->to_ops = &CopyToRoutineBinary; + } + else if (!is_from) + { + /* + * XXX: Currently we support custom COPY format only for COPY + * TO. + * + * XXX: need to check the combination of the existing options + * and a custom format (e.g., FREEZE)? + */ + opts_out->to_ops = GetCopyToRoutineFormat(fmt); + } else ereport(ERROR, (errcode(ERRCODE_INVALID_PARAMETER_VALUE), @@ -864,3 +885,83 @@ CopyGetAttnums(TupleDesc tupDesc, Relation rel, List *attnamelist) return attnums; } + +/* + * Get a CopyRoutine struct by invoking the handler function for the given + * format. + * + * Workhorse for GetCopyToRoutineFormat and GetCopyFromRoutineFormat. + */ +static CopyRoutine * +GetCopyRoutineFormat(char *format_name, bool is_from) +{ + Oid handlerOid; + Oid funcargtypes[1]; + CopyRoutine *routine; + Datum datum; + + funcargtypes[0] = INTERNALOID; + handlerOid = LookupFuncName(list_make1(makeString(format_name)), 1, + funcargtypes, true); + + if (!OidIsValid(handlerOid)) + ereport(ERROR, + (errcode(ERRCODE_UNDEFINED_OBJECT), + errmsg("COPY format \"%s\" not recognized", format_name))); + + datum = OidFunctionCall1(handlerOid, BoolGetDatum(is_from)); + + routine = (CopyRoutine *) DatumGetPointer(datum); + + /* Check returned struct */ + if (routine == NULL || !IsA(routine, CopyRoutine)) + elog(ERROR, "copy handler function %u did not return a CopyRoutine struct", + handlerOid); + + if (!IsA(routine->impl, CopyToRoutine) && !IsA(routine->impl, CopyFromRoutine)) + elog(ERROR, "copy handler function %u returned invalid CopyRoutine struct", + handlerOid); + + if (routine->is_from && !IsA(routine->impl, CopyFromRoutine)) + elog(ERROR, "copy handler function %u returned COPY TO routines but expected COPY FROM routines", + handlerOid); + + if (!routine->is_from && !IsA(routine->impl, CopyToRoutine)) + elog(ERROR, "copy handler function %u returned COPY FROM routines but expected COPY TO routines", + handlerOid); + +#ifdef USE_ASSERT_CHECKING + if (routine->is_from) + { + /* FIXME */ + } + else + { + CopyToRoutine *impl = castNode(CopyToRoutine, routine->impl); + + Assert(impl->start_fn != NULL); + Assert(impl->onerow_fn != NULL); + Assert(impl->end_fn != NULL); + } +#endif + + return routine; +} + +CopyToRoutine * +GetCopyToRoutineFormat(char *format_name) +{ + CopyRoutine *routine; + + routine = GetCopyRoutineFormat(format_name, false); + return (CopyToRoutine *) castNode(CopyToRoutine, routine->impl); +} + +CopyFromRoutine * +GetCopyFromRoutineFormat(char *format_name) +{ + CopyRoutine *routine; + + routine = GetCopyRoutineFormat(format_name, true); + return (CopyFromRoutine *) castNode(CopyFromRoutine, routine->impl); +} diff --git a/src/backend/commands/copyto.c b/src/backend/commands/copyto.c index c66a047c4a..46b8d9638b 100644 --- a/src/backend/commands/copyto.c +++ b/src/backend/commands/copyto.c @@ -131,6 +131,208 @@ static void CopySendEndOfRow(CopyToState cstate); static void CopySendInt32(CopyToState cstate, int32 val); static void CopySendInt16(CopyToState cstate, int16 val); +/* + * CopyToFormatOps implementations. + */ + +/* + * CopyToFormatOps implementation for "text" and "csv". CopyToFormatText*() + * refer cstate->opts.csv_mode and change their behavior. We can split this + * implementation and stop referring cstate->opts.csv_mode later. + */ + +static void +CopyToFormatTextSendEndOfRow(CopyToState cstate) +{ + switch (cstate->copy_dest) + { + case COPY_FILE: + /* Default line termination depends on platform */ +#ifndef WIN32 + CopySendChar(cstate, '\n'); +#else + CopySendString(cstate, "\r\n"); +#endif + break; + case COPY_FRONTEND: + /* The FE/BE protocol uses \n as newline for all platforms */ + CopySendChar(cstate, '\n'); + break; + default: + break; + } + CopySendEndOfRow(cstate); +} + +static void +CopyToFormatTextStart(CopyToState cstate, TupleDesc tupDesc) +{ + int num_phys_attrs; + ListCell *cur; + + num_phys_attrs = tupDesc->natts; + /* Get info about the columns we need to process. */ + cstate->out_functions = (FmgrInfo *) palloc(num_phys_attrs * sizeof(FmgrInfo)); + foreach(cur, cstate->attnumlist) + { + int attnum = lfirst_int(cur); + Oid out_func_oid; + bool isvarlena; + Form_pg_attribute attr = TupleDescAttr(tupDesc, attnum - 1); + + getTypeOutputInfo(attr->atttypid, &out_func_oid, &isvarlena); + fmgr_info(out_func_oid, &cstate->out_functions[attnum - 1]); + } + + /* + * For non-binary copy, we need to convert null_print to file encoding, + * because it will be sent directly with CopySendString. + */ + if (cstate->need_transcoding) + cstate->opts.null_print_client = pg_server_to_any(cstate->opts.null_print, + cstate->opts.null_print_len, + cstate->file_encoding); + + /* if a header has been requested send the line */ + if (cstate->opts.header_line) + { + bool hdr_delim = false; + + foreach(cur, cstate->attnumlist) + { + int attnum = lfirst_int(cur); + char *colname; + + if (hdr_delim) + CopySendChar(cstate, cstate->opts.delim[0]); + hdr_delim = true; + + colname = NameStr(TupleDescAttr(tupDesc, attnum - 1)->attname); + + if (cstate->opts.csv_mode) + CopyAttributeOutCSV(cstate, colname, false, + list_length(cstate->attnumlist) == 1); + else + CopyAttributeOutText(cstate, colname); + } + + CopyToFormatTextSendEndOfRow(cstate); + } +} + +static void +CopyToFormatTextOneRow(CopyToState cstate, TupleTableSlot *slot) +{ + bool need_delim = false; + FmgrInfo *out_functions = cstate->out_functions; + ListCell *cur; + + foreach(cur, cstate->attnumlist) + { + int attnum = lfirst_int(cur); + Datum value = slot->tts_values[attnum - 1]; + bool isnull = slot->tts_isnull[attnum - 1]; + + if (need_delim) + CopySendChar(cstate, cstate->opts.delim[0]); + need_delim = true; + + if (isnull) + CopySendString(cstate, cstate->opts.null_print_client); + else + { + char *string; + + string = OutputFunctionCall(&out_functions[attnum - 1], value); + if (cstate->opts.csv_mode) + CopyAttributeOutCSV(cstate, string, + cstate->opts.force_quote_flags[attnum - 1], + list_length(cstate->attnumlist) == 1); + else + CopyAttributeOutText(cstate, string); + } + } + + CopyToFormatTextSendEndOfRow(cstate); +} + +static void +CopyToFormatTextEnd(CopyToState cstate) +{ +} + +/* + * CopyToFormatOps implementation for "binary". + */ + +static void +CopyToFormatBinaryStart(CopyToState cstate, TupleDesc tupDesc) +{ + int num_phys_attrs; + ListCell *cur; + + num_phys_attrs = tupDesc->natts; + /* Get info about the columns we need to process. */ + cstate->out_functions = (FmgrInfo *) palloc(num_phys_attrs * sizeof(FmgrInfo)); + foreach(cur, cstate->attnumlist) + { + int attnum = lfirst_int(cur); + Oid out_func_oid; + bool isvarlena; + Form_pg_attribute attr = TupleDescAttr(tupDesc, attnum - 1); + + getTypeBinaryOutputInfo(attr->atttypid, &out_func_oid, &isvarlena); + fmgr_info(out_func_oid, &cstate->out_functions[attnum - 1]); + } + + /* Generate header for a binary copy */ + /* Signature */ + CopySendData(cstate, BinarySignature, 11); + /* Flags field */ + CopySendInt32(cstate, 0); + /* No header extension */ + CopySendInt32(cstate, 0); +} + +static void +CopyToFormatBinaryOneRow(CopyToState cstate, TupleTableSlot *slot) +{ + FmgrInfo *out_functions = cstate->out_functions; + ListCell *cur; + + /* Binary per-tuple header */ + CopySendInt16(cstate, list_length(cstate->attnumlist)); + + foreach(cur, cstate->attnumlist) + { + int attnum = lfirst_int(cur); + Datum value = slot->tts_values[attnum - 1]; + bool isnull = slot->tts_isnull[attnum - 1]; + + if (isnull) + CopySendInt32(cstate, -1); + else + { + bytea *outputbytes; + + outputbytes = SendFunctionCall(&out_functions[attnum - 1], value); + CopySendInt32(cstate, VARSIZE(outputbytes) - VARHDRSZ); + CopySendData(cstate, VARDATA(outputbytes), + VARSIZE(outputbytes) - VARHDRSZ); + } + } + + CopySendEndOfRow(cstate); +} + +static void +CopyToFormatBinaryEnd(CopyToState cstate) +{ + /* Generate trailer for a binary copy */ + CopySendInt16(cstate, -1); + /* Need to flush out the trailer */ + CopySendEndOfRow(cstate); +} /* * Send copy start/stop messages for frontend copies. These have changed @@ -198,16 +400,6 @@ CopySendEndOfRow(CopyToState cstate) switch (cstate->copy_dest) { case COPY_FILE: - if (!cstate->opts.binary) - { - /* Default line termination depends on platform */ -#ifndef WIN32 - CopySendChar(cstate, '\n'); -#else - CopySendString(cstate, "\r\n"); -#endif - } - if (fwrite(fe_msgbuf->data, fe_msgbuf->len, 1, cstate->copy_file) != 1 || ferror(cstate->copy_file)) @@ -242,10 +434,6 @@ CopySendEndOfRow(CopyToState cstate) } break; case COPY_FRONTEND: - /* The FE/BE protocol uses \n as newline for all platforms */ - if (!cstate->opts.binary) - CopySendChar(cstate, '\n'); - /* Dump the accumulated row as one CopyData message */ (void) pq_putmessage(PqMsg_CopyData, fe_msgbuf->data, fe_msgbuf->len); break; @@ -748,8 +936,6 @@ DoCopyTo(CopyToState cstate) bool pipe = (cstate->filename == NULL && cstate->data_dest_cb == NULL); bool fe_copy = (pipe && whereToSendOutput == DestRemote); TupleDesc tupDesc; - int num_phys_attrs; - ListCell *cur; uint64 processed; if (fe_copy) @@ -759,32 +945,11 @@ DoCopyTo(CopyToState cstate) tupDesc = RelationGetDescr(cstate->rel); else tupDesc = cstate->queryDesc->tupDesc; - num_phys_attrs = tupDesc->natts; cstate->opts.null_print_client = cstate->opts.null_print; /* default */ /* We use fe_msgbuf as a per-row buffer regardless of copy_dest */ cstate->fe_msgbuf = makeStringInfo(); - /* Get info about the columns we need to process. */ - cstate->out_functions = (FmgrInfo *) palloc(num_phys_attrs * sizeof(FmgrInfo)); - foreach(cur, cstate->attnumlist) - { - int attnum = lfirst_int(cur); - Oid out_func_oid; - bool isvarlena; - Form_pg_attribute attr = TupleDescAttr(tupDesc, attnum - 1); - - if (cstate->opts.binary) - getTypeBinaryOutputInfo(attr->atttypid, - &out_func_oid, - &isvarlena); - else - getTypeOutputInfo(attr->atttypid, - &out_func_oid, - &isvarlena); - fmgr_info(out_func_oid, &cstate->out_functions[attnum - 1]); - } - /* * Create a temporary memory context that we can reset once per row to * recover palloc'd memory. This avoids any problems with leaks inside @@ -795,57 +960,7 @@ DoCopyTo(CopyToState cstate) "COPY TO", ALLOCSET_DEFAULT_SIZES); - if (cstate->opts.binary) - { - /* Generate header for a binary copy */ - int32 tmp; - - /* Signature */ - CopySendData(cstate, BinarySignature, 11); - /* Flags field */ - tmp = 0; - CopySendInt32(cstate, tmp); - /* No header extension */ - tmp = 0; - CopySendInt32(cstate, tmp); - } - else - { - /* - * For non-binary copy, we need to convert null_print to file - * encoding, because it will be sent directly with CopySendString. - */ - if (cstate->need_transcoding) - cstate->opts.null_print_client = pg_server_to_any(cstate->opts.null_print, - cstate->opts.null_print_len, - cstate->file_encoding); - - /* if a header has been requested send the line */ - if (cstate->opts.header_line) - { - bool hdr_delim = false; - - foreach(cur, cstate->attnumlist) - { - int attnum = lfirst_int(cur); - char *colname; - - if (hdr_delim) - CopySendChar(cstate, cstate->opts.delim[0]); - hdr_delim = true; - - colname = NameStr(TupleDescAttr(tupDesc, attnum - 1)->attname); - - if (cstate->opts.csv_mode) - CopyAttributeOutCSV(cstate, colname, false, - list_length(cstate->attnumlist) == 1); - else - CopyAttributeOutText(cstate, colname); - } - - CopySendEndOfRow(cstate); - } - } + cstate->opts.to_ops->start_fn(cstate, tupDesc); if (cstate->rel) { @@ -884,13 +999,7 @@ DoCopyTo(CopyToState cstate) processed = ((DR_copy *) cstate->queryDesc->dest)->processed; } - if (cstate->opts.binary) - { - /* Generate trailer for a binary copy */ - CopySendInt16(cstate, -1); - /* Need to flush out the trailer */ - CopySendEndOfRow(cstate); - } + cstate->opts.to_ops->end_fn(cstate); MemoryContextDelete(cstate->rowcontext); @@ -906,71 +1015,15 @@ DoCopyTo(CopyToState cstate) static void CopyOneRowTo(CopyToState cstate, TupleTableSlot *slot) { - bool need_delim = false; - FmgrInfo *out_functions = cstate->out_functions; MemoryContext oldcontext; - ListCell *cur; - char *string; MemoryContextReset(cstate->rowcontext); oldcontext = MemoryContextSwitchTo(cstate->rowcontext); - if (cstate->opts.binary) - { - /* Binary per-tuple header */ - CopySendInt16(cstate, list_length(cstate->attnumlist)); - } - /* Make sure the tuple is fully deconstructed */ slot_getallattrs(slot); - foreach(cur, cstate->attnumlist) - { - int attnum = lfirst_int(cur); - Datum value = slot->tts_values[attnum - 1]; - bool isnull = slot->tts_isnull[attnum - 1]; - - if (!cstate->opts.binary) - { - if (need_delim) - CopySendChar(cstate, cstate->opts.delim[0]); - need_delim = true; - } - - if (isnull) - { - if (!cstate->opts.binary) - CopySendString(cstate, cstate->opts.null_print_client); - else - CopySendInt32(cstate, -1); - } - else - { - if (!cstate->opts.binary) - { - string = OutputFunctionCall(&out_functions[attnum - 1], - value); - if (cstate->opts.csv_mode) - CopyAttributeOutCSV(cstate, string, - cstate->opts.force_quote_flags[attnum - 1], - list_length(cstate->attnumlist) == 1); - else - CopyAttributeOutText(cstate, string); - } - else - { - bytea *outputbytes; - - outputbytes = SendFunctionCall(&out_functions[attnum - 1], - value); - CopySendInt32(cstate, VARSIZE(outputbytes) - VARHDRSZ); - CopySendData(cstate, VARDATA(outputbytes), - VARSIZE(outputbytes) - VARHDRSZ); - } - } - } - - CopySendEndOfRow(cstate); + cstate->opts.to_ops->onerow_fn(cstate, slot); MemoryContextSwitchTo(oldcontext); } @@ -1284,3 +1337,37 @@ CreateCopyDestReceiver(void) return (DestReceiver *) self; } + +/* + * COPY TO routines for built-in formats. + */ + +const CopyToRoutine CopyToRoutineText = +{ + .type = T_CopyToRoutine, + .start_fn = CopyToFormatTextStart, + .onerow_fn = CopyToFormatTextOneRow, + .end_fn = CopyToFormatTextEnd, +}; + +/* + * We can use the same CopyToFormatOps for both of "text" and "csv" because + * CopyToFormatText*() refer cstate->opts.csv_mode and change their + * behavior. We can split the implementations and stop referring + * cstate->opts.csv_mode later. + */ +const CopyToRoutine CopyToRoutineCSV = +{ + .type = T_CopyToRoutine, + .start_fn = CopyToFormatTextStart, + .onerow_fn = CopyToFormatTextOneRow, + .end_fn = CopyToFormatTextEnd, +}; + +const CopyToRoutine CopyToRoutineBinary = +{ + .type = T_CopyToRoutine, + .start_fn = CopyToFormatBinaryStart, + .onerow_fn = CopyToFormatBinaryOneRow, + .end_fn = CopyToFormatBinaryEnd, +}; diff --git a/src/backend/nodes/Makefile b/src/backend/nodes/Makefile index ebbe9052cb..04d41f5b0a 100644 --- a/src/backend/nodes/Makefile +++ b/src/backend/nodes/Makefile @@ -50,6 +50,7 @@ node_headers = \ access/sdir.h \ access/tableam.h \ access/tsmapi.h \ + commands/copyapi.h \ commands/event_trigger.h \ commands/trigger.h \ executor/tuptable.h \ diff --git a/src/backend/nodes/gen_node_support.pl b/src/backend/nodes/gen_node_support.pl index 72c7963578..c48015a612 100644 --- a/src/backend/nodes/gen_node_support.pl +++ b/src/backend/nodes/gen_node_support.pl @@ -61,6 +61,7 @@ my @all_input_files = qw( access/sdir.h access/tableam.h access/tsmapi.h + commands/copyapi.h commands/event_trigger.h commands/trigger.h executor/tuptable.h @@ -85,6 +86,7 @@ my @nodetag_only_files = qw( access/sdir.h access/tableam.h access/tsmapi.h + commands/copyapi.h commands/event_trigger.h commands/trigger.h executor/tuptable.h diff --git a/src/backend/utils/adt/pseudotypes.c b/src/backend/utils/adt/pseudotypes.c index 3ba8cb192c..4391e5cefc 100644 --- a/src/backend/utils/adt/pseudotypes.c +++ b/src/backend/utils/adt/pseudotypes.c @@ -373,6 +373,7 @@ PSEUDOTYPE_DUMMY_IO_FUNCS(fdw_handler); PSEUDOTYPE_DUMMY_IO_FUNCS(table_am_handler); PSEUDOTYPE_DUMMY_IO_FUNCS(index_am_handler); PSEUDOTYPE_DUMMY_IO_FUNCS(tsm_handler); +PSEUDOTYPE_DUMMY_IO_FUNCS(copy_handler); PSEUDOTYPE_DUMMY_IO_FUNCS(internal); PSEUDOTYPE_DUMMY_IO_FUNCS(anyelement); PSEUDOTYPE_DUMMY_IO_FUNCS(anynonarray); diff --git a/src/include/catalog/pg_proc.dat b/src/include/catalog/pg_proc.dat index 77e8b13764..9e0f33ad9e 100644 --- a/src/include/catalog/pg_proc.dat +++ b/src/include/catalog/pg_proc.dat @@ -7602,6 +7602,12 @@ { oid => '3312', descr => 'I/O', proname => 'tsm_handler_out', prorettype => 'cstring', proargtypes => 'tsm_handler', prosrc => 'tsm_handler_out' }, +{ oid => '8753', descr => 'I/O', + proname => 'copy_handler_in', proisstrict => 'f', prorettype => 'copy_handler', + proargtypes => 'cstring', prosrc => 'copy_handler_in' }, +{ oid => '8754', descr => 'I/O', + proname => 'copy_handler_out', prorettype => 'cstring', + proargtypes => 'copy_handler', prosrc => 'copy_handler_out' }, { oid => '267', descr => 'I/O', proname => 'table_am_handler_in', proisstrict => 'f', prorettype => 'table_am_handler', proargtypes => 'cstring', diff --git a/src/include/catalog/pg_type.dat b/src/include/catalog/pg_type.dat index f6110a850d..4fe5c17818 100644 --- a/src/include/catalog/pg_type.dat +++ b/src/include/catalog/pg_type.dat @@ -632,6 +632,12 @@ typcategory => 'P', typinput => 'tsm_handler_in', typoutput => 'tsm_handler_out', typreceive => '-', typsend => '-', typalign => 'i' }, +{ oid => '8752', + descr => 'pseudo-type for the result of a copy to/from method functoin', + typname => 'copy_handler', typlen => '4', typbyval => 't', typtype => 'p', + typcategory => 'P', typinput => 'copy_handler_in', + typoutput => 'copy_handler_out', typreceive => '-', typsend => '-', + typalign => 'i' }, { oid => '269', typname => 'table_am_handler', descr => 'pseudo-type for the result of a table AM handler function', diff --git a/src/include/commands/copy.h b/src/include/commands/copy.h index f2cca0b90b..2d5784db73 100644 --- a/src/include/commands/copy.h +++ b/src/include/commands/copy.h @@ -18,6 +18,7 @@ #include "nodes/parsenodes.h" #include "parser/parse_node.h" #include "tcop/dest.h" +#include "commands/copyapi.h" /* * Represents whether a header line should be present, and whether it must @@ -63,12 +64,9 @@ typedef struct CopyFormatOptions bool *force_null_flags; /* per-column CSV FN flags */ bool convert_selectively; /* do selective binary conversion? */ List *convert_select; /* list of column names (can be NIL) */ + const CopyToRoutine *to_ops; /* callback routines for COPY TO */ } CopyFormatOptions; -/* These are private in commands/copy[from|to].c */ -typedef struct CopyFromStateData *CopyFromState; -typedef struct CopyToStateData *CopyToState; - typedef int (*copy_data_source_cb) (void *outbuf, int minread, int maxread); typedef void (*copy_data_dest_cb) (void *data, int len); @@ -102,4 +100,9 @@ extern uint64 DoCopyTo(CopyToState cstate); extern List *CopyGetAttnums(TupleDesc tupDesc, Relation rel, List *attnamelist); +/* build-in COPY TO format routines */ +extern const CopyToRoutine CopyToRoutineText; +extern const CopyToRoutine CopyToRoutineCSV; +extern const CopyToRoutine CopyToRoutineBinary; + #endif /* COPY_H */ diff --git a/src/include/commands/copyapi.h b/src/include/commands/copyapi.h new file mode 100644 index 0000000000..49c64493b6 --- /dev/null +++ b/src/include/commands/copyapi.h @@ -0,0 +1,53 @@ +/*------------------------------------------------------------------------- + * + * copyapi.h + * API for COPY TO/FROM + * + * Copyright (c) 2015-2023, PostgreSQL Global Development Group + * + * src/include/command/copyapi.h + * + *------------------------------------------------------------------------- + */ +#ifndef COPYAPI_H +#define COPYAPI_H + +#include "executor/tuptable.h" + +typedef struct CopyFromStateData *CopyFromState; +typedef struct CopyToStateData *CopyToState; + +typedef void (*CopyToStart_function) (CopyToState cstate, TupleDesc tupDesc); +typedef void (*CopyToOneRow_function) (CopyToState cstate, TupleTableSlot *slot); +typedef void (*CopyToEnd_function) (CopyToState cstate); + +typedef struct CopyRoutine +{ + NodeTag type; + + bool is_from; + + /* Either CopyToRoutine or CopyFromRoutine */ + Node *impl; +} CopyRoutine; + +typedef struct CopyToRoutine +{ + NodeTag type; + + CopyToStart_function start_fn; + CopyToOneRow_function onerow_fn; + CopyToEnd_function end_fn; +} CopyToRoutine; + +typedef struct CopyFromRoutine +{ + NodeTag type; + + /* XXX: define routines for COPY FROM */ +} CopyFromRoutine; + +extern CopyToRoutine * GetCopyToRoutineFormat(char *format_name); +extern CopyFromRoutine * GetCopyFromRoutineFormat(char *format_name); + +#endif /* COPYAPI_H */ diff --git a/src/include/nodes/meson.build b/src/include/nodes/meson.build index 626dc696d5..53b262568c 100644 --- a/src/include/nodes/meson.build +++ b/src/include/nodes/meson.build @@ -11,6 +11,7 @@ node_support_input_i = [ 'access/sdir.h', 'access/tableam.h', 'access/tsmapi.h', + 'commands/copyapi.h', 'commands/event_trigger.h', 'commands/trigger.h', 'executor/tuptable.h', diff --git a/src/test/modules/Makefile b/src/test/modules/Makefile index 5d33fa6a9a..204cfd3f49 100644 --- a/src/test/modules/Makefile +++ b/src/test/modules/Makefile @@ -15,6 +15,7 @@ SUBDIRS = \ spgist_name_ops \ test_bloomfilter \ test_copy_callbacks \ + test_copy_format \ test_custom_rmgrs \ test_ddl_deparse \ test_dsa \ diff --git a/src/test/modules/meson.build b/src/test/modules/meson.build index b76f588559..2fbe1abd4a 100644 --- a/src/test/modules/meson.build +++ b/src/test/modules/meson.build @@ -12,6 +12,7 @@ subdir('spgist_name_ops') subdir('ssl_passphrase_callback') subdir('test_bloomfilter') subdir('test_copy_callbacks') +subdir('test_copy_format') subdir('test_custom_rmgrs') subdir('test_ddl_deparse') subdir('test_dsa') diff --git a/src/test/modules/test_copy_format/.gitignore b/src/test/modules/test_copy_format/.gitignore new file mode 100644 index 0000000000..5dcb3ff972 --- /dev/null +++ b/src/test/modules/test_copy_format/.gitignore @@ -0,0 +1,4 @@ +# Generated subdirectories +/log/ +/results/ +/tmp_check/ diff --git a/src/test/modules/test_copy_format/Makefile b/src/test/modules/test_copy_format/Makefile new file mode 100644 index 0000000000..f2b89b56a1 --- /dev/null +++ b/src/test/modules/test_copy_format/Makefile @@ -0,0 +1,23 @@ +# src/test/modules/test_copy_format/Makefile + +MODULE_big = test_copy_format +OBJS = \ + $(WIN32RES) \ + test_copy_format.o +PGFILEDESC = "test_copy_format - test custom COPY format" + +EXTENSION = test_copy_format +DATA = test_copy_format--1.0.sql + +REGRESS = test_copy_format + +ifdef USE_PGXS +PG_CONFIG = pg_config +PGXS := $(shell $(PG_CONFIG) --pgxs) +include $(PGXS) +else +subdir = src/test/modules/test_copy_format +top_builddir = ../../../.. +include $(top_builddir)/src/Makefile.global +include $(top_srcdir)/contrib/contrib-global.mk +endif 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 new file mode 100644 index 0000000000..8becdb6369 --- /dev/null +++ b/src/test/modules/test_copy_format/expected/test_copy_format.out @@ -0,0 +1,14 @@ +CREATE EXTENSION test_copy_format; +CREATE TABLE public.test (a INT, b INT, c INT); +INSERT INTO public.test VALUES (1, 2, 3), (12, 34, 56), (123, 456, 789); +COPY public.test FROM stdin WITH (format 'testfmt'); +ERROR: COPY format "testfmt" not recognized +LINE 1: COPY public.test FROM stdin WITH (format 'testfmt'); + ^ +COPY public.test TO stdout WITH (format 'testfmt'); +NOTICE: testfmt_handler called with is_from 0 +NOTICE: testfmt_copyto_start called +NOTICE: testfmt_copyto_onerow called +NOTICE: testfmt_copyto_onerow called +NOTICE: testfmt_copyto_onerow called +NOTICE: testfmt_copyto_end called diff --git a/src/test/modules/test_copy_format/meson.build b/src/test/modules/test_copy_format/meson.build new file mode 100644 index 0000000000..4adf048280 --- /dev/null +++ b/src/test/modules/test_copy_format/meson.build @@ -0,0 +1,33 @@ +# Copyright (c) 2022-2023, PostgreSQL Global Development Group + +test_copy_format_sources = files( + 'test_copy_format.c', +) + +if host_system == 'windows' + test_copy_format_sources += rc_lib_gen.process(win32ver_rc, extra_args: [ + '--NAME', 'test_copy_format', + '--FILEDESC', 'test_copy_format - test COPY format routines',]) +endif + +test_copy_format = shared_module('test_copy_format', + test_copy_format_sources, + kwargs: pg_test_mod_args, +) +test_install_libs += test_copy_format + +test_install_data += files( + 'test_copy_format.control', + 'test_copy_format--1.0.sql', +) + +tests += { + 'name': 'test_copy_format', + 'sd': meson.current_source_dir(), + 'bd': meson.current_build_dir(), + 'regress': { + 'sql': [ + 'test_copy_format', + ], + }, +} 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 new file mode 100644 index 0000000000..1052135252 --- /dev/null +++ b/src/test/modules/test_copy_format/sql/test_copy_format.sql @@ -0,0 +1,5 @@ +CREATE EXTENSION test_copy_format; +CREATE TABLE public.test (a INT, b INT, c INT); +INSERT INTO public.test VALUES (1, 2, 3), (12, 34, 56), (123, 456, 789); +COPY public.test FROM stdin WITH (format 'testfmt'); +COPY public.test TO stdout WITH (format 'testfmt'); diff --git a/src/test/modules/test_copy_format/test_copy_format--1.0.sql b/src/test/modules/test_copy_format/test_copy_format--1.0.sql new file mode 100644 index 0000000000..2749924831 --- /dev/null +++ b/src/test/modules/test_copy_format/test_copy_format--1.0.sql @@ -0,0 +1,9 @@ +/* src/test/modules/test_copy_format/test_copy_format--1.0.sql */ + +-- complain if script is sourced in psql, rather than via CREATE EXTENSION +\echo Use "CREATE EXTENSION test_copy_format" to load this file. \quit + + +CREATE FUNCTION testfmt(internal) + RETURNS copy_handler + AS 'MODULE_PATHNAME', 'copy_testfmt_handler' LANGUAGE C; diff --git a/src/test/modules/test_copy_format/test_copy_format.c b/src/test/modules/test_copy_format/test_copy_format.c new file mode 100644 index 0000000000..26ecddb76f --- /dev/null +++ b/src/test/modules/test_copy_format/test_copy_format.c @@ -0,0 +1,69 @@ +/*-------------------------------------------------------------------------- + * + * test_copy_format.c + * Code for custom COPY format. + * + * Portions Copyright (c) 1996-2023, PostgreSQL Global Development Group + * Portions Copyright (c) 1994, Regents of the University of California + * + * IDENTIFICATION + * src/test/modules/test_copy_format/test_copy_format.c + * + * ------------------------------------------------------------------------- + */ + +#include "postgres.h" + +#include "commands/copyapi.h" +#include "fmgr.h" + +PG_MODULE_MAGIC; + +static void +testfmt_copyto_start(CopyToState cstate, TupleDesc tupDesc) +{ + ereport(NOTICE, + (errmsg("testfmt_copyto_start called"))); +} + +static void +testfmt_copyto_onerow(CopyToState cstate, TupleTableSlot *slot) +{ + ereport(NOTICE, + (errmsg("testfmt_copyto_onerow called"))); +} + +static void +testfmt_copyto_end(CopyToState cstate) +{ + ereport(NOTICE, + (errmsg("testfmt_copyto_end called"))); +} + +PG_FUNCTION_INFO_V1(copy_testfmt_handler); +Datum +copy_testfmt_handler(PG_FUNCTION_ARGS) +{ + bool is_from = PG_GETARG_BOOL(0); + CopyRoutine *routine = makeNode(CopyRoutine); + + ereport(NOTICE, + (errmsg("testfmt_handler called with is_from %d", is_from))); + + routine->is_from = is_from; + + if (!is_from) + { + CopyToRoutine *impl = makeNode(CopyToRoutine); + + impl->start_fn = testfmt_copyto_start; + impl->onerow_fn = testfmt_copyto_onerow; + impl->end_fn = testfmt_copyto_end; + + routine->impl = (Node *) impl; + } + else + elog(ERROR, "custom COPY format \"testfmt\" does not support COPY FROM"); + + PG_RETURN_POINTER(routine); +} diff --git a/src/test/modules/test_copy_format/test_copy_format.control b/src/test/modules/test_copy_format/test_copy_format.control new file mode 100644 index 0000000000..57e0ef9d91 --- /dev/null +++ b/src/test/modules/test_copy_format/test_copy_format.control @@ -0,0 +1,4 @@ +comment = 'Test code for COPY format' +default_version = '1.0' +module_pathname = '$libdir/test_copy_format' +relocatable = true