Create a FETCH_COUNT parameter that causes psql to execute SELECT-like

queries via a cursor, fetching a limited number of rows at a time and
therefore not risking exhausting memory.  A disadvantage of the scheme
is that 'aligned' output mode will align each group of rows independently
leading to odd-looking output, but all the other output formats work
reasonably well.  Chris Mair, with some additional hacking by moi.

doc/src/sgml/ref/psql-ref.sgml

@@ -1,5 +1,5 @@
 <!--
-$PostgreSQL: pgsql/doc/src/sgml/ref/psql-ref.sgml,v 1.166 2006/07/27 19:52:04 tgl Exp $
+$PostgreSQL: pgsql/doc/src/sgml/ref/psql-ref.sgml,v 1.167 2006/08/29 22:25:04 tgl Exp $
 PostgreSQL documentation
 -->
 
@@ -2007,6 +2007,33 @@ bar
         </listitem>
       </varlistentry>
 
+      <varlistentry>
+        <term><varname>FETCH_COUNT</varname></term>
+        <listitem>
+        <para>
+        If this variable is set to an integer value &gt; 0,
+        the results of <command>SELECT</command> queries are fetched
+        and displayed in groups of that many rows, rather than the
+        default behavior of collecting the entire result set before
+        display.  Therefore only a
+        limited amount of memory is used, regardless of the size of
+        the result set.  Settings of 100 to 1000 are commonly used
+        when enabling this feature.
+        Keep in mind that when using this feature, a query may
+        fail after having already displayed some rows.
+        </para>
+        <tip>
+        <para>
+        Although you can use any output format with this feature,
+        the default <literal>aligned</> format tends to look bad
+        because each group of <varname>FETCH_COUNT</varname> rows
+        will be formatted separately, leading to varying column
+        widths across the row groups.  The other output formats work better.
+        </para>
+        </tip>
+        </listitem>
+      </varlistentry>
+
       <varlistentry>
         <term><varname>HISTCONTROL</varname></term>
         <listitem>

src/bin/psql/common.c

@@ -3,7 +3,7 @@
  *
  * Copyright (c) 2000-2006, PostgreSQL Global Development Group
  *
- * $PostgreSQL: pgsql/src/bin/psql/common.c,v 1.126 2006/08/29 15:19:50 tgl Exp $
+ * $PostgreSQL: pgsql/src/bin/psql/common.c,v 1.127 2006/08/29 22:25:07 tgl Exp $
  */
 #include "postgres_fe.h"
 #include "common.h"
@@ -28,6 +28,7 @@
 #include "command.h"
 #include "copy.h"
 #include "mb/pg_wchar.h"
+#include "mbprint.h"
 
 
 /* Workarounds for Windows */
@@ -53,7 +54,9 @@ typedef struct _timeb TimevalStruct;
 #endif
 
 
+static bool ExecQueryUsingCursor(const char *query, double *elapsed_msec);
 static bool command_no_begin(const char *query);
+static bool is_select_command(const char *query);
 
 /*
  * "Safe" wrapper around strdup()
@@ -450,18 +453,15 @@ ResetCancelConn(void)
  * AcceptResult
  *
  * Checks whether a result is valid, giving an error message if necessary;
- * resets cancelConn as needed, and ensures that the connection to the backend
- * is still up.
+ * and ensures that the connection to the backend is still up.
  *
  * Returns true for valid result, false for error state.
  */
 static bool
