From 2b575846fee609237256fe8eaf38fd45005fe8da 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. Numeric values are printed in plain, booleans as
true/false, and 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               | 182 ++++++++++++++++++++++++++++-
 src/include/fe_utils/print.h       |   6 +-
 src/test/regress/expected/psql.out |  79 +++++++++++++
 src/test/regress/sql/psql.sql      |  28 +++++
 9 files changed, 325 insertions(+), 11 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 5c906e4806..824900819b 100644
--- a/src/bin/psql/command.c
+++ b/src/bin/psql/command.c
@@ -4299,6 +4299,9 @@ _align2string(enum printFormat in)
 		case PRINT_HTML:
 			return "html";
 			break;
+		case PRINT_JSON:
+			return "json";
+			break;
 		case PRINT_LATEX:
 			return "latex";
 			break;
@@ -4389,6 +4392,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},
@@ -4429,7 +4433,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 4e79a819d8..3950ffc2c6 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 036caaec2f..8480304e04 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 ada711d02f..f4ca4d0013 100644
--- a/src/bin/psql/tab-complete.c
+++ b/src/bin/psql/tab-complete.c
@@ -4915,7 +4915,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 d6c4d78f98..028ab8e5c1 100644
--- a/src/fe_utils/print.c
+++ b/src/fe_utils/print.c
@@ -2159,6 +2159,128 @@ print_html_vertical(const printTableContent *cont, FILE *fout)
 }
 
 
+/**********************/
+/* JSON				  */
+/**********************/
+
+/*
+ * Print a JSON value. If in is NULL, it's printed as null. Otherwise, quote =
+ * 'f' passes the value through unchanged, 'b' prints a boolean, and 't'
+ * applies JSON string quoting.
+ */
+static void
+json_escaped_print(const char *in, const char quote, FILE *fout)
+{
+	const char *p;
+
+	if (!in)
+	{
+		fputs("null", fout);
+		return;
+	}
+
+	if (quote == 'f')
+	{
+		fputs(in, fout);
+		return;
+	}
+
+	if (quote == 'b')
+	{
+		fputs(*in == 't' ? "true" : "false", 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], 't', fout);
+		fputs(": ", fout);
+		json_escaped_print(*ptr, cont->aligns[i % cont->ncolumns], 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 +3396,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 +3630,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);
@@ -3568,7 +3699,10 @@ printQuery(const PGresult *result, const printQueryOpt *opt,
 	{
 		printTableAddHeader(&cont, PQfname(result, i),
 							opt->translate_header,
-							column_type_alignment(PQftype(result, i)));
+							cont.opt->format == PRINT_JSON ?
+								column_type_jsonquote(PQftype(result, i)) :
+								column_type_alignment(PQftype(result, i))
+							);
 	}
 
 	/* set cells */
@@ -3581,7 +3715,12 @@ printQuery(const PGresult *result, const printQueryOpt *opt,
 			bool		translate;
 
 			if (PQgetisnull(result, r, c))
-				cell = opt->nullPrint ? opt->nullPrint : "";
+				/* for JSON output, represent NULL as NULL */
+				if (cont.opt->format == PRINT_JSON)
+					cell = NULL;
+				/* otherwise the configured null string */
+				else
+					cell = opt->nullPrint ? opt->nullPrint : "";
 			else
 			{
 				cell = PQgetvalue(result, r, c);
@@ -3637,6 +3776,37 @@ column_type_alignment(Oid ftype)
 	return align;
 }
 
+/*
+ * Whether a type needs quoting when represented as JSON. Numerical types are
+ * passed through, as are JSON and JSONB. Booleans need conversion from t/f to
+ * true/false.
+ */
+char
+column_type_jsonquote(Oid ftype)
+{
+	char		quote;
+
+	switch (ftype)
+	{
+		case BOOLOID:
+			quote = 'b';
+			break;
+		case INT2OID:
+		case INT4OID:
+		case INT8OID:
+		case FLOAT4OID:
+		case FLOAT8OID:
+		case NUMERICOID:
+		case OIDOID:
+			quote = 'f';
+			break;
+		default:
+			quote = 't';
+			break;
+	}
+	return quote;
+}
+
 void
 setDecimalLocale(void)
 {
diff --git a/src/include/fe_utils/print.h b/src/include/fe_utils/print.h
index 72824c5c2f..54ae802897 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,
@@ -176,7 +177,9 @@ typedef struct printTableContent
 	printTableFooter *footers;	/* Pointer to the first footer */
 	printTableFooter *footer;	/* Pointer to the last added footer */
 	char	   *aligns;			/* Array of alignment specifiers; 'l' or 'r',
-								 * one per column */
+								 * one per column. JSON output uses 't' for
+								 * to-be-quoted values, 'b' for booleans, and
+								 * 'f' otherwise. */
 	char	   *align;			/* Pointer to the last added alignment */
 } printTableContent;
 
@@ -228,6 +231,7 @@ extern void printQuery(const PGresult *result, const printQueryOpt *opt,
 					   FILE *fout, bool is_pager, FILE *flog);
 
 extern char column_type_alignment(Oid);
+extern char column_type_jsonquote(Oid);
 
 extern void setDecimalLocale(void);
 extern const printTextFormat *get_line_style(const printTableOpt *opt);
diff --git a/src/test/regress/expected/psql.out b/src/test/regress/expected/psql.out
index ad02772562..e5fc625b48 100644
--- a/src/test/regress/expected/psql.out
+++ b/src/test/regress/expected/psql.out
@@ -3562,6 +3562,85 @@ 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
+}]
+select null null, true true, false false;
+[
+{ "null": null, "true": true, "false": false }
+]
+select 1 int, 1e30::float float, 1.0 numeric;
+[
+{ "int": 1, "float": 1e+30, "numeric": 1.0 }
+]
+select 'text' text, '{"a": null}'::json json, '{"a": "b"}'::jsonb jsonb, 'A\000B'::bytea bytea;
+[
+{ "text": "text", "json": "{\"a\": null}", "jsonb": "{\"a\": \"b\"}", "bytea": "\\x410042" }
+]
+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 129f853353..595d9ef815 100644
--- a/src/test/regress/sql/psql.sql
+++ b/src/test/regress/sql/psql.sql
@@ -674,6 +674,34 @@ 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
+select null null, true true, false false;
+select 1 int, 1e30::float float, 1.0 numeric;
+select 'text' text, '{"a": null}'::json json, '{"a": "b"}'::jsonb jsonb, 'A\000B'::bytea bytea;
+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

