diff --git a/doc/src/sgml/ref/psql-ref.sgml b/doc/src/sgml/ref/psql-ref.sgml index 9915731..1d77569 100644 --- a/doc/src/sgml/ref/psql-ref.sgml +++ b/doc/src/sgml/ref/psql-ref.sgml @@ -2697,6 +2697,33 @@ lo_import 152801 + \set_from_bfile name filename + + + + Sets the psql variable name by content of + binary file filename. + The content is escaped as bytea value. + + + + + + + \set_from_file name filename + + + + Sets the psql variable name by content of + text file filename. + + + + + + \setenv name [ value ] diff --git a/src/bin/psql/command.c b/src/bin/psql/command.c index a9a2fdb..bcc1793 100644 --- a/src/bin/psql/command.c +++ b/src/bin/psql/command.c @@ -82,6 +82,7 @@ static void minimal_error_message(PGresult *res); static void printSSLInfo(void); static bool printPsetInfo(const char *param, struct printQueryOpt *popt); static char *pset_value_string(const char *param, struct printQueryOpt *popt); +static bool read_file(char *fname, PQExpBuffer rawbuf); #ifdef WIN32 static void checkWin32Codepage(void); @@ -1318,6 +1319,82 @@ exec_command(const char *cmd, free(opt0); } + /* \set_from_file, \set_from_bfile -- set variable/option command from file */ + else if (strcmp(cmd, "set_from_file") == 0 || strcmp(cmd, "set_from_bfile") == 0) + { + char *varname = psql_scan_slash_option(scan_state, + OT_NORMAL, NULL, false); + + if (!varname) + { + psql_error("\\%s: missing required argument\n", cmd); + success = false; + } + else + { + char *fname = psql_scan_slash_option(scan_state, + OT_FILEPIPE, NULL, true); + + if (!fname) + { + psql_error("\\%s: missing required argument\n", cmd); + success = false; + } + else + { + PQExpBufferData rawbuf; + + initPQExpBuffer(&rawbuf); + + expand_tilde(&fname); + canonicalize_path(fname); + + if (read_file(fname, &rawbuf)) + { + char *newval; + + /* do bytea escaping when it is required */ + if (strcmp(cmd, "set_from_bfile") == 0) + { + size_t escaped_size; + + newval = (char *) PQescapeByteaConn(pset.db, + (const unsigned char *) rawbuf.data, rawbuf.len, + &escaped_size); + } + else + newval = rawbuf.data; + + if (!newval) + { + psql_error("%s\n", PQerrorMessage(pset.db)); + success = false; + } + else + { + if (!SetVariable(pset.vars, varname, newval)) + { + psql_error("\\%s: error while setting variable\n", cmd); + success = false; + } + + /* release Bytea escaped result */ + if (newval != rawbuf.data) + PQfreemem(newval); + } + } + else + success = false; + + /* release raw content */ + termPQExpBuffer(&rawbuf); + + if (fname) + free(fname); + } + } + free(varname); + } /* \setenv -- set environment command */ else if (strcmp(cmd, "setenv") == 0) @@ -3657,3 +3734,53 @@ minimal_error_message(PGresult *res) destroyPQExpBuffer(msg); } + +/* + * file-content-fetching callback for read file content commands. + */ +static bool +read_file(char *fname, PQExpBuffer rawbuf) +{ + FILE *fd; + bool result = false; + + fd = fopen(fname, PG_BINARY_R); + if (fd) + { + struct stat fst; + + if (fstat(fileno(fd), &fst) != -1) + { + if (S_ISREG(fst.st_mode)) + { + if (fst.st_size <= ((int64) 1024) * 1024 * 1024) + { + size_t size; + char buf[512]; + + while ((size = fread(buf, 1, sizeof(buf), fd)) > 0) + appendBinaryPQExpBuffer(rawbuf, buf, size); + + if (ferror(fd)) + psql_error("%s: %s\n", fname, strerror(errno)); + else if (PQExpBufferBroken(rawbuf)) + psql_error("out of memory\n"); + else + result = true; + } + else + psql_error("%s is too big (greather than 1GB)\n", fname); + } + else + psql_error("%s is not regular file\n", fname); + } + else + psql_error("%s: %s\n", fname, strerror(errno)); + + fclose(fd); + } + else + psql_error("%s: %s\n", fname, strerror(errno)); + + return result; +} diff --git a/src/bin/psql/tab-complete.c b/src/bin/psql/tab-complete.c index cd64c39..c22046a 100644 --- a/src/bin/psql/tab-complete.c +++ b/src/bin/psql/tab-complete.c @@ -1339,8 +1339,9 @@ psql_completion(const char *text, int start, int end) "\\f", "\\g", "\\gexec", "\\gset", "\\h", "\\help", "\\H", "\\i", "\\ir", "\\l", "\\lo_import", "\\lo_export", "\\lo_list", "\\lo_unlink", "\\o", "\\p", "\\password", "\\prompt", "\\pset", "\\q", "\\qecho", "\\r", - "\\s", "\\set", "\\setenv", "\\sf", "\\sv", "\\t", "\\T", - "\\timing", "\\unset", "\\x", "\\w", "\\watch", "\\z", "\\!", NULL + "\\s", "\\set", "\\setenv", "\\set_from_bfile", "\\set_from_file", + "\\sf", "\\sv", "\\t", "\\T", "\\timing", "\\unset", "\\x", + "\\w", "\\watch", "\\z", "\\!", NULL }; (void) end; /* "end" is not used */ @@ -3236,6 +3237,15 @@ psql_completion(const char *text, int start, int end) else if (TailMatchesCS1("VERBOSITY")) COMPLETE_WITH_LIST_CS3("default", "verbose", "terse"); } + else if (TailMatchesCS1("\\set_from_bfile|\\set_from_file")) + { + matches = complete_from_variables(text, "", "", false); + } + else if (TailMatchesCS2("\\set_from_bfile|\\set_from_file", MatchAny)) + { + completion_charp = "\\"; + matches = completion_matches(text, complete_from_files); + } else if (TailMatchesCS1("\\sf*")) COMPLETE_WITH_SCHEMA_QUERY(Query_for_list_of_functions, NULL); else if (TailMatchesCS1("\\sv*")) diff --git a/src/test/regress/input/misc.source b/src/test/regress/input/misc.source index dd2d1b2..eead8ae 100644 --- a/src/test/regress/input/misc.source +++ b/src/test/regress/input/misc.source @@ -273,3 +273,21 @@ drop table oldstyle_test; -- -- rewrite rules -- + +--- +--- load file and store it to variable +--- +CREATE TABLE test_setref(a text, b bytea); + +-- use two different ways for import data - result should be same +\lo_import @abs_builddir@/data/hash.data +\set lo_oid :LASTOID +INSERT INTO test_setref (b) VALUES(lo_get(:lo_oid)); +\lo_unlink :lo_oid +SELECT md5(b) FROM test_setref; +TRUNCATE test_setref; + +\set_from_file var1 @abs_builddir@/data/hash.data +\set_from_bfile var2 @abs_builddir@/data/hash.data +INSERT INTO test_setref(a,b) VALUES(:'var1', :'var2'); +SELECT md5(a), md5(b) FROM test_setref; diff --git a/src/test/regress/output/misc.source b/src/test/regress/output/misc.source index 574ef0d..52c78fe 100644 --- a/src/test/regress/output/misc.source +++ b/src/test/regress/output/misc.source @@ -708,3 +708,28 @@ drop table oldstyle_test; -- -- rewrite rules -- +--- +--- load file and store it to variable +--- +CREATE TABLE test_setref(a text, b bytea); +-- use two different ways for import data - result should be same +\lo_import @abs_builddir@/data/hash.data +\set lo_oid :LASTOID +INSERT INTO test_setref (b) VALUES(lo_get(:lo_oid)); +\lo_unlink :lo_oid +SELECT md5(b) FROM test_setref; + md5 +---------------------------------- + e446fe6ea5a347e69670633412c7f8cb +(1 row) + +TRUNCATE test_setref; +\set_from_file var1 @abs_builddir@/data/hash.data +\set_from_bfile var2 @abs_builddir@/data/hash.data +INSERT INTO test_setref(a,b) VALUES(:'var1', :'var2'); +SELECT md5(a), md5(b) FROM test_setref; + md5 | md5 +----------------------------------+---------------------------------- + e446fe6ea5a347e69670633412c7f8cb | e446fe6ea5a347e69670633412c7f8cb +(1 row) +