-AcceptResult(const PGresult *result, const char *query)
+AcceptResult(const PGresult *result)
 {
 	bool		OK = true;
 
-	ResetCancelConn();
-
 	if (!result)
 		OK = false;
 	else
@@ -560,7 +560,9 @@ PSQLexec(const char *query, bool start_xact)
 
 	res = PQexec(pset.db, query);
 
-	if (!AcceptResult(res, query) && res)
+	ResetCancelConn();
+
+	if (!AcceptResult(res))
 	{
 		PQclear(res);
 		res = NULL;
@@ -602,6 +604,7 @@ PrintQueryTuples(const PGresult *results)
 	/* write output to \g argument, if any */
 	if (pset.gfname)
 	{
+		/* keep this code in sync with ExecQueryUsingCursor */
 		FILE	   *queryFout_copy = pset.queryFout;
 		bool		queryFoutPipe_copy = pset.queryFoutPipe;
 
@@ -782,11 +785,10 @@ bool
 SendQuery(const char *query)
 {
 	PGresult   *results;
-	TimevalStruct before,
-				after;
+	PGTransactionStatusType transaction_status;
+	double		elapsed_msec = 0;
 	bool		OK,
 				on_error_rollback_savepoint = false;
-	PGTransactionStatusType transaction_status;
 	static bool on_error_rollback_warning = false;
 
 	if (!pset.db)
@@ -869,20 +871,38 @@ SendQuery(const char *query)
 		}
 	}
 
-	if (pset.timing)
-		GETTIMEOFDAY(&before);
+	if (pset.fetch_count <= 0 || !is_select_command(query))
+	{
+		/* Default fetch-it-all-and-print mode */
+		TimevalStruct before,
+				after;
 
-	results = PQexec(pset.db, query);
+		if (pset.timing)
+			GETTIMEOFDAY(&before);
 
-	/* these operations are included in the timing result: */
-	OK = (AcceptResult(results, query) && ProcessCopyResult(results));
+		results = PQexec(pset.db, query);
 
-	if (pset.timing)
-		GETTIMEOFDAY(&after);
+		/* these operations are included in the timing result: */
+		ResetCancelConn();
+		OK = (AcceptResult(results) && ProcessCopyResult(results));
 
-	/* but printing results isn't: */
-	if (OK)
-		OK = PrintQueryResults(results);
+		if (pset.timing)
+		{
+			GETTIMEOFDAY(&after);
+			elapsed_msec = DIFF_MSEC(&after, &before);
+		}
+
+		/* but printing results isn't: */
+		if (OK)
+			OK = PrintQueryResults(results);
+	}
+	else
+	{
+		/* Fetch-in-segments mode */
+		OK = ExecQueryUsingCursor(query, &elapsed_msec);
+		ResetCancelConn();
+		results = NULL;			/* PQclear(NULL) does nothing */
+	}
 
 	/* If we made a temporary savepoint, possibly release/rollback */
 	if (on_error_rollback_savepoint)
@@ -904,9 +924,10 @@ SendQuery(const char *query)
 			 * the user did RELEASE or ROLLBACK, our savepoint is gone. If
 			 * they issued a SAVEPOINT, releasing ours would remove theirs.
 			 */
-			if (strcmp(PQcmdStatus(results), "SAVEPOINT") == 0 ||
-				strcmp(PQcmdStatus(results), "RELEASE") == 0 ||
-				strcmp(PQcmdStatus(results), "ROLLBACK") == 0)
+			if (results &&
+				(strcmp(PQcmdStatus(results), "SAVEPOINT") == 0 ||
+				 strcmp(PQcmdStatus(results), "RELEASE") == 0 ||
+				 strcmp(PQcmdStatus(results), "ROLLBACK") == 0))
 				svptres = NULL;
 			else
 				svptres = PQexec(pset.db, "RELEASE pg_psql_temporary_savepoint");
@@ -927,7 +948,7 @@ SendQuery(const char *query)
 
 	/* Possible microtiming output */
 	if (OK && pset.timing)
-		printf(_("Time: %.3f ms\n"), DIFF_MSEC(&after, &before));
+		printf(_("Time: %.3f ms\n"), elapsed_msec);
 
 	/* check for events that may occur during query execution */
 
@@ -947,6 +968,198 @@ SendQuery(const char *query)
 }
 
 
+/*
+ * ExecQueryUsingCursor: run a SELECT-like query using a cursor
+ *
+ * This feature allows result sets larger than RAM to be dealt with.
+ *
+ * Returns true if the query executed successfully, false otherwise.
+ *
+ * If pset.timing is on, total query time (exclusive of result-printing) is
+ * stored into *elapsed_msec.
+ */
+static bool
+ExecQueryUsingCursor(const char *query, double *elapsed_msec)
+{
+	bool		OK = true;
+	PGresult		*results;
+	PQExpBufferData	buf;
+	printQueryOpt my_popt = pset.popt;
+	FILE	   *queryFout_copy = pset.queryFout;
+	bool		queryFoutPipe_copy = pset.queryFoutPipe;
+	bool			started_txn = false;
+	bool			did_pager = false;
+	int				ntuples;
+	char			fetch_cmd[64];
+	TimevalStruct before,
+					after;
+
+	*elapsed_msec = 0;
+
+	/* initialize print options for partial table output */
+	my_popt.topt.start_table = true;
+	my_popt.topt.stop_table = false;
+	my_popt.topt.prior_records = 0;
+
+	if (pset.timing)
+		GETTIMEOFDAY(&before);
+
+	/* if we're not in a transaction, start one */
+	if (PQtransactionStatus(pset.db) == PQTRANS_IDLE)
+	{
+		results = PQexec(pset.db, "BEGIN");
+		OK = AcceptResult(results) &&
+			(PQresultStatus(results) == PGRES_COMMAND_OK);
+		PQclear(results);
+		if (!OK)
+			return false;
+		started_txn = true;
+	}
+
+	/* Send DECLARE CURSOR */
+	initPQExpBuffer(&buf);
+	appendPQExpBuffer(&buf, "DECLARE _psql_cursor NO SCROLL CURSOR FOR\n%s",
+					  query);
+
+	results = PQexec(pset.db, buf.data);
+	OK = AcceptResult(results) &&
+		(PQresultStatus(results) == PGRES_COMMAND_OK);
+	PQclear(results);
+	termPQExpBuffer(&buf);
+	if (!OK)
+		goto cleanup;
+
+	if (pset.timing)
+	{
+		GETTIMEOFDAY(&after);
+		*elapsed_msec += DIFF_MSEC(&after, &before);
+	}
+
+	snprintf(fetch_cmd, sizeof(fetch_cmd),
+			 "FETCH FORWARD %d FROM _psql_cursor",
+			 pset.fetch_count);
+
+	/* prepare to write output to \g argument, if any */
+	if (pset.gfname)
+	{
+		/* keep this code in sync with PrintQueryTuples */
+		pset.queryFout = stdout;    /* so it doesn't get closed */
+
+		/* open file/pipe */
+		if (!setQFout(pset.gfname))
+		{
+			pset.queryFout = queryFout_copy;
+			pset.queryFoutPipe = queryFoutPipe_copy;
+			OK = false;
+			goto cleanup;
+		}
+	}
+
+	for (;;)
+	{
+		if (pset.timing)
+			GETTIMEOFDAY(&before);
+
+		/* get FETCH_COUNT tuples at a time */
+		results = PQexec(pset.db, fetch_cmd);
+		OK = AcceptResult(results) &&
+			(PQresultStatus(results) == PGRES_TUPLES_OK);
+
+		if (pset.timing)
+		{
+			GETTIMEOFDAY(&after);
+			*elapsed_msec += DIFF_MSEC(&after, &before);
+		}
+
+		if (!OK)
+		{
+			PQclear(results);
+			break;
+		}
+
+		ntuples = PQntuples(results);
+
+		if (ntuples < pset.fetch_count)
+		{
+			/* this is the last result set, so allow footer decoration */
+			my_popt.topt.stop_table = true;
+		}
+		else if (pset.queryFout == stdout && !did_pager)
+		{
+			/*
+			 * If query requires multiple result sets, hack to ensure that
+			 * only one pager instance is used for the whole mess
+			 */
+			pset.queryFout = PageOutput(100000, my_popt.topt.pager);
+			did_pager = true;
+		}
+
+		printQuery(results, &my_popt, pset.queryFout, pset.logfile);
+
+		/* after the first result set, disallow header decoration */
+		my_popt.topt.start_table = false;
+		my_popt.topt.prior_records += ntuples;
+
+		PQclear(results);
+
+		if (ntuples < pset.fetch_count || cancel_pressed)
+			break;
+	}
+
+	/* close \g argument file/pipe, restore old setting */
+	if (pset.gfname)
+	{
+		/* keep this code in sync with PrintQueryTuples */
+		setQFout(NULL);
+
+		pset.queryFout = queryFout_copy;
+		pset.queryFoutPipe = queryFoutPipe_copy;
+
+		free(pset.gfname);
+		pset.gfname = NULL;
+	}
+	else if (did_pager)
+	{
+		ClosePager(pset.queryFout);
+		pset.queryFout = queryFout_copy;
+		pset.queryFoutPipe = queryFoutPipe_copy;
+	}
+
+cleanup:
+	if (pset.timing)
+		GETTIMEOFDAY(&before);
+
+	/*
+	 * We try to close the cursor on either success or failure, but on
+	 * failure ignore the result (it's probably just a bleat about
+	 * being in an aborted transaction)
+	 */
+	results = PQexec(pset.db, "CLOSE _psql_cursor");
+	if (OK)
+	{
+		OK = AcceptResult(results) &&
+			(PQresultStatus(results) == PGRES_COMMAND_OK);
+	}
+	PQclear(results);
+
+	if (started_txn)
+	{
+		results = PQexec(pset.db, OK ? "COMMIT" : "ROLLBACK");
+		OK &= AcceptResult(results) &&
+			(PQresultStatus(results) == PGRES_COMMAND_OK);
+		PQclear(results);
+	}
+
+	if (pset.timing)
+	{
+		GETTIMEOFDAY(&after);
+		*elapsed_msec += DIFF_MSEC(&after, &before);
+	}
+
+	return OK;
+}
+
+
 /*
  * Advance the given char pointer over white space and SQL comments.
  */
@@ -1158,6 +1371,43 @@ command_no_begin(const char *query)
 }
 
 
