From 5de0629e3664b981b2a246c480a90636ddfc8dd7 Mon Sep 17 00:00:00 2001
From: Christoph Berg <myon@debian.org>
Date: Fri, 8 Sep 2023 15:59:29 +0200
Subject: [PATCH] Add JSON output format to psql

Query results are formatted as an array of JSON objects. In non-expanded
mode, one object per line is printed, and in expanded mode, one
key-value pair. Use `\pset format json` to enable, or run `psql --json`
or `psql -J`. NULL values are printed as `null`, independently from the
psql null setting, all other values are printed as quoted strings.
---
 doc/src/sgml/ref/psql-ref.sgml     |  26 +++++-
 src/bin/psql/command.c             |   6 +-
 src/bin/psql/help.c                |   1 +
 src/bin/psql/startup.c             |   6 +-
 src/bin/psql/tab-complete.c        |   2 +-
 src/fe_utils/print.c               | 128 +++++++++++++++++++++++++++--
 src/include/fe_utils/print.h       |   1 +
 src/test/regress/expected/psql.out |  67 +++++++++++++++
 src/test/regress/sql/psql.sql      |  25 ++++++
 9 files changed, 253 insertions(+), 9 deletions(-)

diff --git a/doc/src/sgml/ref/psql-ref.sgml b/doc/src/sgml/ref/psql-ref.sgml
index cc7d797159..f1f7eda082 100644
--- a/doc/src/sgml/ref/psql-ref.sgml
+++ b/doc/src/sgml/ref/psql-ref.sgml
@@ -274,6 +274,17 @@ EOF
       </listitem>
     </varlistentry>
 
+    <varlistentry id="app-psql-option-json">
+      <term><option>-J</option></term>
+      <term><option>--json</option></term>
+      <listitem>
+      <para>
+      Switches to <acronym>JSON</acronym> output mode.  This is
+      equivalent to <command>\pset format json</command>.
+      </para>
+      </listitem>
+    </varlistentry>
+
     <varlistentry id="app-psql-option-list">
       <term><option>-l</option></term>
       <term><option>--list</option></term>
@@ -2956,6 +2967,7 @@ lo_import 152801
           <literal>asciidoc</literal>,
           <literal>csv</literal>,
           <literal>html</literal>,
+          <literal>json</literal>,
           <literal>latex</literal>,
           <literal>latex-longtable</literal>, <literal>troff-ms</literal>,
           <literal>unaligned</literal>, or <literal>wrapped</literal>.
@@ -2972,7 +2984,7 @@ lo_import 152801
           in by other programs, for example, tab-separated or comma-separated
           format.  However, the field separator character is not treated
           specially if it appears in a column's value;
-          so <acronym>CSV</acronym> format may be better suited for such
+          the <acronym>CSV</acronym> or <acronym>JSON</acronym> formats may be better suited for such
           purposes.
           </para>
 
@@ -2997,6 +3009,18 @@ lo_import 152801
           <command>\pset csv_fieldsep</command>.
           </para>
 
+          <para><literal>json</literal> format
+          <indexterm>
+           <primary>JSON format</primary>
+           <secondary>in psql</secondary>
+          </indexterm>
+          writes output in JSON format, described in
+          <ulink url="https://www.ietf.org/rfc/rfc4627.txt">RFC 4627</ulink>.
+          The result is an array of JSON objects. In non-expanded mode, one
+          object per line is printed, while in expanded mode, each key-value
+          pair is a separate line.
+          </para>
+
           <para><literal>wrapped</literal> format is like <literal>aligned</literal> but wraps
           wide data values across lines to make the output fit in the target
           column width.  The target width is determined as described under
diff --git a/src/bin/psql/command.c b/src/bin/psql/command.c
index 82cc091568..d8d11e9519 100644
--- a/src/bin/psql/command.c
+++ b/src/bin/psql/command.c
@@ -4318,6 +4318,9 @@ _align2string(enum printFormat in)
 		case PRINT_HTML:
 			return "html";
 			break;
+		case PRINT_JSON:
+			return "json";
+			break;
 		case PRINT_LATEX:
 			return "latex";
 			break;
