gset.patch

application/octet-stream

Filename: gset.patch
Type: application/octet-stream
Part: 0
Message: Re: proposal - assign result of query to psql variable

Patch

Same data as JSON: GET /api/v1/attachments/:id/patch the parsed metadata as JSON — format, series position, per-file stats; never the diff bytes. API reference →
Format: unified
File+
src/bin/psql/command.c 23 0
src/bin/psql/command.h 1 0
src/bin/psql/common.c 153 1
src/bin/psql/common.h 11 0
src/bin/psql/help.c 2 1
src/bin/psql/mainloop.c 8 0
src/bin/psql/psqlscan.h 7 1
src/bin/psql/psqlscan.l 99 3
src/bin/psql/settings.h 2 0
src/bin/psql/tab-complete.c 1 1
diff --git a/src/bin/psql/command.c b/src/bin/psql/command.c
index 8abadb2..2e8309d 100644
--- a/src/bin/psql/command.c
+++ b/src/bin/psql/command.c
@@ -748,6 +748,29 @@ exec_command(const char *cmd,
 		status = PSQL_CMD_SEND;
 	}
 
+	/* \gset send query and store result */
+	else if (strcmp(cmd, "gset") == 0)
+	{
+		bool	error;
+
+		pset.gvars = psql_scan_slash_vars(scan_state, &error);
+
+		if (!pset.gvars)
+		{
+			psql_error("\\%s: missing required argument\n", cmd);
+			status = PSQL_CMD_NOSEND;
+		}
+		else if (error)
+		{
+			psql_error("\\%s: syntax error\n", cmd);
+			status = PSQL_CMD_NOSEND;
+			tglist_free(pset.gvars);
+			pset.gvars = NULL;
+		}
+		else
+			status = PSQL_CMD_SEND;
+	}
+
 	/* help */
 	else if (strcmp(cmd, "h") == 0 || strcmp(cmd, "help") == 0)
 	{
diff --git a/src/bin/psql/command.h b/src/bin/psql/command.h
index f0bcea0..57cbe23 100644
--- a/src/bin/psql/command.h
+++ b/src/bin/psql/command.h
@@ -16,6 +16,7 @@ typedef enum _backslashResult
 {
 	PSQL_CMD_UNKNOWN = 0,		/* not done parsing yet (internal only) */
 	PSQL_CMD_SEND,				/* query complete; send off */
+	PSQL_CMD_NOSEND,			/* query complete, don't send */
 	PSQL_CMD_SKIP_LINE,			/* keep building query */
 	PSQL_CMD_TERMINATE,			/* quit program */
 	PSQL_CMD_NEWEDIT,			/* query buffer was changed (e.g., via \e) */
diff --git a/src/bin/psql/common.c b/src/bin/psql/common.c
index 330d5ce..5248721 100644
--- a/src/bin/psql/common.c
+++ b/src/bin/psql/common.c
@@ -826,6 +826,103 @@ PrintQueryResults(PGresult *results)
 	return success;
 }
 
+/*
+ * StoreQueryResult: store first row of result to selected variables
+ *
+ * Note: Utility function for use by SendQuery() only.
+ *
+ * Returns true if the query executed sucessfully, false otherwise.
+ */
+static bool
+StoreQueryResult(PGresult *result)
+{
+	bool success;
+
+	switch (PQresultStatus(result))
+	{
+		case PGRES_TUPLES_OK:
+			{
+				int	i;
+
+				if (PQntuples(result) < 1)
+				{
+					psql_error("no data found\n");
+					success = false;
+				}
+				else if (PQntuples(result) > 1)
+				{
+					psql_error("too many rows\n");
+					success = false;
+				}
+				else
+				{
+					TargetListData *iter = (TargetListData *) pset.gvars;
+
+					success = true;
+
+					for (i = 0; i < PQnfields(result); i++)
+					{
+						if (!iter)
+						{
+							psql_error("to few target variables\n");
+							success = false;
+							break;
+						}
+
+						if (iter->name)
+						{
+							if (!SetVariable(pset.vars, iter->name,
+										   PQgetvalue(result, 0, i)))
+							{
+								psql_error("invalid variable name: \"%s\"",
+													iter->name);
+								success = false;
+								break;
+							}
+						}
+
+						iter = iter->next;
+					}
+
+					if (success && iter != NULL)
+					{
+						psql_error("too many target variables\n");
+						success = false;
+					}
+				}
+			}
+			break;
+
+		case PGRES_COMMAND_OK:
+		case PGRES_EMPTY_QUERY:
+			psql_error("no data found\n");
+			success = false;
+			break;
+
+		case PGRES_COPY_OUT:
+		case PGRES_COPY_IN:
+			psql_error("COPY isnot supported by \\gset command\n");
+			success = false;
+			break;
+
+		case PGRES_BAD_RESPONSE:
+		case PGRES_NONFATAL_ERROR:
+		case PGRES_FATAL_ERROR:
+			success = false;
+			break;
+
+		default:
+			success = false;
+			psql_error("unexpected PQresultStatus: %d\n",
+					   PQresultStatus(result));
+			break;
+	}
+
+	tglist_free(pset.gvars);
+	pset.gvars = NULL;
+
+	return success;
+}
 
 /*
  * SendQuery: send the query string to the backend
@@ -953,7 +1050,12 @@ SendQuery(const char *query)
 
 		/* but printing results isn't: */
 		if (OK && results)
-			OK = PrintQueryResults(results);
+		{
+			if (pset.gvars)
+				OK = StoreQueryResult(results);
+			else
+				OK = PrintQueryResults(results);
+		}
 	}
 	else
 	{
@@ -1668,3 +1770,53 @@ expand_tilde(char **filename)
 
 	return *filename;
 }