+/*
+ * Check whether the specified command is a SELECT (or VALUES).
+ */
+static bool
+is_select_command(const char *query)
+{
+	int			wordlen;
+
+	/*
+	 * First advance over any whitespace, comments and left parentheses.
+	 */
+	for (;;)
+	{
+		query = skip_white_space(query);
+		if (query[0] == '(')
+			query++;
+		else
+			break;
+	}
+
+	/*
+	 * Check word length (since "selectx" is not "select").
+	 */
+	wordlen = 0;
+	while (isalpha((unsigned char) query[wordlen]))
+		wordlen += PQmblen(&query[wordlen], pset.encoding);
+
+	if (wordlen == 6 && pg_strncasecmp(query, "select", 6) == 0)
+		return true;
+	
+	if (wordlen == 6 && pg_strncasecmp(query, "values", 6) == 0)
+		return true;
+
+	return false;
+}
+
+
 /*
  * Test if the current user is a database superuser.
  *

src/bin/psql/print.h

@@ -3,7 +3,7 @@
  *
  * Copyright (c) 2000-2006, PostgreSQL Global Development Group
  *
- * $PostgreSQL: pgsql/src/bin/psql/print.h,v 1.31 2006/03/05 15:58:51 momjian Exp $
+ * $PostgreSQL: pgsql/src/bin/psql/print.h,v 1.32 2006/08/29 22:25:07 tgl Exp $
  */
 #ifndef PRINT_H
 #define PRINT_H
