Hi,
Currently \copy cannot span multiple lines (like any meta-command)
and cannot use psql variables whereas \g can do both.
The POC patch attached implements two meta-commands \copyfrom
and \copyto that are to COPY what \g is to any other query:
- they take the COPY query already var-expanded from the query buffer,
which must mention FROM STDIN or TO STDOUT.
- they accept an argument declaring the local data source or destination,
either a filename or a program (|command args) or empty for stdin/stdout.
By contrast \copy has a specific parser to extract the data source
or dest from its line of arguments, plus whether it's a COPY FROM or TO,
and build a COPY query from that.
Examples of use
1. $ psql -v filename="/path/data-$(date -I).csv"
COPY (SELECT *
FROM table
WHERE ...)
TO STDOUT (FORMAT csv) \copyto :filename
2. -- copy only the first 100 lines
COPY table FROM stdin \copyfrom |head -n 100 /data/datafile.txt
3. $ cat script.sql
COPY table1 FROM stdin; -- copy inline data
data line
data line
\.
-- copy data from psql's stdin
COPY table2 FROM stdin \copyfrom
# copy both in-script and out-of-script data
$ psql -f script.sql < table2.data
Comments? Does that look useful as an alternative to \copy ?
Best regards,
--
Daniel Vérité
PostgreSQL-powered mailer: http://www.manitou-mail.org
Twitter: @DanielVerite
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,