@@ -4408,6 +4411,7 @@ do_pset(const char *param, const char *value, printQueryOpt *popt, bool quiet)
 			{"asciidoc", PRINT_ASCIIDOC},
 			{"csv", PRINT_CSV},
 			{"html", PRINT_HTML},
+			{"json", PRINT_JSON},
 			{"latex", PRINT_LATEX},
 			{"troff-ms", PRINT_TROFF_MS},
 			{"unaligned", PRINT_UNALIGNED},
@@ -4448,7 +4452,7 @@ do_pset(const char *param, const char *value, printQueryOpt *popt, bool quiet)
 			}
 			else
 			{
-				pg_log_error("\\pset: allowed formats are aligned, asciidoc, csv, html, latex, latex-longtable, troff-ms, unaligned, wrapped");
+				pg_log_error("\\pset: allowed formats are aligned, asciidoc, csv, html, json, latex, latex-longtable, troff-ms, unaligned, wrapped");
 				return false;
 			}
 		}
diff --git a/src/bin/psql/help.c b/src/bin/psql/help.c
index 3b2d59e2ee..1f5f081c98 100644
--- a/src/bin/psql/help.c
+++ b/src/bin/psql/help.c
@@ -99,6 +99,7 @@ usage(unsigned short int pager)
 		  "                           field separator for unaligned output (default: \"%s\")\n",
 		  DEFAULT_FIELD_SEP);
 	HELP0("  -H, --html               HTML table output mode\n");
+	HELP0("  -J, --json               JSON output mode\n");
 	HELP0("  -P, --pset=VAR[=ARG]     set printing option VAR to ARG (see \\pset command)\n");
 	HELP0("  -R, --record-separator=STRING\n"
 		  "                           record separator for unaligned output (default: newline)\n");
diff --git a/src/bin/psql/startup.c b/src/bin/psql/startup.c
index 78526eb9da..b1eed9d3fe 100644
--- a/src/bin/psql/startup.c
+++ b/src/bin/psql/startup.c
@@ -496,6 +496,7 @@ parse_psql_options(int argc, char *argv[], struct adhoc_opts *options)
 		{"field-separator-zero", no_argument, NULL, 'z'},
 		{"host", required_argument, NULL, 'h'},
 		{"html", no_argument, NULL, 'H'},
+		{"json", no_argument, NULL, 'J'},
 		{"list", no_argument, NULL, 'l'},
 		{"log-file", required_argument, NULL, 'L'},
 		{"no-readline", no_argument, NULL, 'n'},
@@ -528,7 +529,7 @@ parse_psql_options(int argc, char *argv[], struct adhoc_opts *options)
 
 	memset(options, 0, sizeof *options);
 
