Re: ToDo: allow to get a number of processed rows by COPY statement [Review of Patch]

*** a/contrib/pg_stat_statements/pg_stat_statements.c
--- b/contrib/pg_stat_statements/pg_stat_statements.c
***************
*** 242,248 **** static void pgss_ExecutorEnd(QueryDesc *queryDesc);
  static void pgss_ProcessUtility(Node *parsetree,
  			  const char *queryString, ParamListInfo params,
  					DestReceiver *dest, char *completionTag,
! 					ProcessUtilityContext context);
  static uint32 pgss_hash_fn(const void *key, Size keysize);
  static int	pgss_match_fn(const void *key1, const void *key2, Size keysize);
  static uint32 pgss_hash_string(const char *str);
--- 242,248 ----
  static void pgss_ProcessUtility(Node *parsetree,
  			  const char *queryString, ParamListInfo params,
  					DestReceiver *dest, char *completionTag,
! 					ProcessUtilityContext context, uint64 *processed);
  static uint32 pgss_hash_fn(const void *key, Size keysize);
  static int	pgss_match_fn(const void *key1, const void *key2, Size keysize);
  static uint32 pgss_hash_string(const char *str);
***************
*** 787,793 **** pgss_ExecutorEnd(QueryDesc *queryDesc)
  static void
  pgss_ProcessUtility(Node *parsetree, const char *queryString,
  					ParamListInfo params, DestReceiver *dest,
! 					char *completionTag, ProcessUtilityContext context)
  {
  	/*
  	 * If it's an EXECUTE statement, we don't track it and don't increment the
--- 787,794 ----
  static void
  pgss_ProcessUtility(Node *parsetree, const char *queryString,
  					ParamListInfo params, DestReceiver *dest,
! 					char *completionTag, ProcessUtilityContext context,
! 					uint64 *processed)
  {
  	/*
  	 * If it's an EXECUTE statement, we don't track it and don't increment the
***************
*** 820,829 **** pgss_ProcessUtility(Node *parsetree, const char *queryString,
  		{
  			if (prev_ProcessUtility)
  				prev_ProcessUtility(parsetree, queryString, params,
! 									dest, completionTag, context);
  			else
  				standard_ProcessUtility(parsetree, queryString, params,
! 										dest, completionTag, context);
  			nested_level--;
  		}
  		PG_CATCH();
--- 821,830 ----
  		{
  			if (prev_ProcessUtility)
  				prev_ProcessUtility(parsetree, queryString, params,
! 									dest, completionTag, context, NULL);
  			else
  				standard_ProcessUtility(parsetree, queryString, params,
! 										dest, completionTag, context, NULL);
  			nested_level--;
  		}
  		PG_CATCH();
***************
*** 881,890 **** pgss_ProcessUtility(Node *parsetree, const char *queryString,
  	{
  		if (prev_ProcessUtility)
  			prev_ProcessUtility(parsetree, queryString, params,
! 								dest, completionTag, context);
  		else
  			standard_ProcessUtility(parsetree, queryString, params,
! 									dest, completionTag, context);
  	}
  }
  
--- 882,891 ----
  	{
  		if (prev_ProcessUtility)
  			prev_ProcessUtility(parsetree, queryString, params,
! 								dest, completionTag, context, NULL);
  		else
  			standard_ProcessUtility(parsetree, queryString, params,
! 									dest, completionTag, context, NULL);
  	}
  }
  
*** a/contrib/sepgsql/hooks.c
--- b/contrib/sepgsql/hooks.c
***************
*** 267,273 **** sepgsql_utility_command(Node *parsetree,
  						ParamListInfo params,
  						DestReceiver *dest,
  						char *completionTag,
! 						ProcessUtilityContext context)
  {
  	sepgsql_context_info_t saved_context_info = sepgsql_context_info;
  	ListCell   *cell;
--- 267,274 ----
  						ParamListInfo params,
  						DestReceiver *dest,
  						char *completionTag,
! 						ProcessUtilityContext context,
! 						uint64 *processed)
  {
  	sepgsql_context_info_t saved_context_info = sepgsql_context_info;
  	ListCell   *cell;
***************
*** 328,337 **** sepgsql_utility_command(Node *parsetree,
  
  		if (next_ProcessUtility_hook)
  			(*next_ProcessUtility_hook) (parsetree, queryString, params,
! 										 dest, completionTag, context);
  		else
  			standard_ProcessUtility(parsetree, queryString, params,
! 									dest, completionTag, context);
  	}
  	PG_CATCH();
  	{
--- 329,338 ----
  
  		if (next_ProcessUtility_hook)
  			(*next_ProcessUtility_hook) (parsetree, queryString, params,
! 										 dest, completionTag, context, NULL);
  		else
  			standard_ProcessUtility(parsetree, queryString, params,
! 									dest, completionTag, context, NULL);
  	}
  	PG_CATCH();
  	{
*** a/src/backend/commands/extension.c
--- b/src/backend/commands/extension.c
***************
*** 753,759 **** execute_sql_string(const char *sql, const char *filename)
  							   NULL,
  							   dest,
  							   NULL,
! 							   PROCESS_UTILITY_QUERY);
  			}
  
  			PopActiveSnapshot();
--- 753,760 ----
  							   NULL,
  							   dest,
  							   NULL,
! 							   PROCESS_UTILITY_QUERY,
! 							   NULL);
  			}
  
  			PopActiveSnapshot();
*** a/src/backend/commands/schemacmds.c
--- b/src/backend/commands/schemacmds.c
***************
*** 135,141 **** CreateSchemaCommand(CreateSchemaStmt *stmt, const char *queryString)
  					   NULL,
  					   None_Receiver,
  					   NULL,
! 					   PROCESS_UTILITY_SUBCOMMAND);
  		/* make sure later steps can see the object created here */
  		CommandCounterIncrement();
  	}