+
+
+/*
+ * Add name of internal variable to query targer list
+ *
+ */
+TargetList
+tglist_add(TargetList tglist, const char *name)
+{
+	TargetListData		*tgf;
+
+	tgf = pg_malloc(sizeof(TargetListData));
+	tgf->name = name ? pg_strdup(name) : NULL;
+	tgf->next = NULL;
+
+	if (tglist)
+	{
+		TargetListData *iter = (TargetListData *) tglist;
+
+		while (iter->next)
+			iter = iter->next;
+
+		iter->next = tgf;
+
+		return tglist;
+	}
+	else
+		return (TargetList) tgf;
+}
+
+/*
+ * Release target list
+ *
+ */
+void
+tglist_free(TargetList tglist)
+{
+	TargetListData *iter = (TargetListData *) tglist;
+
+	while (iter)
+	{
+		TargetListData *next = iter->next;
+
+		if (iter->name)
+			free(iter->name);
+
+		free(iter);
+		iter = next;
+	}
+}
diff --git a/src/bin/psql/common.h b/src/bin/psql/common.h
index 8037cbc..f955fec 100644
--- a/src/bin/psql/common.h
+++ b/src/bin/psql/common.h
@@ -21,6 +21,14 @@
 
 #define atooid(x)  ((Oid) strtoul((x), NULL, 10))
 
+typedef struct _target_field
+{
+	char	   *name;
+	struct _target_field *next;
+} TargetListData;
+
+typedef struct TargetListData *TargetList;
+
 /*
  * Safer versions of some standard C library functions. If an
  * out-of-memory condition occurs, these functions will bail out
@@ -63,4 +71,7 @@ extern const char *session_username(void);
 
 extern char *expand_tilde(char **filename);
 
+extern TargetList  tglist_add(TargetList tglist, const char *name);
+extern void tglist_free(TargetList tglist);
+
 #endif   /* COMMON_H */
diff --git a/src/bin/psql/help.c b/src/bin/psql/help.c
index 3ebf7cc..0f34f41 100644
--- a/src/bin/psql/help.c
+++ b/src/bin/psql/help.c
@@ -162,13 +162,14 @@ slashUsage(unsigned short int pager)
 {
 	FILE	   *output;
 
-	output = PageOutput(94, pager);
+	output = PageOutput(95, pager);
 
 	/* if you add/remove a line here, change the row count above */
 
 	fprintf(output, _("General\n"));
 	fprintf(output, _("  \\copyright             show PostgreSQL usage and distribution terms\n"));
 	fprintf(output, _("  \\g [FILE] or ;         execute query (and send results to file or |pipe)\n"));
+	fprintf(output, _("  \\gset NAME [, NAME [..]] execute query and store result in internal variables\n"));
 	fprintf(output, _("  \\h [NAME]              help on syntax of SQL commands, * for all commands\n"));
 	fprintf(output, _("  \\q                     quit psql\n"));
 	fprintf(output, "\n");
diff --git a/src/bin/psql/mainloop.c b/src/bin/psql/mainloop.c
index 16c65ec..33faff8 100644
--- a/src/bin/psql/mainloop.c
+++ b/src/bin/psql/mainloop.c
@@ -327,6 +327,14 @@ MainLoop(FILE *source)
 					/* flush any paren nesting info after forced send */
 					psql_scan_reset(scan_state);
 				}
