Add a hook to let loadable modules get control at ProcessUtility execution,

and use it to extend contrib/pg_stat_statements to track utility commands.

Itagaki Takahiro, reviewed by Euler Taveira de Oliveira.

contrib/pg_stat_statements/pg_stat_statements.c

@@ -14,7 +14,7 @@
  * Copyright (c) 2008-2009, PostgreSQL Global Development Group
  *
  * IDENTIFICATION
- *	  $PostgreSQL: pgsql/contrib/pg_stat_statements/pg_stat_statements.c,v 1.8 2009/12/15 04:57:46 rhaas Exp $
+ *	  $PostgreSQL: pgsql/contrib/pg_stat_statements/pg_stat_statements.c,v 1.9 2009/12/15 20:04:49 tgl Exp $
  *
  *-------------------------------------------------------------------------
  */
@@ -32,6 +32,7 @@
 #include "storage/fd.h"
 #include "storage/ipc.h"
 #include "storage/spin.h"
+#include "tcop/utility.h"
 #include "utils/builtins.h"
 #include "utils/hsearch.h"
 #include "utils/guc.h"
@@ -113,6 +114,7 @@ static shmem_startup_hook_type prev_shmem_startup_hook = NULL;
 static ExecutorStart_hook_type prev_ExecutorStart = NULL;
 static ExecutorRun_hook_type prev_ExecutorRun = NULL;
 static ExecutorEnd_hook_type prev_ExecutorEnd = NULL;
+static ProcessUtility_hook_type prev_ProcessUtility = NULL;
 
 /* Links to shared memory state */
 static pgssSharedState *pgss = NULL;
