Extensible copy

contrib/Makefile

@@ -32,6 +32,7 @@ SUBDIRS = \
 		pageinspect	\
 		passwordcheck	\
 		pg_buffercache	\
+		pg_copy_json	\
 		pg_freespacemap \
 		pg_prewarm	\
 		pg_stat_statements \

contrib/file_fdw/file_fdw.c

@@ -329,7 +329,7 @@ file_fdw_validator(PG_FUNCTION_ARGS)
 	/*
 	 * Now apply the core COPY code's validation logic for more checks.
 	 */
-	ProcessCopyOptions(NULL, NULL, true, other_options);
+	ProcessCopyOptions(NULL, NULL, true, NULL, other_options);
 
 	/*
 	 * Either filename or program option is required for file_fdw foreign

contrib/meson.build

@@ -45,6 +45,7 @@ subdir('oid2name')
 subdir('pageinspect')
 subdir('passwordcheck')
 subdir('pg_buffercache')
+subdir('pg_copy_json')
 subdir('pgcrypto')
 subdir('pg_freespacemap')
 subdir('pg_prewarm')

contrib/pg_copy_json/.gitignore

@@ -0,0 +1,4 @@
+# Generated subdirectories
+/log/
+/results/
+/tmp_check/

contrib/pg_copy_json/Makefile

@@ -0,0 +1,23 @@
+# contrib/pg_copy_json//Makefile
+
+MODULE_big = pg_copy_json
+OBJS = \
+	$(WIN32RES) \
+	pg_copy_json.o
+PGFILEDESC = "pg_copy_json - COPY TO JSON (JavaScript Object Notation) format"
+
+EXTENSION = pg_copy_json
+DATA = pg_copy_json--1.0.sql
+
+REGRESS = test_copy_format
+
+ifdef USE_PGXS
+PG_CONFIG = pg_config
+PGXS := $(shell $(PG_CONFIG) --pgxs)
+include $(PGXS)
+else
+subdir = contrib/pg_copy_json
+top_builddir = ../..
+include $(top_builddir)/src/Makefile.global
+include $(top_srcdir)/contrib/contrib-global.mk
+endif

contrib/pg_copy_json/expected/pg_copy_json.out

@@ -0,0 +1,80 @@
+--
+-- COPY TO JSON
+--
+CREATE EXTENSION pg_copy_json;
+-- test copying in JSON format with various styles
+-- of embedded line ending characters
+create temp table copytest (
+	style	text,
+	test 	text,
+	filler	int);
+insert into copytest values('DOS',E'abc\r\ndef',1);
+insert into copytest values('Unix',E'abc\ndef',2);
+insert into copytest values('Mac',E'abc\rdef',3);
+insert into copytest values(E'esc\\ape',E'a\\r\\\r\\\n\\nb',4);
+copy copytest to stdout with (format 'json');
+{"style":"DOS","test":"abc\r\ndef","filler":1}
+{"style":"Unix","test":"abc\ndef","filler":2}
+{"style":"Mac","test":"abc\rdef","filler":3}
+{"style":"esc\\ape","test":"a\\r\\\r\\\n\\nb","filler":4}
+-- pg_copy_json do not support COPY FROM
+copy copytest from stdout with (format 'json');
+ERROR:  cannot use JSON mode in COPY FROM
+-- test copying in JSON format with various styles
+-- of embedded escaped characters
+create temp table copyjsontest (
+    id bigserial,
+    f1 text,
+    f2 timestamptz);
+insert into copyjsontest
+  select g.i,
+         CASE WHEN g.i % 2 = 0 THEN
+           'line with '' in it: ' || g.i::text
+         ELSE
+           'line with " in it: ' || g.i::text
+         END,
+         'Mon Feb 10 17:32:01 1997 PST'
+  from generate_series(1,5) as g(i);
+insert into copyjsontest (f1) values
+(E'aaa\"bbb'::text),
+(E'aaa\\bbb'::text),
+(E'aaa\/bbb'::text),
+(E'aaa\bbbb'::text),
+(E'aaa\fbbb'::text),
+(E'aaa\nbbb'::text),
+(E'aaa\rbbb'::text),
+(E'aaa\tbbb'::text);
+copy copyjsontest to stdout with (format 'json');
+{"id":1,"f1":"line with \" in it: 1","f2":"1997-02-10T17:32:01-08:00"}
+{"id":2,"f1":"line with ' in it: 2","f2":"1997-02-10T17:32:01-08:00"}
+{"id":3,"f1":"line with \" in it: 3","f2":"1997-02-10T17:32:01-08:00"}
+{"id":4,"f1":"line with ' in it: 4","f2":"1997-02-10T17:32:01-08:00"}
+{"id":5,"f1":"line with \" in it: 5","f2":"1997-02-10T17:32:01-08:00"}
+{"id":1,"f1":"aaa\"bbb","f2":null}
+{"id":2,"f1":"aaa\\bbb","f2":null}
+{"id":3,"f1":"aaa/bbb","f2":null}
+{"id":4,"f1":"aaa\bbbb","f2":null}
+{"id":5,"f1":"aaa\fbbb","f2":null}
+{"id":6,"f1":"aaa\nbbb","f2":null}
+{"id":7,"f1":"aaa\rbbb","f2":null}
+{"id":8,"f1":"aaa\tbbb","f2":null}
+-- test force array
+copy copytest to stdout (format 'json', force_array);
+[
+ {"style":"DOS","test":"abc\r\ndef","filler":1}
+,{"style":"Unix","test":"abc\ndef","filler":2}
+,{"style":"Mac","test":"abc\rdef","filler":3}
+,{"style":"esc\\ape","test":"a\\r\\\r\\\n\\nb","filler":4}
+]
+copy copytest to stdout (format 'json', force_array true);
+[
+ {"style":"DOS","test":"abc\r\ndef","filler":1}
+,{"style":"Unix","test":"abc\ndef","filler":2}
+,{"style":"Mac","test":"abc\rdef","filler":3}
+,{"style":"esc\\ape","test":"a\\r\\\r\\\n\\nb","filler":4}
+]
+copy copytest to stdout (format 'json', force_array false);
+{"style":"DOS","test":"abc\r\ndef","filler":1}
+{"style":"Unix","test":"abc\ndef","filler":2}
+{"style":"Mac","test":"abc\rdef","filler":3}
+{"style":"esc\\ape","test":"a\\r\\\r\\\n\\nb","filler":4}

contrib/pg_copy_json/meson.build

@@ -0,0 +1,34 @@
+# Copyright (c) 2024, PostgreSQL Global Development Group
+
+pg_copy_json_sources = files(
+  'pg_copy_json.c',
+)
+
+if host_system == 'windows'
+  pg_copy_json_sources += rc_lib_gen.process(win32ver_rc, extra_args: [
+    '--NAME', 'pg_copy_json',
+    '--FILEDESC', 'pg_copy_json - COPY TO JSON format',])
+endif
+
+pg_copy_json = shared_module('pg_copy_json',
+  pg_copy_json_sources,
+  kwargs: contrib_mod_args,
+)
+contrib_targets += pg_copy_json
+
+install_data(
+  'pg_copy_json--1.0.sql',
+  'pg_copy_json.control',
+  kwargs: contrib_data_args,
+)
+
+tests += {
+  'name': 'pg_copy_json',
+  'sd': meson.current_source_dir(),
+  'bd': meson.current_build_dir(),
+  'regress': {
+    'sql': [
+      'pg_copy_json',
+    ],
+  },
+}

contrib/pg_copy_json/pg_copy_json--1.0.sql

@@ -0,0 +1,9 @@
+/* contrib/pg_copy_json/copy_json--1.0.sql */
+
+-- complain if script is sourced in psql, rather than via CREATE EXTENSION
+\echo Use "CREATE EXTENSION pg_copy_json" to load this file. \quit
+
+CREATE FUNCTION pg_catalog.json(internal)
+	RETURNS copy_handler
+	AS 'MODULE_PATHNAME', 'copy_json'
+	LANGUAGE C;