--- 135,142 ----
  					   NULL,
  					   None_Receiver,
  					   NULL,
! 					   PROCESS_UTILITY_SUBCOMMAND,
! 					   NULL);
  		/* make sure later steps can see the object created here */
  		CommandCounterIncrement();
  	}
*** a/src/backend/commands/trigger.c
--- b/src/backend/commands/trigger.c
***************
*** 1012,1018 **** ConvertTriggerToFK(CreateTrigStmt *stmt, Oid funcoid)
  		/* ... and execute it */
  		ProcessUtility((Node *) atstmt,
  					   "(generated ALTER TABLE ADD FOREIGN KEY command)",
! 					   NULL, None_Receiver, NULL, PROCESS_UTILITY_GENERATED);
  
  		/* Remove the matched item from the list */
  		info_list = list_delete_ptr(info_list, info);
--- 1012,1019 ----
  		/* ... and execute it */
  		ProcessUtility((Node *) atstmt,
  					   "(generated ALTER TABLE ADD FOREIGN KEY command)",
! 					   NULL, None_Receiver, NULL, PROCESS_UTILITY_GENERATED,
! 					   NULL);
  
  		/* Remove the matched item from the list */
  		info_list = list_delete_ptr(info_list, info);
*** a/src/backend/executor/functions.c
--- b/src/backend/executor/functions.c
***************
*** 786,792 **** postquel_getnext(execution_state *es, SQLFunctionCachePtr fcache)
  					   es->qd->params,
  					   es->qd->dest,
  					   NULL,
! 					   PROCESS_UTILITY_QUERY);
  		result = true;			/* never stops early */
  	}
  	else
--- 786,793 ----
  					   es->qd->params,
  					   es->qd->dest,
  					   NULL,
! 					   PROCESS_UTILITY_QUERY,
! 					   NULL);
  		result = true;			/* never stops early */
  	}
  	else
