Hello

There is not possible to get a number of processed rows when COPY is
evaluated via SPI. Client can use a tag, but SPI doesn't use a tag.

I propose a small change a ProcessUtility to return a processed rows.

Note: I found a small inconsistency between SPI and Utility interface.
SPI still use a 4 byte unsign int for storing a number of processed
rows. Utility use a 8bytes unsign int.

Motivation:

 postgres=# \sf fx
CREATE OR REPLACE FUNCTION public.fx(tablename text, filename text)
 RETURNS integer
 LANGUAGE plpgsql
AS $function$
declare r int;
begin
  execute format('COPY %s FROM %s', quote_ident(tablename),
quote_literal(filename));
  get diagnostics r = row_count;
  return r;
end;
$function$

Regards

Pavel Stehule
diff --git a/src/backend/executor/functions.c b/src/backend/executor/functions.c
index 398bc40..a7c2b8f 100644
--- a/src/backend/executor/functions.c
+++ b/src/backend/executor/functions.c
@@ -600,6 +600,7 @@ postquel_getnext(execution_state *es, SQLFunctionCachePtr fcache)
 					   es->qd->params,
 					   false,	/* not top level */
 					   es->qd->dest,
+					   NULL,
 					   NULL);
 		result = true;			/* never stops early */
 	}
diff --git a/src/backend/executor/spi.c b/src/backend/executor/spi.c
index 688279c..21cabcc 100644
--- a/src/backend/executor/spi.c
+++ b/src/backend/executor/spi.c
@@ -1838,6 +1838,7 @@ _SPI_execute_plan(SPIPlanPtr plan, ParamListInfo paramLI,
 		{
 			Node	   *stmt = (Node *) lfirst(lc2);
 			bool		canSetTag;
+			bool		isCopyStmt = false;
 			DestReceiver *dest;
 
 			_SPI_current->processed = 0;
@@ -1857,6 +1858,7 @@ _SPI_execute_plan(SPIPlanPtr plan, ParamListInfo paramLI,
 				{
 					CopyStmt   *cstmt = (CopyStmt *) stmt;
 
+					isCopyStmt = true;
 					if (cstmt->filename == NULL)
 					{
 						my_res = SPI_ERROR_COPY;
@@ -1911,16 +1913,23 @@ _SPI_execute_plan(SPIPlanPtr plan, ParamListInfo paramLI,
 			}
 			else
 			{
+				uint32 processed;
+
 				ProcessUtility(stmt,
 							   plansource->query_string,
 							   paramLI,
 							   false,	/* not top level */
 							   dest,
-							   NULL);
+							   NULL,
+							   &processed);
 				/* Update "processed" if stmt returned tuples */
+
 				if (_SPI_current->tuptable)
 					_SPI_current->processed = _SPI_current->tuptable->alloced -
 						_SPI_current->tuptable->free;
+				else if (canSetTag && isCopyStmt)
+					_SPI_current->processed = processed;
+
 				res = SPI_OK_UTILITY;
 			}
 
diff --git a/src/backend/tcop/pquery.c b/src/backend/tcop/pquery.c
index 466727b..1a861ee 100644
--- a/src/backend/tcop/pquery.c
+++ b/src/backend/tcop/pquery.c
@@ -1184,7 +1184,8 @@ PortalRunUtility(Portal portal, Node *utilityStmt, bool isTopLevel,
 				   portal->portalParams,
 				   isTopLevel,
 				   dest,
-				   completionTag);
+				   completionTag,
+				   NULL);
 
 	/* Some utility statements may change context on us */
 	MemoryContextSwitchTo(PortalGetHeapMemory(portal));
diff --git a/src/backend/tcop/utility.c b/src/backend/tcop/utility.c
index 0749227..35db28c 100644
--- a/src/backend/tcop/utility.c
+++ b/src/backend/tcop/utility.c
@@ -319,6 +319,9 @@ CheckRestrictedOperation(const char *cmdname)
  * completionTag is only set nonempty if we want to return a nondefault status.
  *
  * completionTag may be NULL if caller doesn't want a status string.
+ *
+ * processed may be NULL if caller doesn't want a number of processed rows 
+ * by COPY statement
  */
 void
 ProcessUtility(Node *parsetree,
@@ -326,7 +329,8 @@ ProcessUtility(Node *parsetree,
 			   ParamListInfo params,
 			   bool isTopLevel,
 			   DestReceiver *dest,
-			   char *completionTag)
+			   char *completionTag,
+			   uint32 *processed)
 {
 	Assert(queryString != NULL);	/* required as of 8.4 */
 
@@ -337,10 +341,10 @@ ProcessUtility(Node *parsetree,
 	 */
 	if (ProcessUtility_hook)
 		(*ProcessUtility_hook) (parsetree, queryString, params,
-								isTopLevel, dest, completionTag);
+								isTopLevel, dest, completionTag, processed);
 	else
 		standard_ProcessUtility(parsetree, queryString, params,
-								isTopLevel, dest, completionTag);
+								isTopLevel, dest, completionTag, processed);
 }
 
 void
@@ -349,7 +353,8 @@ standard_ProcessUtility(Node *parsetree,
 						ParamListInfo params,
 						bool isTopLevel,
 						DestReceiver *dest,
-						char *completionTag)
+						char *completionTag,
+						uint32 *processed)
 {
 	check_xact_readonly(parsetree);
 
@@ -571,6 +576,7 @@ standard_ProcessUtility(Node *parsetree,
 									   params,
 									   false,
 									   None_Receiver,
+									   NULL,
 									   NULL);
 					}
 
@@ -716,12 +722,14 @@ standard_ProcessUtility(Node *parsetree,
 
 		case T_CopyStmt:
 			{
-				uint64		processed;
+				uint64		_processed;
 
-				processed = DoCopy((CopyStmt *) parsetree, queryString);
+				_processed = DoCopy((CopyStmt *) parsetree, queryString);
 				if (completionTag)
 					snprintf(completionTag, COMPLETION_TAG_BUFSIZE,
-							 "COPY " UINT64_FORMAT, processed);
+							 "COPY " UINT64_FORMAT, _processed);
+				if (processed != NULL)
+					*processed = (uint32) _processed;
 			}
 			break;
 
@@ -782,6 +790,7 @@ standard_ProcessUtility(Node *parsetree,
 									   params,
 									   false,
 									   None_Receiver,
+									   NULL,
 									   NULL);
 					}
 
