@@ -122,7 +122,9 @@ static void make_directory(const char *dir);
122122
123123static void header (const char * fmt ,...) pg_attribute_printf (1 , 2 );
124124static void status (const char * fmt ,...) pg_attribute_printf (1 , 2 );
125- static void psql_command (const char * database , const char * query ,...) pg_attribute_printf (2 , 3 );
125+ static StringInfo psql_start_command (void );
126+ static void psql_add_command (StringInfo buf , const char * query ,...) pg_attribute_printf (2 , 3 );
127+ static void psql_end_command (StringInfo buf , const char * database );
126128
127129/*
128130 * allow core files if possible.
@@ -1136,51 +1138,94 @@ config_sspi_auth(const char *pgdata, const char *superuser_name)
11361138#endif /* ENABLE_SSPI */
11371139
11381140/*
1139- * Issue a command via psql, connecting to the specified database
1141+ * psql_start_command, psql_add_command, psql_end_command
1142+ *
1143+ * Issue one or more commands within one psql call.
1144+ * Set up with psql_start_command, then add commands one at a time
1145+ * with psql_add_command, and finally execute with psql_end_command.
11401146 *
11411147 * Since we use system(), this doesn't return until the operation finishes
11421148 */
1149+ static StringInfo
1150+ psql_start_command (void )
1151+ {
1152+ StringInfo buf = makeStringInfo ();
1153+
1154+ appendStringInfo (buf ,
1155+ "\"%s%spsql\" -X" ,
1156+ bindir ? bindir : "" ,
1157+ bindir ? "/" : "" );
1158+ return buf ;
1159+ }
1160+
11431161static void
1144- psql_command ( const char * database , const char * query ,...)
1162+ psql_add_command ( StringInfo buf , const char * query ,...)
11451163{
1146- char query_formatted [1024 ];
1147- char query_escaped [2048 ];
1148- char psql_cmd [MAXPGPATH + 2048 ];
1149- va_list args ;
1150- char * s ;
1151- char * d ;
1164+ StringInfoData cmdbuf ;
1165+ const char * cmdptr ;
1166+
1167+ /* Add each command as a -c argument in the psql call */
1168+ appendStringInfoString (buf , " -c \"" );
11521169
11531170/* Generate the query with insertion of sprintf arguments */
1154- va_start (args , query );
1155- vsnprintf (query_formatted , sizeof (query_formatted ), query , args );
1156- va_end (args );
1171+ initStringInfo (& cmdbuf );
1172+ for (;;)
1173+ {
1174+ va_list args ;
1175+ int needed ;
1176+
1177+ va_start (args , query );
1178+ needed = appendStringInfoVA (& cmdbuf , query , args );
1179+ va_end (args );
1180+ if (needed == 0 )
1181+ break ;/* success */
1182+ enlargeStringInfo (& cmdbuf , needed );
1183+ }
11571184
11581185/* Now escape any shell double-quote metacharacters */
1159- d = query_escaped ;
1160- for (s = query_formatted ; * s ; s ++ )
1186+ for (cmdptr = cmdbuf .data ; * cmdptr ; cmdptr ++ )
11611187{
1162- if (strchr ("\\\"$`" , * s ))
1163- * d ++ = '\\' ;
1164- * d ++ = * s ;
1188+ if (strchr ("\\\"$`" , * cmdptr ))
1189+ appendStringInfoChar ( buf , '\\' ) ;
1190+ appendStringInfoChar ( buf , * cmdptr ) ;
11651191}
1166- * d = '\0' ;
11671192
1168- /* And now we can build and execute the shell command */
1169- snprintf (psql_cmd , sizeof (psql_cmd ),
1170- "\"%s%spsql\" -X -c \"%s\" \"%s\"" ,
1171- bindir ? bindir : "" ,
1172- bindir ? "/" : "" ,
1173- query_escaped ,
1174- database );
1193+ appendStringInfoChar (buf , '"' );
1194+
1195+ pfree (cmdbuf .data );
1196+ }
1197+
1198+ static void
1199+ psql_end_command (StringInfo buf , const char * database )
1200+ {
1201+ /* Add the database name --- assume it needs no extra escaping */
1202+ appendStringInfo (buf ,
1203+ " \"%s\"" ,
1204+ database );
11751205
1176- if (system (psql_cmd ) != 0 )
1206+ /* And now we can execute the shell command */
1207+ if (system (buf -> data ) != 0 )
11771208{
11781209/* psql probably already reported the error */
1179- fprintf (stderr , _ ("command failed: %s\n" ), psql_cmd );
1210+ fprintf (stderr , _ ("command failed: %s\n" ), buf -> data );
11801211exit (2 );
11811212}
1213+
1214+ /* Clean up */
1215+ pfree (buf -> data );
1216+ pfree (buf );
11821217}
11831218
1219+ /*
1220+ * Shorthand macro for the common case of a single command
1221+ */
1222+ #define psql_command (database , ...) \
1223+ do { \
1224+ StringInfo cmdbuf = psql_start_command(); \
1225+ psql_add_command(cmdbuf, __VA_ARGS__); \
1226+ psql_end_command(cmdbuf, database); \
1227+ } while (0)
1228+
11841229/*
11851230 * Spawn a process to execute the given shell command; don't wait for it
11861231 *
@@ -2012,13 +2057,19 @@ open_result_files(void)
20122057static void
20132058drop_database_if_exists (const char * dbname )
20142059{
2060+ StringInfo buf = psql_start_command ();
2061+
20152062header (_ ("dropping database \"%s\"" ), dbname );
2016- psql_command ("postgres" , "DROP DATABASE IF EXISTS \"%s\"" , dbname );
2063+ /* Set warning level so we don't see chatter about nonexistent DB */
2064+ psql_add_command (buf , "SET client_min_messages = warning" );
2065+ psql_add_command (buf , "DROP DATABASE IF EXISTS \"%s\"" , dbname );
2066+ psql_end_command (buf , "postgres" );
20172067}
20182068
20192069static void
20202070create_database (const char * dbname )
20212071{
2072+ StringInfo buf = psql_start_command ();
20222073_stringlist * sl ;
20232074
20242075/*
@@ -2027,19 +2078,20 @@ create_database(const char *dbname)
20272078 */
20282079header (_ ("creating database \"%s\"" ), dbname );
20292080if (encoding )
2030- psql_command ( "postgres" , "CREATE DATABASE \"%s\" TEMPLATE=template0 ENCODING='%s'%s" , dbname , encoding ,
2031- (nolocale ) ? " LC_COLLATE='C' LC_CTYPE='C'" : "" );
2081+ psql_add_command ( buf , "CREATE DATABASE \"%s\" TEMPLATE=template0 ENCODING='%s'%s" , dbname , encoding ,
2082+ (nolocale ) ? " LC_COLLATE='C' LC_CTYPE='C'" : "" );
20322083else
2033- psql_command ("postgres" , "CREATE DATABASE \"%s\" TEMPLATE=template0%s" , dbname ,
2034- (nolocale ) ? " LC_COLLATE='C' LC_CTYPE='C'" : "" );
2035- psql_command (dbname ,
2036- "ALTER DATABASE \"%s\" SET lc_messages TO 'C';"
2037- "ALTER DATABASE \"%s\" SET lc_monetary TO 'C';"
2038- "ALTER DATABASE \"%s\" SET lc_numeric TO 'C';"
2039- "ALTER DATABASE \"%s\" SET lc_time TO 'C';"
2040- "ALTER DATABASE \"%s\" SET bytea_output TO 'hex';"
2041- "ALTER DATABASE \"%s\" SET timezone_abbreviations TO 'Default';" ,
2042- dbname , dbname , dbname , dbname , dbname , dbname );
2084+ psql_add_command (buf , "CREATE DATABASE \"%s\" TEMPLATE=template0%s" , dbname ,
2085+ (nolocale ) ? " LC_COLLATE='C' LC_CTYPE='C'" : "" );
2086+ psql_add_command (buf ,
2087+ "ALTER DATABASE \"%s\" SET lc_messages TO 'C';"
2088+ "ALTER DATABASE \"%s\" SET lc_monetary TO 'C';"
2089+ "ALTER DATABASE \"%s\" SET lc_numeric TO 'C';"
2090+ "ALTER DATABASE \"%s\" SET lc_time TO 'C';"
2091+ "ALTER DATABASE \"%s\" SET bytea_output TO 'hex';"
2092+ "ALTER DATABASE \"%s\" SET timezone_abbreviations TO 'Default';" ,
2093+ dbname , dbname , dbname , dbname , dbname , dbname );
2094+ psql_end_command (buf , "postgres" );
20432095
20442096/*
20452097 * Install any requested extensions. We use CREATE IF NOT EXISTS so that
@@ -2055,20 +2107,28 @@ create_database(const char *dbname)
20552107static void
20562108drop_role_if_exists (const char * rolename )
20572109{
2110+ StringInfo buf = psql_start_command ();
2111+
20582112header (_ ("dropping role \"%s\"" ), rolename );
2059- psql_command ("postgres" , "DROP ROLE IF EXISTS \"%s\"" , rolename );
2113+ /* Set warning level so we don't see chatter about nonexistent role */
2114+ psql_add_command (buf , "SET client_min_messages = warning" );
2115+ psql_add_command (buf , "DROP ROLE IF EXISTS \"%s\"" , rolename );
2116+ psql_end_command (buf , "postgres" );
20602117}
20612118
20622119static void
20632120create_role (const char * rolename , const _stringlist * granted_dbs )
20642121{
2122+ StringInfo buf = psql_start_command ();
2123+
20652124header (_ ("creating role \"%s\"" ), rolename );
2066- psql_command ( "postgres" , "CREATE ROLE \"%s\" WITH LOGIN" , rolename );
2125+ psql_add_command ( buf , "CREATE ROLE \"%s\" WITH LOGIN" , rolename );
20672126for (; granted_dbs != NULL ; granted_dbs = granted_dbs -> next )
20682127{
2069- psql_command ( "postgres" , "GRANT ALL ON DATABASE \"%s\" TO \"%s\"" ,
2070- granted_dbs -> str , rolename );
2128+ psql_add_command ( buf , "GRANT ALL ON DATABASE \"%s\" TO \"%s\"" ,
2129+ granted_dbs -> str , rolename );
20712130}
2131+ psql_end_command (buf , "postgres" );
20722132}
20732133
20742134static void
0 commit comments