@@ -124,7 +126,7 @@ typedef enum
 {
 	PGSS_TRACK_NONE,			/* track no statements */
 	PGSS_TRACK_TOP,				/* only top level statements */
-	PGSS_TRACK_ALL,				/* all statements, including nested ones */
+	PGSS_TRACK_ALL				/* all statements, including nested ones */
 } PGSSTrackLevel;
 
 static const struct config_enum_entry track_options[] = {
@@ -136,6 +138,7 @@ static const struct config_enum_entry track_options[] = {
 
 static int	pgss_max;			/* max # statements to track */
 static int	pgss_track;			/* tracking level */
+static bool pgss_track_utility;	/* whether to track utility commands */
 static bool pgss_save;			/* whether to save stats across shutdown */
 
 
@@ -161,10 +164,12 @@ static void pgss_ExecutorRun(QueryDesc *queryDesc,
 				 ScanDirection direction,
 				 long count);
 static void pgss_ExecutorEnd(QueryDesc *queryDesc);
+static void pgss_ProcessUtility(Node *parsetree,
+			   const char *queryString, ParamListInfo params, bool isTopLevel,
+			   DestReceiver *dest, char *completionTag);
 static uint32 pgss_hash_fn(const void *key, Size keysize);
 static int	pgss_match_fn(const void *key1, const void *key2, Size keysize);
-static void pgss_store(const char *query,
-		   const Instrumentation *instr, uint32 rows);
+static void pgss_store(const char *query, double total_time, uint64 rows);
 static Size pgss_memsize(void);
 static pgssEntry *entry_alloc(pgssHashKey *key);
 static void entry_dealloc(void);
@@ -214,6 +219,16 @@ _PG_init(void)
 							 NULL,
 							 NULL);
 
+	DefineCustomBoolVariable("pg_stat_statements.track_utility",
+			   "Selects whether utility commands are tracked by pg_stat_statements.",
+							 NULL,
+							 &pgss_track_utility,
+							 true,
+							 PGC_SUSET,
+							 0,
+							 NULL,
+							 NULL);
+
 	DefineCustomBoolVariable("pg_stat_statements.save",
 			   "Save pg_stat_statements statistics across server shutdowns.",
 							 NULL,
@@ -245,6 +260,8 @@ _PG_init(void)
 	ExecutorRun_hook = pgss_ExecutorRun;
 	prev_ExecutorEnd = ExecutorEnd_hook;
 	ExecutorEnd_hook = pgss_ExecutorEnd;
+	prev_ProcessUtility = ProcessUtility_hook;
+	ProcessUtility_hook = pgss_ProcessUtility;
 }
 
 /*
@@ -254,10 +271,11 @@ void
 _PG_fini(void)
 {
 	/* Uninstall hooks. */
+	shmem_startup_hook = prev_shmem_startup_hook;
 	ExecutorStart_hook = prev_ExecutorStart;
 	ExecutorRun_hook = prev_ExecutorRun;
 	ExecutorEnd_hook = prev_ExecutorEnd;
-	shmem_startup_hook = prev_shmem_startup_hook;
+	ProcessUtility_hook = prev_ProcessUtility;
 }
 
 /*
@@ -539,7 +557,7 @@ pgss_ExecutorEnd(QueryDesc *queryDesc)
 		InstrEndLoop(queryDesc->totaltime);
 
 		pgss_store(queryDesc->sourceText,
-				   queryDesc->totaltime,
+				   queryDesc->totaltime->total,
 				   queryDesc->estate->es_processed);
 	}
 
@@ -549,6 +567,61 @@ pgss_ExecutorEnd(QueryDesc *queryDesc)
 		standard_ExecutorEnd(queryDesc);
 }
 
+/*
+ * ProcessUtility hook
+ */
+static void
+pgss_ProcessUtility(Node *parsetree, const char *queryString,
+					ParamListInfo params, bool isTopLevel,
+					DestReceiver *dest, char *completionTag)
+{
+	if (pgss_track_utility && pgss_enabled())
+	{
+		instr_time	start;
+		instr_time	duration;
+		uint64		rows = 0;
+
+		INSTR_TIME_SET_CURRENT(start);
+
+		nested_level++;
+		PG_TRY();
+		{
+			if (prev_ProcessUtility)
+				prev_ProcessUtility(parsetree, queryString, params,
+									isTopLevel, dest, completionTag);
+			else
+				standard_ProcessUtility(parsetree, queryString, params,
+										isTopLevel, dest, completionTag);
+			nested_level--;
+		}
+		PG_CATCH();
+		{
+			nested_level--;
+			PG_RE_THROW();
+		}
+		PG_END_TRY();
+
+		INSTR_TIME_SET_CURRENT(duration);
+		INSTR_TIME_SUBTRACT(duration, start);
+
+		/* parse command tag to retrieve the number of affected rows. */
+		if (completionTag &&
+			sscanf(completionTag, "COPY " UINT64_FORMAT, &rows) != 1)
+			rows = 0;
+
+		pgss_store(queryString, INSTR_TIME_GET_DOUBLE(duration), rows);
+	}
+	else
+	{
+		if (prev_ProcessUtility)
+			prev_ProcessUtility(parsetree, queryString, params,
+								isTopLevel, dest, completionTag);
+		else
+			standard_ProcessUtility(parsetree, queryString, params,
+									isTopLevel, dest, completionTag);
+	}
+}
+
 /*
  * Calculate hash value for a key
  */
@@ -587,7 +660,7 @@ pgss_match_fn(const void *key1, const void *key2, Size keysize)
  * Store some statistics for a statement.
  */
 static void