diff --git a/src/include/tcop/utility.h b/src/include/tcop/utility.h
index c21857a..86fad4b 100644
--- a/src/include/tcop/utility.h
+++ b/src/include/tcop/utility.h
@@ -20,15 +20,16 @@
 /* Hook for plugins to get control in ProcessUtility() */
 typedef void (*ProcessUtility_hook_type) (Node *parsetree,
 			  const char *queryString, ParamListInfo params, bool isTopLevel,
-									DestReceiver *dest, char *completionTag);
+									DestReceiver *dest, char *completionTag, 
+													    uint32 *processed);
 extern PGDLLIMPORT ProcessUtility_hook_type ProcessUtility_hook;
 
 extern void ProcessUtility(Node *parsetree, const char *queryString,
 			   ParamListInfo params, bool isTopLevel,
-			   DestReceiver *dest, char *completionTag);
+			   DestReceiver *dest, char *completionTag, uint32 *processed);
 extern void standard_ProcessUtility(Node *parsetree, const char *queryString,
 						ParamListInfo params, bool isTopLevel,
-						DestReceiver *dest, char *completionTag);
+						DestReceiver *dest, char *completionTag, uint32 *processed);
 
 extern bool UtilityReturnsTuples(Node *parsetree);
 
-- 
Sent via pgsql-hackers mailing list (pgsql-hackers@postgresql.org)
To make changes to your subscription:
http://www.postgresql.org/mailpref/pgsql-hackers

Reply via email to