+				else if (slashCmdStatus == PSQL_CMD_NOSEND)
+				{
+					resetPQExpBuffer(query_buf);
+					/* reset parsing state since we are rescanning whole line */
+					psql_scan_reset(scan_state);
+					line_saved_in_history = true;
+					success = false;
+				}
 				else if (slashCmdStatus == PSQL_CMD_NEWEDIT)
 				{
 					/* rescan query_buf as new input */
diff --git a/src/bin/psql/psqlscan.h b/src/bin/psql/psqlscan.h
index 0b5ef0c..742de7f 100644
--- a/src/bin/psql/psqlscan.h
+++ b/src/bin/psql/psqlscan.h
@@ -8,6 +8,7 @@
 #ifndef PSQLSCAN_H
 #define PSQLSCAN_H
 
+#include "common.h"
 #include "pqexpbuffer.h"
 
 #include "prompt.h"
@@ -33,7 +34,8 @@ enum slash_option_type
 	OT_SQLIDHACK,				/* SQL identifier, but don't downcase */
 	OT_FILEPIPE,				/* it's a filename or pipe */
 	OT_WHOLE_LINE,				/* just snarf the rest of the line */
-	OT_NO_EVAL					/* no expansion of backticks or variables */
+	OT_NO_EVAL,					/* no expansion of backticks or variables */
+	OT_VARLIST				/* returns variable's identifier or comma */
 };
 
 
@@ -59,6 +61,10 @@ extern char *psql_scan_slash_option(PsqlScanState state,
 					   char *quote,
 					   bool semicolon);
 
+extern TargetList psql_scan_slash_vars(PsqlScanState state,
+							bool *error);
+
+
 extern void psql_scan_slash_command_end(PsqlScanState state);
 
 #endif   /* PSQLSCAN_H */
diff --git a/src/bin/psql/psqlscan.l b/src/bin/psql/psqlscan.l
index 1208c8f..15018af 100644
--- a/src/bin/psql/psqlscan.l
+++ b/src/bin/psql/psqlscan.l
@@ -106,13 +106,18 @@ static char *option_quote;
 static int	unquoted_option_chars;
 static int	backtick_start_offset;
 
+TargetList gvars;
+
 
 /* Return values from yylex() */
 #define LEXRES_EOL			0	/* end of input */
 #define LEXRES_SEMI			1	/* command-terminating semicolon found */
-#define LEXRES_BACKSLASH	2	/* backslash command start */
-#define LEXRES_OK			3	/* OK completion of backslash argument */
-
+#define LEXRES_COMMA			2
+#define LEXRES_BACKSLASH	3	/* backslash command start */
+#define LEXRES_IDENT			4	/* valid internal variable identifier */
+#define LEXRES_OK			5	/* OK completion of backslash argument */
+#define LEXRES_ERROR			6
+#define LEXRES_NONE			7
 
 int	yylex(void);
 
@@ -167,6 +172,7 @@ static void escape_variable(bool as_ident);
  *  <xdolq> $foo$ quoted strings
  *  <xui> quoted identifier with Unicode escapes
  *  <xus> quoted string with Unicode escapes
+ *  <xvl> comma separated list of variables
  *
  * Note: we intentionally don't mimic the backend's <xeu> state; we have
  * no need to distinguish it from <xe> state, and no good way to get out
@@ -183,6 +189,7 @@ static void escape_variable(bool as_ident);
 %x xdolq
 %x xui
 %x xus
+%x xvl
 /* Additional exclusive states for psql only: lex backslash commands */
 %x xslashcmd
 %x xslashargstart
@@ -628,6 +635,26 @@ other			.
 					ECHO;
 				}
 