-pgss_store(const char *query, const Instrumentation *instr, uint32 rows)
+pgss_store(const char *query, double total_time, uint64 rows)
 {
 	pgssHashKey key;
 	double		usage;
@@ -631,7 +704,7 @@ pgss_store(const char *query, const Instrumentation *instr, uint32 rows)
 
 		SpinLockAcquire(&e->mutex);
 		e->counters.calls += 1;
-		e->counters.total_time += instr->total;
+		e->counters.total_time += total_time;
 		e->counters.rows += rows;
 		e->counters.usage += usage;
 		SpinLockRelease(&e->mutex);

doc/src/sgml/pgstatstatements.sgml

-<!-- $PostgreSQL: pgsql/doc/src/sgml/pgstatstatements.sgml,v 1.4 2009/12/01 02:31:11 momjian Exp $ -->
+<!-- $PostgreSQL: pgsql/doc/src/sgml/pgstatstatements.sgml,v 1.5 2009/12/15 20:04:49 tgl Exp $ -->
 
 <sect1 id="pgstatstatements">
  <title>pg_stat_statements</title>
@@ -174,6 +174,23 @@
     </listitem>
    </varlistentry>
 
+   <varlistentry>
+    <term>
+     <varname>pg_stat_statements.track_utility</varname> (<type>boolean</type>)
+    </term>
+
+    <listitem>
+     <para>
+      <varname>pg_stat_statements.track_utility</varname> controls whether
+      utility commands are tracked by the module.  Utility commands are
+      all those other than <command>SELECT</>, <command>INSERT</>,
+      <command>UPDATE</> and <command>DELETE</>.
+      The default value is <literal>on</>.
+      Only superusers can change this setting.
+     </para>
+    </listitem>
+   </varlistentry>
+
    <varlistentry>
     <term>
      <varname>pg_stat_statements.save</varname> (<type>boolean</type>)

src/backend/tcop/utility.c

@@ -10,7 +10,7 @@
  *
  *
  * IDENTIFICATION
- *	  $PostgreSQL: pgsql/src/backend/tcop/utility.c,v 1.323 2009/12/11 03:34:55 itagaki Exp $
+ *	  $PostgreSQL: pgsql/src/backend/tcop/utility.c,v 1.324 2009/12/15 20:04:49 tgl Exp $
  *
  *-------------------------------------------------------------------------
  */
@@ -58,6 +58,10 @@
 #include "utils/syscache.h"
 
 
+/* Hook for plugins to get control in ProcessUtility() */
+ProcessUtility_hook_type ProcessUtility_hook = NULL;
+
+
 /*
  * Verify user has ownership of specified relation, else ereport.
  *
@@ -274,6 +278,27 @@ ProcessUtility(Node *parsetree,
 {
 	Assert(queryString != NULL);	/* required as of 8.4 */
 
+	/*
+	 * We provide a function hook variable that lets loadable plugins
+	 * get control when ProcessUtility is called.  Such a plugin would
+	 * normally call standard_ProcessUtility().
+	 */
+	if (ProcessUtility_hook)
+		(*ProcessUtility_hook) (parsetree, queryString, params,
+								isTopLevel, dest, completionTag);
+	else
+		standard_ProcessUtility(parsetree, queryString, params,
+								isTopLevel, dest, completionTag);
+}
+
+void
+standard_ProcessUtility(Node *parsetree,
+						const char *queryString,
+						ParamListInfo params,
+						bool isTopLevel,
+						DestReceiver *dest,
+						char *completionTag)
+{
 	check_xact_readonly(parsetree);
 
 	if (completionTag)

src/include/tcop/utility.h

@@ -7,7 +7,7 @@
  * Portions Copyright (c) 1996-2009, PostgreSQL Global Development Group
  * Portions Copyright (c) 1994, Regents of the University of California
  *
- * $PostgreSQL: pgsql/src/include/tcop/utility.h,v 1.37 2009/12/01 02:31:13 momjian Exp $
+ * $PostgreSQL: pgsql/src/include/tcop/utility.h,v 1.38 2009/12/15 20:04:49 tgl Exp $
  *
  *-------------------------------------------------------------------------
  */
@@ -17,9 +17,18 @@
 #include "tcop/tcopprot.h"
 
 
+/* 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);
+extern PGDLLIMPORT ProcessUtility_hook_type ProcessUtility_hook;
+
 extern void ProcessUtility(Node *parsetree, const char *queryString,
 			   ParamListInfo params, bool isTopLevel,
 			   DestReceiver *dest, char *completionTag);
+extern void standard_ProcessUtility(Node *parsetree, const char *queryString,
+			   ParamListInfo params, bool isTopLevel,
+			   DestReceiver *dest, char *completionTag);
 
 extern bool UtilityReturnsTuples(Node *parsetree);