@@ -12,6 +12,7 @@
 
 
 extern FILE *PageOutput(int lines, unsigned short int pager);
+extern void ClosePager(FILE *pagerpipe);
 
 extern void html_escaped_print(const char *in, FILE *fout);
 
@@ -32,11 +33,14 @@ typedef struct _printTableOpt
 	enum printFormat format;	/* one of the above */
 	bool		expanded;		/* expanded/vertical output (if supported by
 								 * output format) */
+	unsigned short int border;	/* Print a border around the table. 0=none,
+								 * 1=dividing lines, 2=full */
 	unsigned short int pager;	/* use pager for output (if to stdout and
 								 * stdout is a tty) 0=off 1=on 2=always */
 	bool		tuples_only;	/* don't output headers, row counts, etc. */
-	unsigned short int border;	/* Print a border around the table. 0=none,
-								 * 1=dividing lines, 2=full */
+	bool		start_table;	/* print start decoration, eg <table> */
+	bool		stop_table;		/* print stop decoration, eg </table> */
+	unsigned long prior_records;	/* start offset for record counters */
 	char	   *fieldSep;		/* field separator for unaligned text mode */
 	char	   *recordSep;		/* record separator for unaligned text mode */
 	bool		numericLocale;	/* locale-aware numeric units separator and

src/bin/psql/settings.h

@@ -3,7 +3,7 @@
  *
  * Copyright (c) 2000-2006, PostgreSQL Global Development Group
  *
- * $PostgreSQL: pgsql/src/bin/psql/settings.h,v 1.29 2006/08/29 15:19:51 tgl Exp $
+ * $PostgreSQL: pgsql/src/bin/psql/settings.h,v 1.30 2006/08/29 22:25:07 tgl Exp $
  */
 #ifndef SETTINGS_H
 #define SETTINGS_H