+<xvl>{
+
+{identifier}	{
+					gvars = tglist_add(gvars, yytext);
+					return LEXRES_IDENT;
+				}
+
+","		{
+					return LEXRES_COMMA;
+				}
+
+{horiz_space}+	{ }
+
+. |
+\n		{
+					return LEXRES_ERROR;
+				}
+
+}
+
 {xufailed}	{
 					/* throw back all but the initial u/U */
 					yyless(1);
@@ -1449,6 +1476,75 @@ psql_scan_slash_command(PsqlScanState state)
 }
 
 /*
+ * returns identifier from identifier list
+ *
+ *    when comma_expected is true, when we require comma before identifier
+ *    error is true, when unexpected char was identified or missing
+ *        comma.
+ *
+ */
+TargetList
+psql_scan_slash_vars(PsqlScanState state,
+						bool *error)
+{
+	PQExpBufferData mybuf;
+	int			lexresult = LEXRES_NONE;
+	int		lexresult_prev;
+	bool	is_lexer_error = false;
+
+
+	/* Must be scanning already */
+	psql_assert(state->scanbufhandle);
+
+	/* Set up static variables that will be used by yylex */
+	cur_state = state;
+	output_buf = &mybuf;
+
+	if (state->buffer_stack != NULL)
+		yy_switch_to_buffer(state->buffer_stack->buf);
+	else
+		yy_switch_to_buffer(state->scanbufhandle);
+
+	gvars = NULL;
+
+	do
+	{
+		lexresult_prev = lexresult;
+
+		BEGIN(xvl);
+		lexresult = yylex();
+
+		/* read to end */
+		if (is_lexer_error)
+			continue;
+
+		if (lexresult == LEXRES_ERROR)
+		{
+			is_lexer_error = true;
+			continue;
+		}
+
+		if (lexresult_prev == LEXRES_IDENT && lexresult == LEXRES_IDENT)
+		{
+			is_lexer_error = true;
+			continue;
+		}
+
+		if (!is_lexer_error && (lexresult_prev == LEXRES_COMMA || lexresult_prev == LEXRES_NONE) 
+				&& lexresult == LEXRES_COMMA)
+			gvars = tglist_add(gvars, NULL);
+
+	} while (lexresult != LEXRES_EOL);
+
+	if (!is_lexer_error && lexresult_prev == LEXRES_COMMA)
+		gvars = tglist_add(gvars, NULL);
+
+	*error = is_lexer_error;
+
+	return gvars;
+}
+
+/*
  * Parse off the next argument for a backslash command, and return it as a
  * malloc'd string.  If there are no more arguments, returns NULL.
  *
diff --git a/src/bin/psql/settings.h b/src/bin/psql/settings.h
index c907fa0..e3381b2 100644
--- a/src/bin/psql/settings.h
+++ b/src/bin/psql/settings.h
@@ -9,6 +9,7 @@
 #define SETTINGS_H
 
 
+#include "common.h"
 #include "variables.h"
 #include "print.h"
 
@@ -73,6 +74,7 @@ typedef struct _psqlSettings
 	printQueryOpt popt;
 
 	char	   *gfname;			/* one-shot file output argument for \g */
+	TargetList  gvars;			/* one-shot target list argument for \gset */
 
 	bool		notty;			/* stdin or stdout is not a tty (as determined
 								 * on startup) */
diff --git a/src/bin/psql/tab-complete.c b/src/bin/psql/tab-complete.c
index a1bb230..b47a44a 100644
--- a/src/bin/psql/tab-complete.c
+++ b/src/bin/psql/tab-complete.c
@@ -811,7 +811,7 @@ psql_completion(char *text, int start, int end)
 		"\\dF", "\\dFd", "\\dFp", "\\dFt", "\\dg", "\\di", "\\dl", "\\dL",
 		"\\dn", "\\do", "\\dp", "\\drds", "\\ds", "\\dS", "\\dt", "\\dT", "\\dv", "\\du",
 		"\\e", "\\echo", "\\ef", "\\encoding",
-		"\\f", "\\g", "\\h", "\\help", "\\H", "\\i", "\\ir", "\\l",
+		"\\f", "\\g", "\\gset", "\\h", "\\help", "\\H", "\\i", "\\ir", "\\l",
 		"\\lo_import", "\\lo_export", "\\lo_list", "\\lo_unlink",
 		"\\o", "\\p", "\\password", "\\prompt", "\\pset", "\\q", "\\qecho", "\\r",
 		"\\set", "\\sf", "\\t", "\\T",