*** a/src/backend/executor/spi.c
--- b/src/backend/executor/spi.c
***************
*** 1909,1921 **** _SPI_execute_plan(SPIPlanPtr plan, ParamListInfo paramLI,
  			else
  			{
  				char		completionTag[COMPLETION_TAG_BUFSIZE];
  
  				ProcessUtility(stmt,
  							   plansource->query_string,
  							   paramLI,
  							   dest,
  							   completionTag,
! 							   PROCESS_UTILITY_QUERY);
  
  				/* Update "processed" if stmt returned tuples */
  				if (_SPI_current->tuptable)
--- 1909,1925 ----
  			else
  			{
  				char		completionTag[COMPLETION_TAG_BUFSIZE];
+ 				uint64		processed;
  
  				ProcessUtility(stmt,
  							   plansource->query_string,
  							   paramLI,
  							   dest,
  							   completionTag,
! 							   PROCESS_UTILITY_QUERY,
! 							   &processed);
! 
! 				_SPI_current->processed = (uint32) processed;
  
  				/* Update "processed" if stmt returned tuples */
  				if (_SPI_current->tuptable)
*** a/src/backend/tcop/pquery.c
--- b/src/backend/tcop/pquery.c
***************
*** 44,53 **** static uint32 RunFromStore(Portal portal, ScanDirection direction, long count,
  static long PortalRunSelect(Portal portal, bool forward, long count,
  				DestReceiver *dest);
  static void PortalRunUtility(Portal portal, Node *utilityStmt, bool isTopLevel,
! 				 DestReceiver *dest, char *completionTag);
  static void PortalRunMulti(Portal portal, bool isTopLevel,
  			   DestReceiver *dest, DestReceiver *altdest,
! 			   char *completionTag);
  static long DoPortalRunFetch(Portal portal,
  				 FetchDirection fdirection,
  				 long count,
--- 44,54 ----
  static long PortalRunSelect(Portal portal, bool forward, long count,
  				DestReceiver *dest);
  static void PortalRunUtility(Portal portal, Node *utilityStmt, bool isTopLevel,
! 				 DestReceiver *dest, char *completionTag,
! 				 uint64 *processed);
  static void PortalRunMulti(Portal portal, bool isTopLevel,
  			   DestReceiver *dest, DestReceiver *altdest,
! 			   char *completionTag, uint64 *processed);
  static long DoPortalRunFetch(Portal portal,
  				 FetchDirection fdirection,
  				 long count,
***************
*** 813,819 **** PortalRun(Portal portal, long count, bool isTopLevel,
  
  			case PORTAL_MULTI_QUERY:
  				PortalRunMulti(portal, isTopLevel,
! 							   dest, altdest, completionTag);
  
  				/* Prevent portal's commands from being re-executed */
  				MarkPortalDone(portal);
--- 814,821 ----
  
  			case PORTAL_MULTI_QUERY:
  				PortalRunMulti(portal, isTopLevel,
! 							   dest, altdest, completionTag,
! 							   &portal->processed);
  
  				/* Prevent portal's commands from being re-executed */
  				MarkPortalDone(portal);
***************
*** 1053,1064 **** FillPortalStore(Portal portal, bool isTopLevel)
  			 * tuplestore. Auxiliary query outputs are discarded.
  			 */
  			PortalRunMulti(portal, isTopLevel,
! 						   treceiver, None_Receiver, completionTag);
  			break;
  
  		case PORTAL_UTIL_SELECT:
  			PortalRunUtility(portal, (Node *) linitial(portal->stmts),
! 							 isTopLevel, treceiver, completionTag);
  			break;
  
  		default:
--- 1055,1068 ----
  			 * tuplestore. Auxiliary query outputs are discarded.
  			 */
  			PortalRunMulti(portal, isTopLevel,
! 						   treceiver, None_Receiver, completionTag,
! 						   &portal->processed);
  			break;
  
  		case PORTAL_UTIL_SELECT:
  			PortalRunUtility(portal, (Node *) linitial(portal->stmts),
! 							 isTopLevel, treceiver, completionTag,
! 							 &portal->processed);
  			break;
  
  		default:
***************
*** 1148,1154 **** RunFromStore(Portal portal, ScanDirection direction, long count,
   */
  static void
  PortalRunUtility(Portal portal, Node *utilityStmt, bool isTopLevel,
! 				 DestReceiver *dest, char *completionTag)
  {
  	bool		active_snapshot_set;
  
--- 1152,1159 ----
   */
  static void
  PortalRunUtility(Portal portal, Node *utilityStmt, bool isTopLevel,
! 				 DestReceiver *dest, char *completionTag,
! 				 uint64 *processed)
  {
  	bool		active_snapshot_set;
  
***************
*** 1189,1195 **** PortalRunUtility(Portal portal, Node *utilityStmt, bool isTopLevel,
  				   dest,
  				   completionTag,
  				   isTopLevel ?
! 					PROCESS_UTILITY_TOPLEVEL : PROCESS_UTILITY_QUERY);
  
  	/* Some utility statements may change context on us */
  	MemoryContextSwitchTo(PortalGetHeapMemory(portal));
--- 1194,1201 ----
  				   dest,
  				   completionTag,
  				   isTopLevel ?
! 					PROCESS_UTILITY_TOPLEVEL : PROCESS_UTILITY_QUERY,
! 				   processed);
  
  	/* Some utility statements may change context on us */
  	MemoryContextSwitchTo(PortalGetHeapMemory(portal));
***************
*** 1213,1219 **** PortalRunUtility(Portal portal, Node *utilityStmt, bool isTopLevel,
  static void
  PortalRunMulti(Portal portal, bool isTopLevel,
  			   DestReceiver *dest, DestReceiver *altdest,
! 			   char *completionTag)
  {
  	bool		active_snapshot_set = false;
  	ListCell   *stmtlist_item;
--- 1219,1226 ----
  static void
  PortalRunMulti(Portal portal, bool isTopLevel,
  			   DestReceiver *dest, DestReceiver *altdest,
! 			   char *completionTag,
! 			   uint64 *processed)
  {
  	bool		active_snapshot_set = false;
  	ListCell   *stmtlist_item;
***************
*** 1316,1329 **** PortalRunMulti(Portal portal, bool isTopLevel,
  				Assert(!active_snapshot_set);
  				/* statement can set tag string */
  				PortalRunUtility(portal, stmt, isTopLevel,
! 								 dest, completionTag);
  			}
  			else
  			{
  				Assert(IsA(stmt, NotifyStmt));
  				/* stmt added by rewrite cannot set tag */
  				PortalRunUtility(portal, stmt, isTopLevel,
! 								 altdest, NULL);
  			}
  		}
  
--- 1323,1338 ----
  				Assert(!active_snapshot_set);
  				/* statement can set tag string */
  				PortalRunUtility(portal, stmt, isTopLevel,
! 								 dest, completionTag,
! 								 processed);
  			}
  			else
  			{
  				Assert(IsA(stmt, NotifyStmt));
  				/* stmt added by rewrite cannot set tag */
  				PortalRunUtility(portal, stmt, isTopLevel,
! 								 altdest, NULL,
! 								 processed);
  			}
  		}
  
*** a/src/backend/tcop/utility.c
--- b/src/backend/tcop/utility.c
***************
*** 323,329 **** ProcessUtility(Node *parsetree,
  			   ParamListInfo params,
  			   DestReceiver *dest,
  			   char *completionTag,
! 			   ProcessUtilityContext context)
  {
  	Assert(queryString != NULL);	/* required as of 8.4 */
  
--- 323,330 ----
  			   ParamListInfo params,
  			   DestReceiver *dest,
  			   char *completionTag,
! 			   ProcessUtilityContext context,
! 			   uint64 *processed)
  {
  	Assert(queryString != NULL);	/* required as of 8.4 */
  
***************
*** 334,343 **** ProcessUtility(Node *parsetree,
  	 */
  	if (ProcessUtility_hook)
  		(*ProcessUtility_hook) (parsetree, queryString, params,
! 								dest, completionTag, context);
  	else
  		standard_ProcessUtility(parsetree, queryString, params,
! 								dest, completionTag, context);
  }
  
  void
--- 335,346 ----
  	 */
  	if (ProcessUtility_hook)
  		(*ProcessUtility_hook) (parsetree, queryString, params,
! 								dest, completionTag, context,
! 								processed);
  	else
  		standard_ProcessUtility(parsetree, queryString, params,
! 								dest, completionTag, context,
! 								processed);
  }
  
  void
***************
*** 346,352 **** standard_ProcessUtility(Node *parsetree,
  						ParamListInfo params,
  						DestReceiver *dest,
  						char *completionTag,
! 						ProcessUtilityContext context)
  {
  	bool		isTopLevel = (context == PROCESS_UTILITY_TOPLEVEL);
  	bool		isCompleteQuery = (context <= PROCESS_UTILITY_QUERY);
--- 349,356 ----
  						ParamListInfo params,
  						DestReceiver *dest,
  						char *completionTag,
! 						ProcessUtilityContext context,
! 						uint64 *processed)
  {
  	bool		isTopLevel = (context == PROCESS_UTILITY_TOPLEVEL);
  	bool		isCompleteQuery = (context <= PROCESS_UTILITY_QUERY);
***************
*** 356,361 **** standard_ProcessUtility(Node *parsetree,
--- 360,368 ----
  	if (completionTag)
  		completionTag[0] = '\0';
  
+ 	if (processed != NULL)
+ 		*processed = 0;
+ 
  	switch (nodeTag(parsetree))
  	{
  			/*
***************
*** 576,582 **** standard_ProcessUtility(Node *parsetree,
  									   params,
  									   None_Receiver,
  									   NULL,
! 									   PROCESS_UTILITY_GENERATED);
  					}
  
  					/* Need CCI between commands */
--- 583,590 ----
  									   params,
  									   None_Receiver,
  									   NULL,
! 									   PROCESS_UTILITY_GENERATED,
! 									   processed);
  					}
  
  					/* Need CCI between commands */
***************
*** 706,717 **** standard_ProcessUtility(Node *parsetree,
  
  		case T_CopyStmt:
  			{
! 				uint64		processed;
  
! 				processed = DoCopy((CopyStmt *) parsetree, queryString);
  				if (completionTag)
  					snprintf(completionTag, COMPLETION_TAG_BUFSIZE,
! 							 "COPY " UINT64_FORMAT, processed);
  			}
  			break;
  
--- 714,728 ----
  
  		case T_CopyStmt:
  			{
! 				uint64 processed_rows;
  
! 				processed_rows = DoCopy((CopyStmt *) parsetree, queryString);
  				if (completionTag)
  					snprintf(completionTag, COMPLETION_TAG_BUFSIZE,
! 							 "COPY " UINT64_FORMAT, processed_rows);
! 
! 				if (processed != NULL)
! 					*processed = processed_rows;
  			}
  			break;
  
***************
*** 813,819 **** standard_ProcessUtility(Node *parsetree,
  										   params,
  										   None_Receiver,
  										   NULL,
! 										   PROCESS_UTILITY_GENERATED);
  						}
  
  						/* Need CCI between commands */
--- 824,831 ----
  										   params,
  										   None_Receiver,
  										   NULL,
! 										   PROCESS_UTILITY_GENERATED,
! 										   processed);
  						}
  
  						/* Need CCI between commands */
*** a/src/include/tcop/utility.h
--- b/src/include/tcop/utility.h
***************
*** 28,42 **** typedef enum
  typedef void (*ProcessUtility_hook_type) (Node *parsetree,
  			  const char *queryString, ParamListInfo params,
  			  DestReceiver *dest, char *completionTag,
! 			  ProcessUtilityContext context);
  extern PGDLLIMPORT ProcessUtility_hook_type ProcessUtility_hook;
  
  extern void ProcessUtility(Node *parsetree, const char *queryString,
  			   ParamListInfo params, DestReceiver *dest, char *completionTag,
! 			   ProcessUtilityContext context);
  extern void standard_ProcessUtility(Node *parsetree, const char *queryString,
  						ParamListInfo params, DestReceiver *dest,
! 						char *completionTag, ProcessUtilityContext context);
  
  extern bool UtilityReturnsTuples(Node *parsetree);
  
--- 28,45 ----
  typedef void (*ProcessUtility_hook_type) (Node *parsetree,
  			  const char *queryString, ParamListInfo params,
  			  DestReceiver *dest, char *completionTag,
! 			  ProcessUtilityContext context,
! 			  uint64 *processed);
  extern PGDLLIMPORT ProcessUtility_hook_type ProcessUtility_hook;
  
  extern void ProcessUtility(Node *parsetree, const char *queryString,
  			   ParamListInfo params, DestReceiver *dest, char *completionTag,
! 			   ProcessUtilityContext context,
! 			   uint64 *processed);
  extern void standard_ProcessUtility(Node *parsetree, const char *queryString,
  						ParamListInfo params, DestReceiver *dest,
! 						char *completionTag, ProcessUtilityContext context,
! 						uint64 *processed);
  
  extern bool UtilityReturnsTuples(Node *parsetree);
  
*** a/src/include/utils/portal.h
--- b/src/include/utils/portal.h
***************
*** 131,136 **** typedef struct PortalData
--- 131,137 ----
  	const char *commandTag;		/* command tag for original query */
  	List	   *stmts;			/* PlannedStmts and/or utility statements */
  	CachedPlan *cplan;			/* CachedPlan, if stmts are from one */
+ 	uint64	   processed;			/* returned rows */
  
  	ParamListInfo portalParams; /* params to pass to query */
  
*** a/src/test/regress/input/copy.source
--- b/src/test/regress/input/copy.source
***************
*** 106,108 **** this is just a line full of junk that would error out if parsed
--- 106,121 ----
  \.
  
  copy copytest3 to stdout csv header;
+ 
+ do $$
+ declare r int;
+ begin
+   copy copytest2 to '@abs_builddir@/results/copytest.csv' csv;
+   get diagnostics r = row_count;
+   raise notice 'exported % rows', r;
+   truncate copytest2;
+   copy copytest2 from '@abs_builddir@/results/copytest.csv' csv;
+   get diagnostics r = row_count;
+   raise notice 'imported % rows', r;
+ end;
+ $$ language plpgsql;
*** a/src/test/regress/output/copy.source
--- b/src/test/regress/output/copy.source
***************
*** 71,73 **** copy copytest3 to stdout csv header;
--- 71,87 ----
  c1,"col with , comma","col with "" quote"
  1,a,1
  2,b,2
+ do $$
+ declare r int;
+ begin
+   copy copytest2 to '@abs_builddir@/results/copytest.csv' csv;
+   get diagnostics r = row_count;
+   raise notice 'exported % rows', r;
+   truncate copytest2;
+   copy copytest2 from '@abs_builddir@/results/copytest.csv' csv;
+   get diagnostics r = row_count;
+   raise notice 'imported % rows', r;
+ end;
+ $$ language plpgsql;
+ NOTICE:  exported 4 rows
+ NOTICE:  imported 4 rows