contrib/pg_copy_json/pg_copy_json.c

@@ -0,0 +1,218 @@
+/*--------------------------------------------------------------------------
+ *
+ * pg_copy_json.c
+ *		COPY TO JSON (JavaScript Object Notation) format.
+ *
+ * Portions Copyright (c) 2024, PostgreSQL Global Development Group
+ *
+ * IDENTIFICATION
+ *		contrib/test_copy_format.c
+ *
+ * -------------------------------------------------------------------------
+ */
+
+#include "postgres.h"
+
+#include "commands/copy.h"
+#include "commands/defrem.h"
+#include "funcapi.h"
+#include "libpq/libpq.h"
+#include "libpq/pqformat.h"
+#include "utils/json.h"
+
+PG_MODULE_MAGIC;
+
+typedef struct
+{
+	/*
+	 * Force output of square brackets as array decorations at the beginning
+	 * and end of output, with commas between the rows.
+	 */
+	bool	force_array;
+	bool	force_array_specified;
+
+	/* need delimiter to start next json array element */
+	bool	json_row_delim_needed;
+} CopyJsonData;
+
+static inline void
+InitCopyJsonData(CopyJsonData *p)
+{
+	Assert(p);
+	p->force_array = false;
+	p->force_array_specified = false;
+	p->json_row_delim_needed = false;
+}
+
+static void
+CopyToJsonSendEndOfRow(CopyToState cstate)
+{
+	switch (cstate->copy_dest)
+	{
+		case COPY_DEST_FILE:
+			/* Default line termination depends on platform */
+#ifndef WIN32
+			CopySendChar(cstate, '\n');
+#else
+			CopySendString(cstate, "\r\n");
+#endif
+			break;
+		case COPY_DEST_FRONTEND:
+			/* The FE/BE protocol uses \n as newline for all platforms */
+			CopySendChar(cstate, '\n');
+			break;
+		default:
+			break;
+	}
+	CopyToStateFlush(cstate);
+}
+
+static bool
+CopyToJsonProcessOption(CopyToState cstate, DefElem *defel)
+{
+	CopyJsonData	   *p;
+
+	if (cstate->opaque == NULL)
+	{
+		MemoryContext oldcontext;
+		oldcontext = MemoryContextSwitchTo(cstate->copycontext);
+		cstate->opaque = palloc0(sizeof(CopyJsonData));
+		MemoryContextSwitchTo(oldcontext);
+		InitCopyJsonData(cstate->opaque);
+	}
+
+	p = (CopyJsonData *)cstate->opaque;
+
+	if (strcmp(defel->defname, "force_array") == 0)
+	{
+		if (p->force_array_specified)
+			ereport(ERROR,
+					errcode(ERRCODE_SYNTAX_ERROR),
+					errmsg("CopyToJsonProcessOption: redundant options \"%s\"=\"%s\"",
+						   defel->defname, defGetString(defel)));
+		p->force_array_specified = true;
+		p->force_array = defGetBoolean(defel);
+
+		return true;
+	}
+
+	return false;
+}
+
+static void
+CopyToJsonSendCopyBegin(CopyToState cstate)
+{
+	StringInfoData buf;
+	int16		format = 0;
+
+	pq_beginmessage(&buf, PqMsg_CopyOutResponse);
+	pq_sendbyte(&buf, format);	/* overall format */
+	/*
+	 * JSON mode is always one non-binary column
+	 */
+	pq_sendint16(&buf, 1);
+	pq_sendint16(&buf, 0);
+	pq_endmessage(&buf);
+}
+
+static void
+CopyToJsonStart(CopyToState cstate, TupleDesc tupDesc)
+{
+	CopyJsonData	   *p;
+
+	if (cstate->opaque == NULL)
+	{
+		MemoryContext oldcontext;
+		oldcontext = MemoryContextSwitchTo(cstate->copycontext);
+		cstate->opaque = palloc0(sizeof(CopyJsonData));
+		MemoryContextSwitchTo(oldcontext);
+		InitCopyJsonData(cstate->opaque);
+	}
+
+	/* No need to alloc cstate->out_functions */
+
+	p = (CopyJsonData *)cstate->opaque;
+
+	/* If FORCE_ARRAY has been specified send the open bracket. */
+	if (p->force_array)
+	{
+		CopySendChar(cstate, '[');
+		CopyToJsonSendEndOfRow(cstate);
+	}
+}
+
+static void
+CopyToJsonOneRow(CopyToState cstate, TupleTableSlot *slot)
+{
+	Datum				rowdata;
+	StringInfo			result;
+	CopyJsonData	   *p;
+
+	Assert(cstate->opaque);
+	p = (CopyJsonData *)cstate->opaque;
+
+	if(!cstate->rel)
+	{
+		for (int i = 0; i < slot->tts_tupleDescriptor->natts; i++)
+		{
+			/* Flat-copy the attribute array */
+			memcpy(TupleDescAttr(slot->tts_tupleDescriptor, i),
+			TupleDescAttr(cstate->queryDesc->tupDesc, i),
+							1 * sizeof(FormData_pg_attribute));
+		}
+		BlessTupleDesc(slot->tts_tupleDescriptor);
+	}
+	rowdata = ExecFetchSlotHeapTupleDatum(slot);
+	result = makeStringInfo();
+	composite_to_json(rowdata, result, false);
+
+	if (p->json_row_delim_needed)
+		CopySendChar(cstate, ',');
+	else if (p->force_array)
+	{
+		/* first row needs no delimiter */
+		CopySendChar(cstate, ' ');
+		p->json_row_delim_needed = true;
+	}
+	CopySendData(cstate, result->data, result->len);
+	CopyToJsonSendEndOfRow(cstate);
+}
+
+static void
+CopyToJsonEnd(CopyToState cstate)
+{
+	CopyJsonData	   *p;
+
+	Assert(cstate->opaque);
+	p = (CopyJsonData *)cstate->opaque;
+
+	/* If FORCE_ARRAY has been specified send the close bracket. */
+	if (p->force_array)
+	{
+		CopySendChar(cstate, ']');
+		CopyToJsonSendEndOfRow(cstate);
+	}
+}
+
+static const CopyToRoutine CopyToRoutineJson = {
+	.type = T_CopyToRoutine,
+	.CopyToProcessOption = CopyToJsonProcessOption,
+	.CopyToSendCopyBegin = CopyToJsonSendCopyBegin,
+	.CopyToStart = CopyToJsonStart,
+	.CopyToOneRow = CopyToJsonOneRow,
+	.CopyToEnd = CopyToJsonEnd,
+};
+
+PG_FUNCTION_INFO_V1(copy_json);
+Datum
+copy_json(PG_FUNCTION_ARGS)
+{
+	bool		is_from = PG_GETARG_BOOL(0);
+
+	if (is_from)
+		ereport(ERROR,
+				(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+				 errmsg("cannot use JSON mode in COPY FROM")));
+
+	PG_RETURN_POINTER(&CopyToRoutineJson);
+}