poc-patch-psql-copy-like-g.diff

diff --git a/src/bin/psql/command.c b/src/bin/psql/command.c
index 0dea54d..7398dec 100644
--- a/src/bin/psql/command.c
+++ b/src/bin/psql/command.c
@@ -65,6 +65,8 @@ static backslashResult exec_command_cd(PsqlScanState scan_state, bool active_bra
 				const char *cmd);
 static backslashResult exec_command_conninfo(PsqlScanState scan_state, bool active_branch);
 static backslashResult exec_command_copy(PsqlScanState scan_state, bool active_branch);
+static backslashResult exec_command_copyfrom(PQExpBuffer query_buf, PsqlScanState scan_state, bool active_branch);
+static backslashResult exec_command_copyto(PQExpBuffer query_buf, PsqlScanState scan_state, bool active_branch);
 static backslashResult exec_command_copyright(PsqlScanState scan_state, bool active_branch);
 static backslashResult exec_command_crosstabview(PsqlScanState scan_state, bool active_branch);
 static backslashResult exec_command_d(PsqlScanState scan_state, bool active_branch,
@@ -305,6 +307,10 @@ exec_command(const char *cmd,
 		status = exec_command_conninfo(scan_state, active_branch);
 	else if (pg_strcasecmp(cmd, "copy") == 0)
 		status = exec_command_copy(scan_state, active_branch);
+	else if (pg_strcasecmp(cmd, "copyfrom") == 0)
+		status = exec_command_copyfrom(query_buf, scan_state, active_branch);
+	else if (pg_strcasecmp(cmd, "copyto") == 0)
+		status = exec_command_copyto(query_buf, scan_state, active_branch);
 	else if (strcmp(cmd, "copyright") == 0)
 		status = exec_command_copyright(scan_state, active_branch);
 	else if (strcmp(cmd, "crosstabview") == 0)
@@ -634,6 +640,53 @@ exec_command_copy(PsqlScanState scan_state, bool active_branch)
 }
 
 /*
+ * \copyfrom -- run a COPY FROM command
+ * The COPY query is obtained from the query buffer
+ * The argument is 'filename' as in \copy
+ */
+static backslashResult
+exec_command_copyfrom(PQExpBuffer query_buf, PsqlScanState scan_state, bool active_branch)
+{
+	bool		success = true;
+
+	if (active_branch)
+	{
+		char	   *opt = psql_scan_slash_option(scan_state,
+												 OT_FILEPIPE, NULL, false);
+
+		success = do_copy_query(query_buf->data, opt, copy_from);
+		resetPQExpBuffer(query_buf);
+		free(opt);
+	}
+	else
+		ignore_slash_filepipe(scan_state);
+
+	return success ? PSQL_CMD_SKIP_LINE : PSQL_CMD_ERROR;
+}
+/*
+ * \copyto -- run a COPY TO command
+ */
+static backslashResult
+exec_command_copyto(PQExpBuffer query_buf, PsqlScanState scan_state, bool active_branch)
+{
+	bool		success = true;
+
+	if (active_branch)
+	{
+		char	   *opt = psql_scan_slash_option(scan_state,
+												 OT_FILEPIPE, NULL, false);
+
+		success = do_copy_query(query_buf->data, opt, copy_to);
+		resetPQExpBuffer(query_buf);
+		free(opt);
+	}
+	else
+		ignore_slash_filepipe(scan_state);
+
+	return success ? PSQL_CMD_SKIP_LINE : PSQL_CMD_ERROR;
+}
+
+/*
  * \copyright -- print copyright notice
  */
 static backslashResult
diff --git a/src/bin/psql/copy.c b/src/bin/psql/copy.c
index 555c633..95c5c5e 100644
--- a/src/bin/psql/copy.c
+++ b/src/bin/psql/copy.c
@@ -260,28 +260,11 @@ error:
 }
 
 
-/*
- * Execute a \copy command (frontend copy). We have to open a file (or execute
- * a command), then submit a COPY query to the backend and either feed it data
- * from the file or route its response into the file.
- */
-bool
-do_copy(const char *args)
+static
+FILE*
+open_copy_stream(struct copy_options *options)
 {
-	PQExpBufferData query;
-	FILE	   *copystream;
-	struct copy_options *options;
-	bool		success;
-
-	/* parse options */
-	options = parse_slash_copy(args);
-
-	if (!options)
-		return false;
-
-	/* prepare to read or write the target file */
-	if (options->file && !options->program)
-		canonicalize_path(options->file);
+	FILE *copystream = NULL;
 
 	if (options->from)
 	{
@@ -331,8 +314,7 @@ do_copy(const char *args)
 		else
 			psql_error("%s: %s\n",
 					   options->file, strerror(errno));
-		free_copy_options(options);
-		return false;
+		return NULL;
 	}
 
 	if (!options->program)
@@ -352,27 +334,18 @@ do_copy(const char *args)
 		if (result < 0 || S_ISDIR(st.st_mode))
 		{
 			fclose(copystream);
-			free_copy_options(options);
-			return false;
+			return NULL;
 		}
 	}
 