@@ -96,6 +96,7 @@ typedef struct _psqlSettings
 	bool		quiet;
 	bool		singleline;
 	bool		singlestep;
+	int			fetch_count;
 	PSQL_ECHO	echo;
 	PSQL_ECHO_HIDDEN echo_hidden;
 	PSQL_ERROR_ROLLBACK on_error_rollback;

src/bin/psql/startup.c

@@ -3,7 +3,7 @@
  *
  * Copyright (c) 2000-2006, PostgreSQL Global Development Group
  *
- * $PostgreSQL: pgsql/src/bin/psql/startup.c,v 1.136 2006/08/29 15:19:51 tgl Exp $
+ * $PostgreSQL: pgsql/src/bin/psql/startup.c,v 1.137 2006/08/29 22:25:08 tgl Exp $
  */
 #include "postgres_fe.h"
 
@@ -145,6 +145,8 @@ main(int argc, char *argv[])
 	pset.popt.topt.format = PRINT_ALIGNED;
 	pset.popt.topt.border = 1;
 	pset.popt.topt.pager = 1;
+	pset.popt.topt.start_table = true;
+	pset.popt.topt.stop_table = true;
 	pset.popt.default_footer = true;
 
 	pset.notty = (!isatty(fileno(stdin)) || !isatty(fileno(stdout)));
@@ -798,6 +800,12 @@ singlestep_hook(const char *newval)
 	pset.singlestep = ParseVariableBool(newval);
 }
 
+static void
+fetch_count_hook(const char *newval)
+{
+	pset.fetch_count = ParseVariableNum(newval, -1, -1, false);
+}
+
 static void
 echo_hook(const char *newval)
 {
@@ -899,6 +907,7 @@ EstablishVariableSpace(void)
 	SetVariableAssignHook(pset.vars, "QUIET", quiet_hook);
 	SetVariableAssignHook(pset.vars, "SINGLELINE", singleline_hook);
 	SetVariableAssignHook(pset.vars, "SINGLESTEP", singlestep_hook);
+	SetVariableAssignHook(pset.vars, "FETCH_COUNT", fetch_count_hook);
 	SetVariableAssignHook(pset.vars, "ECHO", echo_hook);
 	SetVariableAssignHook(pset.vars, "ECHO_HIDDEN", echo_hidden_hook);
 	SetVariableAssignHook(pset.vars, "ON_ERROR_ROLLBACK", on_error_rollback_hook);