[RFC] Add an until-0 loop in psql

[Pierre Ducroquet]

Hi

When running database migrations with .sql files on a live database, it's not 
uncommon to have to run a migration in a loop to prevent a big lock on a 
table.
For instance if one want to delete some old datas from a big table one would 
write :

DELETE FROM big_table WHERE id IN (SELECT id FROM big_table WHERE bad = true 
LIMIT 1000);
VACUUM big_table;

Right now, doing this is quite inefficient. We either have to write a script 
in another language, or run psql in a shell loop and wait for the migration to 
stop altering rows.

The attached **proof of concept** patch (I insist, it's a 15 minutes hack 
sprint with no previous knowledge of psql code) implements an 'until-0' loop 
in psql.
The previous migration could be simply written as :

\until-0
BEGIN;
DELETE FROM big_table WHERE id IN (SELECT id FROM big_table WHERE bad = true 
LIMIT 1000);
VACUUM big_table;
COMMIT;
\end-until

And psql will execute it until there is no row affected in the inner queries.

I am willing to write a proper patch for this (I hope the tell/seek is an 
acceptable implementation…), but I prefer having some feedback first.

Thanks

 Pierre
diff --git a/src/bin/psql/command.c b/src/bin/psql/command.c
index 4c85f43f09..d706e38ffc 100644
--- a/src/bin/psql/command.c
+++ b/src/bin/psql/command.c
@@ -401,6 +401,19 @@ exec_command(const char *cmd,
 		status = exec_command_shell_escape(scan_state, active_branch);
 	else if (strcmp(cmd, "?") == 0)
 		status = exec_command_slash_command_help(scan_state, active_branch);
+	else if (strcmp(cmd, "until-0") == 0) {
+		status = PSQL_CMD_SKIP_LINE;
+		pset.is_in_until = ftell(pset.cur_cmd_source);
+		pset.has_affected_rows = false;
+	} else if (strcmp(cmd, "end-until") == 0) {
+		status = PSQL_CMD_SKIP_LINE;
+		if (pset.has_affected_rows) {
+			fseek(pset.cur_cmd_source, pset.is_in_until, SEEK_SET);
+			pset.has_affected_rows = false;
+		} else {
+			pset.is_in_until = 0;
+		}
+	}
 	else
 		status = PSQL_CMD_UNKNOWN;
 
diff --git a/src/bin/psql/mainloop.c b/src/bin/psql/mainloop.c
index c06ce3ca09..869dbf6dcd 100644
--- a/src/bin/psql/mainloop.c
+++ b/src/bin/psql/mainloop.c
@@ -430,6 +430,12 @@ MainLoop(FILE *source)
 				{
 					success = SendQuery(query_buf->data);
 					slashCmdStatus = success ? PSQL_CMD_SEND : PSQL_CMD_ERROR;
+					if (success && pset.is_in_until && !pset.has_affected_rows)
+					{
+						const char *row_count = GetVariable(pset.vars, "ROW_COUNT");
+						if (row_count != NULL && strcmp(row_count, "0") != 0)
+							pset.has_affected_rows = true;
+					}
 					pset.stmt_lineno = 1;
 
 					/* transfer query to previous_buf by pointer-swapping */
diff --git a/src/bin/psql/settings.h b/src/bin/psql/settings.h
index 69e617e6b5..f3f92fe899 100644
--- a/src/bin/psql/settings.h
+++ b/src/bin/psql/settings.h
@@ -139,6 +139,9 @@ typedef struct _psqlSettings
 	const char *prompt3;
 	PGVerbosity verbosity;		/* current error verbosity level */
 	PGContextVisibility show_context;	/* current context display level */
+	
+	long		is_in_until;
+	int			has_affected_rows;
 } PsqlSettings;
 
 extern PsqlSettings pset;
diff --git a/src/bin/psql/startup.c b/src/bin/psql/startup.c
index be57574cd3..beba8851a3 100644
--- a/src/bin/psql/startup.c
+++ b/src/bin/psql/startup.c
@@ -134,6 +134,8 @@ main(int argc, char *argv[])
 	pset.last_error_result = NULL;
 	pset.cur_cmd_source = stdin;
 	pset.cur_cmd_interactive = false;
+	pset.is_in_until = 0;
+	pset.has_affected_rows = false;
 
 	/* We rely on unmentioned fields of pset.popt to start out 0/false/NULL */
 	pset.popt.topt.format = PRINT_ALIGNED;