-	/* build the command we will send to the backend */
-	initPQExpBuffer(&query);
-	printfPQExpBuffer(&query, "COPY ");
-	appendPQExpBufferStr(&query, options->before_tofrom);
-	if (options->from)
-		appendPQExpBufferStr(&query, " FROM STDIN ");
-	else
-		appendPQExpBufferStr(&query, " TO STDOUT ");
-	if (options->after_tofrom)
-		appendPQExpBufferStr(&query, options->after_tofrom);
+	return copystream;
+}
 
-	/* run it like a user command, but with copystream as data source/sink */
-	pset.copyStream = copystream;
-	success = SendQuery(query.data);
-	pset.copyStream = NULL;
-	termPQExpBuffer(&query);
+static
+bool
+close_copy_stream(FILE* copystream, struct copy_options *options)
+{
+	bool success = true;
 
 	if (options->file != NULL)
 	{
@@ -407,6 +380,119 @@ do_copy(const char *args)
 			}
 		}
 	}
+	return success;
+}
+
+/*
+ * Execute a \copy command (frontend copy). We have to open a file (or execute
+ * a command), then submit a COPY query to the backend and either feed it data
+ * from the file or route its response into the file.
+ */
+bool
+do_copy(const char *args)
+{
+	PQExpBufferData query;
+	FILE	   *copystream;
+	struct copy_options *options;
+	bool		success;
+
+	/* parse options */
+	options = parse_slash_copy(args);
+
+	if (!options)
+		return false;
+
+	/* prepare to read or write the target file */
+	if (options->file && !options->program)
+		canonicalize_path(options->file);
+
+	/* open the file or program from which or to which data goes */
+	copystream = open_copy_stream(options);
+
+	if (!copystream)
+	{
+		/* a user-visible error has already been emitted at this point, so
+		 * just clean up and return */
+		free_copy_options(options);
+		return false;
+	}
+
+	/* build the command we will send to the backend */
+	initPQExpBuffer(&query);
+	printfPQExpBuffer(&query, "COPY ");
+	appendPQExpBufferStr(&query, options->before_tofrom);
+	if (options->from)
+		appendPQExpBufferStr(&query, " FROM STDIN ");
+	else
+		appendPQExpBufferStr(&query, " TO STDOUT ");
+	if (options->after_tofrom)
+		appendPQExpBufferStr(&query, options->after_tofrom);
+
+	/* run it like a user command, but with copystream as data source/sink */
+	pset.copyStream = copystream;
+	success = SendQuery(query.data);
+
+	pset.copyStream = NULL;
+	termPQExpBuffer(&query);
+	success = close_copy_stream(copystream, options);
+	free_copy_options(options);
+	return success;
+}
+
+/*
+ * Send the COPY query (frontend copy) in the query buffer, and feed it data
+ * from the file or program given as argument (when dir=copy_from),
+ * or route its response into the file or program (when dir=copy_to).
+ */
+bool
+do_copy_query(const char* querystring, const char *args, enum copy_direction dir)
+{
+	FILE	   *copystream;
+	struct copy_options *options;
+	bool		success;
+
+	options = pg_malloc0(sizeof(struct copy_options));
+	if (!options)
+		return false;
+
+	options->from = (dir == copy_from);
+
+	if (args)
+	{
+		if (args[0] == '|')
+		{
+			options->program = true;
+			options->file = pg_strdup(args+1);
+		}
+		else
+		{
+			options->file = pg_strdup(args);
+			canonicalize_path(options->file);
+		}
+	}
+	else
+	{
+		/* Use psql's stdin or stdout if no file or program given */
+		options->psql_inout = true;
+	}
+
+	/* open the file or program from which or to which data goes */
+	copystream = open_copy_stream(options);
+	if (!copystream)
+	{
+		/* a user-visible error has already been emitted at this point, so
+		 * just clean up and return */
+		free_copy_options(options);
+		return false;
+	}
+
+	/* run the COPY query with copystream as data source/sink */
+	pset.copyStream = copystream;
+	success = SendQuery(querystring);
+	pset.copyStream = NULL;
+
+	success = close_copy_stream(copystream, options);
+
 	free_copy_options(options);
 	return success;
 }
diff --git a/src/bin/psql/copy.h b/src/bin/psql/copy.h
index f4107d7..7755ac7 100644
--- a/src/bin/psql/copy.h
+++ b/src/bin/psql/copy.h
@@ -10,10 +10,18 @@
 
 #include "libpq-fe.h"
 
+enum copy_direction
+{
+	copy_from = 1,
+	copy_to
+};
 
 /* handler for \copy */
 extern bool do_copy(const char *args);
 
+/* handler for \copyfrom and \copyto */
+extern bool do_copy_query(const char* querystring, const char *args, enum copy_direction dir);
+
 /* lower level processors for copy in/out streams */
 
 extern bool handleCopyOut(PGconn *conn, FILE *copystream,