-	while ((c = getopt_long(argc, argv, "aAbc:d:eEf:F:h:HlL:no:p:P:qR:sStT:U:v:VwWxXz?01",
+	while ((c = getopt_long(argc, argv, "aAbc:d:eEf:F:h:HJlL:no:p:P:qR:sStT:U:v:VwWxXz?01",
 							long_options, &optindex)) != -1)
 	{
 		switch (c)
@@ -576,6 +577,9 @@ parse_psql_options(int argc, char *argv[], struct adhoc_opts *options)
 			case 'H':
 				pset.popt.topt.format = PRINT_HTML;
 				break;
+			case 'J':
+				pset.popt.topt.format = PRINT_JSON;
+				break;
 			case 'l':
 				options->list_dbs = true;
 				break;
diff --git a/src/bin/psql/tab-complete.c b/src/bin/psql/tab-complete.c
index 049801186c..cf4b84b31e 100644
--- a/src/bin/psql/tab-complete.c
+++ b/src/bin/psql/tab-complete.c
@@ -4902,7 +4902,7 @@ psql_completion(const char *text, int start, int end)
 	else if (TailMatchesCS("\\pset", MatchAny))
 	{
 		if (TailMatchesCS("format"))
-			COMPLETE_WITH_CS("aligned", "asciidoc", "csv", "html", "latex",
+			COMPLETE_WITH_CS("aligned", "asciidoc", "csv", "html", "json", "latex",
 							 "latex-longtable", "troff-ms", "unaligned",
 							 "wrapped");
 		else if (TailMatchesCS("xheader_width"))
diff --git a/src/fe_utils/print.c b/src/fe_utils/print.c
index a91292ab6d..8beb8a3188 100644
--- a/src/fe_utils/print.c
+++ b/src/fe_utils/print.c
@@ -2159,6 +2159,112 @@ print_html_vertical(const printTableContent *cont, FILE *fout)
 }
 
 
+/**********************/
+/* JSON				  */
+/**********************/
+
+
+static void
+json_escaped_print(const char *in, FILE *fout)
+{
+	const char *p;
+
+	if (!in)
+	{
+		fputs("null", fout);
+		return;
+	}
+
+	fputc('"', fout);
+
+	for (p = in; *p; p++)
+	{
+		switch (*p)
+		{
+			case '\\':
+				fputs("\\\\", fout);
+				break;
+			case '"':
+				fputs("\\\"", fout);
+				break;
+			case '\b':
+				fputs("\\b", fout);
+				break;
+			case '\f':
+				fputs("\\f", fout);
+				break;
+			case '\n':
+				fputs("\\n", fout);
+				break;
+			case '\r':
+				fputs("\\r", fout);
+				break;
+			case '\t':
+				fputs("\\t", fout);
+				break;
+			default:
+				if (*p >= 0 && *p < 0x20)
+					fprintf(fout, "\\u%04X", *p);
+				else
+					fputc(*p, fout);
+		}
+	}
+
+	fputc('"', fout);
+}
+
+static void
+print_json_text(const printTableContent *cont, FILE *fout, bool vertical)
+{
+	unsigned int i;
+	const char *const *ptr;
+
+	if (cancel_pressed)
+		return;
+
+	/*
+	 * The title and footer are never printed in json format.
+	 */
+
+	if (cont->opt->start_table)
+		fputc('[', fout);
+	else if (cont->nrows > 0)
+		/* print trailing comma after previous FETCH_COUNT iteration */
+		fputc(',', fout);
+
+	/* print records */
+	for (i = 0, ptr = cont->cells; i < cont->ncolumns * cont->nrows; i++, ptr++)
+	{
+		/* start new record */
+		if (i % cont->ncolumns == 0)
+		{
+			if (cancel_pressed)
+				break;
+			/* print trailing comma after previous record */
+			if (i > 0)
+				fputc(',', fout);
+			fputs(vertical ? "{\n" : "\n{", fout);
+		}
+
+		/* cell */
+		fputs(vertical ? "  " : " ", fout);
+		json_escaped_print(cont->headers[i % cont->ncolumns], fout);
+		fputs(": ", fout);
+		json_escaped_print(*ptr, fout);
+
+		/* next cell */
+		if (i % cont->ncolumns != cont->ncolumns - 1)
+			fputs(vertical ? ",\n" : ",", fout);
+		/* end of record */
+		else
+			fputs(vertical ? "\n}" : " }", fout);
+	}
+
+	if (cont->opt->stop_table)
+		fputs(vertical ? "]\n" : "\n]\n", fout);
+}
+
+
 /*************************/
 /* ASCIIDOC				 */
 /*************************/
@@ -3274,13 +3380,19 @@ printTableAddCell(printTableContent *const content, char *cell,
 		exit(EXIT_FAILURE);
 	}
 
-	*content->cell = (char *) mbvalidate((unsigned char *) cell,
-										 content->opt->encoding);
+	/* cell can be NULL in JSON format */
+	if (cell)
+	{
+		*content->cell = (char *) mbvalidate((unsigned char *) cell,
+											 content->opt->encoding);
 
 #ifdef ENABLE_NLS
-	if (translate)
-		*content->cell = _(*content->cell);
+		if (translate)
+			*content->cell = _(*content->cell);
 #endif
+	}
+	else
+		*content->cell = cell;
 
 	if (mustfree)
 	{
@@ -3502,6 +3614,9 @@ printTable(const printTableContent *cont,
 			else
 				print_html_text(cont, fout);
 			break;
+		case PRINT_JSON:
+			print_json_text(cont, fout, cont->opt->expanded == 1);
+			break;
 		case PRINT_ASCIIDOC:
 			if (cont->opt->expanded == 1)
 				print_asciidoc_vertical(cont, fout);
@@ -3581,7 +3696,10 @@ printQuery(const PGresult *result, const printQueryOpt *opt,
 			bool		translate;
 
 			if (PQgetisnull(result, r, c))
-				cell = opt->nullPrint ? opt->nullPrint : "";
+				if (cont.opt->format == PRINT_JSON)
+					cell = NULL;
+				else
+					cell = opt->nullPrint ? opt->nullPrint : "";
 			else
 			{
 				cell = PQgetvalue(result, r, c);
diff --git a/src/include/fe_utils/print.h b/src/include/fe_utils/print.h
index cfc26b9520..d0d0fe1224 100644
--- a/src/include/fe_utils/print.h
+++ b/src/include/fe_utils/print.h
@@ -32,6 +32,7 @@ enum printFormat
 	PRINT_ASCIIDOC,
 	PRINT_CSV,
 	PRINT_HTML,
+	PRINT_JSON,
 	PRINT_LATEX,
 	PRINT_LATEX_LONGTABLE,
 	PRINT_TROFF_MS,
diff --git a/src/test/regress/expected/psql.out b/src/test/regress/expected/psql.out
index 13e4f6db7b..0ca0ccb67f 100644
--- a/src/test/regress/expected/psql.out
+++ b/src/test/regress/expected/psql.out
@@ -3562,6 +3562,73 @@ execute q;
 
 \pset tableattr
 deallocate q;
+\pset expanded off
+-- test json output format
+\pset format json
+prepare q as
+  select 'some"text' as "a\title", E'  <foo>\n<bar>' as "junk",
+         '' as "empty", n as int, null as "null"
+  from generate_series(1,2) as n;
+execute q;
+[
+{ "a\\title": "some\"text", "junk": "  <foo>\n<bar>", "empty": "", "int": "1", "null": null },
+{ "a\\title": "some\"text", "junk": "  <foo>\n<bar>", "empty": "", "int": "2", "null": null }
+]
+execute q \gx
+[{
+  "a\\title": "some\"text",
+  "junk": "  <foo>\n<bar>",
+  "empty": "",
+  "int": "1",
+  "null": null
+},{
+  "a\\title": "some\"text",
+  "junk": "  <foo>\n<bar>",
+  "empty": "",
+  "int": "2",
+  "null": null
+}]
+deallocate q;
+-- check if commas are placed correctly with FETCH_COUNT
+\set FETCH_COUNT 2
+select i from generate_series(1, 3) g(i);
+[
+{ "i": "1" },
+{ "i": "2" },
+{ "i": "3" }
+]
+select i from generate_series(1, 3) g(i) \gx
+[{
+  "i": "1"
+},{
+  "i": "2"
+},{
+  "i": "3"
+}]
+select i from generate_series(1, 4) g(i);
+[
+{ "i": "1" },
+{ "i": "2" },
+{ "i": "3" },
+{ "i": "4" }
+]
+select i from generate_series(1, 4) g(i) \gx
+[{
+  "i": "1"
+},{
+  "i": "2"
+},{
+  "i": "3"
+},{
+  "i": "4"
+}]
+\unset FETCH_COUNT
+-- empty output
+select;
+[
+]
+select \gx
+[]
 -- test latex output format
 \pset format latex
 \pset border 1
diff --git a/src/test/regress/sql/psql.sql b/src/test/regress/sql/psql.sql
index 695c72d866..c8e1bc4ed3 100644
--- a/src/test/regress/sql/psql.sql
+++ b/src/test/regress/sql/psql.sql
@@ -674,6 +674,31 @@ execute q;
 \pset tableattr
 
 deallocate q;
+\pset expanded off
+
+-- test json output format
+
+\pset format json
+
+prepare q as
+  select 'some"text' as "a\title", E'  <foo>\n<bar>' as "junk",
+         '' as "empty", n as int, null as "null"
+  from generate_series(1,2) as n;
+execute q;
+execute q \gx
+deallocate q;
+
+-- check if commas are placed correctly with FETCH_COUNT
+\set FETCH_COUNT 2
+select i from generate_series(1, 3) g(i);
+select i from generate_series(1, 3) g(i) \gx
+select i from generate_series(1, 4) g(i);
+select i from generate_series(1, 4) g(i) \gx
+\unset FETCH_COUNT
+
+-- empty output
+select;
+select \gx
 
 -- test latex output format
 
-- 
2.43.0

