From 14d806a9e0b5783d49965e308c1fcb2ac75c624b Mon Sep 17 00:00:00 2001
From: Hari Babu <kommi.haribabu@gmail.com>
Date: Thu, 15 Nov 2018 07:11:05 +1100
Subject: [PATCH] pg_stat_statements_reset to reset specific query/user/db
 statistics

Now the pg_stat_statements_reset() can accept input parameters of
userid, dbid and queryid. Corresponding to the specified input values,
specific statistics entries are released. The new behaviour is useful
to get the fresh new statistics of a specific query/user/database
to observation without resetting all the existing statistics.
If no parameters are passed, it will release all entries as per the
old behaviour.
---
 contrib/pg_stat_statements/Makefile           |   8 +-
 .../expected/pg_stat_statements.out           | 148 ++++++++++++++++++
 .../pg_stat_statements--1.6--1.7.sql          |  22 +++
 .../pg_stat_statements/pg_stat_statements.c   |  96 ++++++++++--
 .../pg_stat_statements.control                |   2 +-
 .../sql/pg_stat_statements.sql                |  53 +++++++
 doc/src/sgml/pgstatstatements.sgml            |  15 +-
 7 files changed, 323 insertions(+), 21 deletions(-)
 create mode 100644 contrib/pg_stat_statements/pg_stat_statements--1.6--1.7.sql

diff --git a/contrib/pg_stat_statements/Makefile b/contrib/pg_stat_statements/Makefile
index 14a50380dc..051ce46f0c 100644
--- a/contrib/pg_stat_statements/Makefile
+++ b/contrib/pg_stat_statements/Makefile
@@ -4,10 +4,10 @@ MODULE_big = pg_stat_statements
 OBJS = pg_stat_statements.o $(WIN32RES)
 
 EXTENSION = pg_stat_statements
-DATA = pg_stat_statements--1.4.sql pg_stat_statements--1.5--1.6.sql \
-	pg_stat_statements--1.4--1.5.sql pg_stat_statements--1.3--1.4.sql \
-	pg_stat_statements--1.2--1.3.sql pg_stat_statements--1.1--1.2.sql \
-	pg_stat_statements--1.0--1.1.sql \
+DATA = pg_stat_statements--1.4.sql pg_stat_statements--1.6--1.7.sql \
+	pg_stat_statements--1.5--1.6.sql pg_stat_statements--1.4--1.5.sql \
+	pg_stat_statements--1.3--1.4.sql pg_stat_statements--1.2--1.3.sql \
+	pg_stat_statements--1.1--1.2.sql pg_stat_statements--1.0--1.1.sql \
 	pg_stat_statements--unpackaged--1.0.sql
 PGFILEDESC = "pg_stat_statements - execution statistics of SQL statements"
 
diff --git a/contrib/pg_stat_statements/expected/pg_stat_statements.out b/contrib/pg_stat_statements/expected/pg_stat_statements.out
index 5318c3550c..40e2a4230a 100644
--- a/contrib/pg_stat_statements/expected/pg_stat_statements.out
+++ b/contrib/pg_stat_statements/expected/pg_stat_statements.out
@@ -395,4 +395,152 @@ SELECT query, calls, rows FROM pg_stat_statements ORDER BY query COLLATE "C";
  SELECT pg_stat_statements_reset()         |     1 |    1
 (8 rows)
 
+--
+-- Track user activity and reset them
+--
+SELECT pg_stat_statements_reset();
+ pg_stat_statements_reset 
+--------------------------
+ 
+(1 row)
+
+CREATE ROLE regress_stats_user1;
+CREATE ROLE regress_stats_user2;
+SET ROLE regress_stats_user1;
+SELECT 1 AS "ONE";
+ ONE 
+-----
+   1
+(1 row)
+
+SELECT 1+1 AS "TWO";
+ TWO 
+-----
+   2
+(1 row)
+
+RESET ROLE;
+SET ROLE regress_stats_user2;
+SELECT 1 AS "ONE";
+ ONE 
+-----
+   1
+(1 row)
+
+SELECT 1+1 AS "TWO";
+ TWO 
+-----
+   2
+(1 row)
+
+RESET ROLE;
+SELECT query, calls, rows FROM pg_stat_statements ORDER BY query COLLATE "C";
+               query               | calls | rows 
+-----------------------------------+-------+------
+ CREATE ROLE regress_stats_user1   |     1 |    0
+ CREATE ROLE regress_stats_user2   |     1 |    0
+ RESET ROLE                        |     2 |    0
+ SELECT $1 AS "ONE"                |     1 |    1
+ SELECT $1 AS "ONE"                |     1 |    1
+ SELECT $1+$2 AS "TWO"             |     1 |    1
+ SELECT $1+$2 AS "TWO"             |     1 |    1
+ SELECT pg_stat_statements_reset() |     1 |    1
+ SET ROLE regress_stats_user1      |     1 |    0
+ SET ROLE regress_stats_user2      |     1 |    0
+(10 rows)
+
+--
+-- remove query ('SELECT $1 AS "ONE"') executed by two users
+--
+SELECT pg_stat_statements_reset(NULL,NULL,s.queryid) FROM pg_stat_statements AS s WHERE s.query = 'SELECT $1 AS "ONE"';
+ pg_stat_statements_reset 
+--------------------------
+ 
+ 
+(2 rows)
+
+SELECT query, calls, rows FROM pg_stat_statements ORDER BY query COLLATE "C";
+                                              query                                               | calls | rows 
+--------------------------------------------------------------------------------------------------+-------+------
+ CREATE ROLE regress_stats_user1                                                                  |     1 |    0
+ CREATE ROLE regress_stats_user2                                                                  |     1 |    0
+ RESET ROLE                                                                                       |     2 |    0
+ SELECT $1+$2 AS "TWO"                                                                            |     1 |    1
+ SELECT $1+$2 AS "TWO"                                                                            |     1 |    1
+ SELECT pg_stat_statements_reset($1,$2,s.queryid) FROM pg_stat_statements AS s WHERE s.query = $3 |     1 |    2
+ SELECT pg_stat_statements_reset()                                                                |     1 |    1
+ SELECT query, calls, rows FROM pg_stat_statements ORDER BY query COLLATE "C"                     |     1 |   10
+ SET ROLE regress_stats_user1                                                                     |     1 |    0
+ SET ROLE regress_stats_user2                                                                     |     1 |    0
+(10 rows)
+
+--
+-- remove query of a user (regress_stats_user1)
+--
+SELECT pg_stat_statements_reset(r.oid) FROM pg_roles AS r WHERE r.rolname = 'regress_stats_user1';
+ pg_stat_statements_reset 
+--------------------------
+ 
+(1 row)
+
+SELECT query, calls, rows FROM pg_stat_statements ORDER BY query COLLATE "C";
+                                              query                                               | calls | rows 
+--------------------------------------------------------------------------------------------------+-------+------
+ CREATE ROLE regress_stats_user1                                                                  |     1 |    0
+ CREATE ROLE regress_stats_user2                                                                  |     1 |    0
+ RESET ROLE                                                                                       |     2 |    0
+ SELECT $1+$2 AS "TWO"                                                                            |     1 |    1
+ SELECT pg_stat_statements_reset($1,$2,s.queryid) FROM pg_stat_statements AS s WHERE s.query = $3 |     1 |    2
+ SELECT pg_stat_statements_reset()                                                                |     1 |    1
+ SELECT pg_stat_statements_reset(r.oid) FROM pg_roles AS r WHERE r.rolname = $1                   |     1 |    1
+ SELECT query, calls, rows FROM pg_stat_statements ORDER BY query COLLATE "C"                     |     2 |   20
+ SET ROLE regress_stats_user2                                                                     |     1 |    0
+(9 rows)
+
+--
+-- remove query ('SELECT $1 + $2 AS "TWO"') executed by regress_stats_user2 in the current_database
+--
+SELECT pg_stat_statements_reset((SELECT r.oid FROM pg_roles AS r WHERE r.rolname = 'regress_stats_user2'),
+								(SELECT d.oid from pg_database As d where datname = current_database()),
+								(SELECT s.queryid FROM pg_stat_statements AS s WHERE s.query = 'SELECT $1 AS "ONE"'));
+ pg_stat_statements_reset 
+--------------------------
+ 
+(1 row)
+
+SELECT query, calls, rows FROM pg_stat_statements ORDER BY query COLLATE "C";
+                                                                  query                                                                   | calls | rows 
+------------------------------------------------------------------------------------------------------------------------------------------+-------+------
+ CREATE ROLE regress_stats_user1                                                                                                          |     1 |    0
+ CREATE ROLE regress_stats_user2                                                                                                          |     1 |    0
+ RESET ROLE                                                                                                                               |     2 |    0
+ SELECT pg_stat_statements_reset($1,$2,s.queryid) FROM pg_stat_statements AS s WHERE s.query = $3                                         |     1 |    2
+ SELECT pg_stat_statements_reset((SELECT r.oid FROM pg_roles AS r WHERE r.rolname = $1),                                                 +|     1 |    1
+                                                                 (SELECT d.oid from pg_database As d where datname = current_database()),+|       | 
+                                                                 (SELECT s.queryid FROM pg_stat_statements AS s WHERE s.query = $2))      |       | 
+ SELECT pg_stat_statements_reset()                                                                                                        |     1 |    1
+ SELECT pg_stat_statements_reset(r.oid) FROM pg_roles AS r WHERE r.rolname = $1                                                           |     1 |    1
+ SELECT query, calls, rows FROM pg_stat_statements ORDER BY query COLLATE "C"                                                             |     3 |   29
+(8 rows)
+
+--
+-- reset all
+--
+SELECT pg_stat_statements_reset();
+ pg_stat_statements_reset 
+--------------------------
+ 
+(1 row)
+
+SELECT query, calls, rows FROM pg_stat_statements ORDER BY query COLLATE "C";
+               query               | calls | rows 
+-----------------------------------+-------+------
+ SELECT pg_stat_statements_reset() |     1 |    1
+(1 row)
+
+--
+-- cleanup
+--
+DROP ROLE regress_stats_user1;
+DROP ROLE regress_stats_user2;
 DROP EXTENSION pg_stat_statements;
diff --git a/contrib/pg_stat_statements/pg_stat_statements--1.6--1.7.sql b/contrib/pg_stat_statements/pg_stat_statements--1.6--1.7.sql
new file mode 100644
index 0000000000..c862de20a6
--- /dev/null
+++ b/contrib/pg_stat_statements/pg_stat_statements--1.6--1.7.sql
@@ -0,0 +1,22 @@
+/* contrib/pg_stat_statements/pg_stat_statements--1.6--1.7.sql */
+
+-- complain if script is sourced in psql, rather than via ALTER EXTENSION
+\echo Use "ALTER EXTENSION pg_stat_statements UPDATE TO '1.7'" to load this file. \quit
+
+/* First we have to remove them from the extension */
+ALTER EXTENSION pg_stat_statements DROP FUNCTION pg_stat_statements_reset();
+
+/* Then we can drop them */
+DROP FUNCTION pg_stat_statements_reset();
+
+/* Now redefine */
+CREATE FUNCTION pg_stat_statements_reset(IN userid Oid DEFAULT NULL,
+	IN dbid Oid DEFAULT NULL,
+	IN queryid bigint DEFAULT NULL
+)
+RETURNS void
+AS 'MODULE_PATHNAME', 'pg_stat_statements_reset_1_7'
+LANGUAGE C PARALLEL SAFE;
+
+-- Don't want this to be available to non-superusers.
+REVOKE ALL ON FUNCTION pg_stat_statements_reset(Oid, Oid, bigint) FROM PUBLIC;
diff --git a/contrib/pg_stat_statements/pg_stat_statements.c b/contrib/pg_stat_statements/pg_stat_statements.c
index 33f9a79f54..578e1571b6 100644
--- a/contrib/pg_stat_statements/pg_stat_statements.c
+++ b/contrib/pg_stat_statements/pg_stat_statements.c
@@ -290,6 +290,7 @@ void		_PG_init(void);
 void		_PG_fini(void);
 
 PG_FUNCTION_INFO_V1(pg_stat_statements_reset);
+PG_FUNCTION_INFO_V1(pg_stat_statements_reset_1_7);
 PG_FUNCTION_INFO_V1(pg_stat_statements_1_2);
 PG_FUNCTION_INFO_V1(pg_stat_statements_1_3);
 PG_FUNCTION_INFO_V1(pg_stat_statements);
@@ -327,7 +328,7 @@ static char *qtext_fetch(Size query_offset, int query_len,
 			char *buffer, Size buffer_size);
 static bool need_gc_qtexts(void);
 static void gc_qtexts(void);
-static void entry_reset(void);
+static void entry_reset(Oid userid, Oid dbid, uint64 queryid);
 static void AppendJumble(pgssJumbleState *jstate,
 			 const unsigned char *item, Size size);
 static void JumbleQuery(pgssJumbleState *jstate, Query *query);
@@ -1148,8 +1149,17 @@ pgss_store(const char *query, uint64 queryId,
 	 * For utility statements, we just hash the query string to get an ID.
 	 */
 	if (queryId == UINT64CONST(0))
+	{
 		queryId = pgss_hash_string(query, query_len);
 
+		/*
+		 * If we are unlucky enough to get a hash of zero(invalid), use queryID
+		 * as 2 instead, queryID 1 is already in use for normal statements.
+		 */
+		if (queryId == UINT64CONST(0))
+			queryId = UINT64CONST(2);
+	}
+
 	/* Set up key for hashtable search */
 	key.userid = GetUserId();
 	key.dbid = MyDatabaseId;
@@ -1293,16 +1303,32 @@ done:
 }
 
 /*
- * Reset all statement statistics.
+ * Reset statement statistics corresponding to userid, dbid, and queryid.
+ */
+Datum
+pg_stat_statements_reset_1_7(PG_FUNCTION_ARGS)
+{
+	Oid			userid;
+	Oid			dbid;
+	uint64		queryid;
+
+	userid = PG_ARGISNULL(0) ? 0 : PG_GETARG_OID(0);
+	dbid = PG_ARGISNULL(1) ? 0 : PG_GETARG_OID(1);
+	queryid = (uint64) (PG_ARGISNULL(2) ? 0 : PG_GETARG_INT64(2));
+
+	entry_reset(userid, dbid, queryid);
+
+	PG_RETURN_VOID();
+}
+
+/*
+ * Reset statement statistics.
  */
 Datum
 pg_stat_statements_reset(PG_FUNCTION_ARGS)
 {
-	if (!pgss || !pgss_hash)
-		ereport(ERROR,
-				(errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
-				 errmsg("pg_stat_statements must be loaded via shared_preload_libraries")));
-	entry_reset();
+	entry_reset(0, 0, 0);
+
 	PG_RETURN_VOID();
 }
 
@@ -2229,23 +2255,68 @@ gc_fail:
 }
 
 /*
- * Release all entries.
+ * Release entries corresponding to parameters passed.
  */
 static void
-entry_reset(void)
+entry_reset(Oid userid, Oid dbid, uint64 queryid)
 {
 	HASH_SEQ_STATUS hash_seq;
 	pgssEntry  *entry;
 	FILE	   *qfile;
+	long		num_entries;
+	long		num_remove = 0;
+	pgssHashKey key;
+
+	if (!pgss || !pgss_hash)
+		ereport(ERROR,
+				(errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
+				 errmsg("pg_stat_statements must be loaded via shared_preload_libraries")));
 
 	LWLockAcquire(pgss->lock, LW_EXCLUSIVE);
+	num_entries = hash_get_num_entries(pgss_hash);
 
-	hash_seq_init(&hash_seq, pgss_hash);
-	while ((entry = hash_seq_search(&hash_seq)) != NULL)
+	if (userid != 0 && dbid != 0 && queryid != UINT64CONST(0))
+	{
+		/* If all the parameters are available, use the fast path. */
+		key.userid = userid;
+		key.dbid = dbid;
+		key.queryid = queryid;
+
+		/* Remove the key if exists */
+		entry = (pgssEntry *) hash_search(pgss_hash, &key, HASH_REMOVE, NULL);
+		if (entry)				/* found */
+			num_remove++;
+	}
+	else if (userid != 0 || dbid != 0 || queryid != UINT64CONST(0))
 	{
-		hash_search(pgss_hash, &entry->key, HASH_REMOVE, NULL);
+		/* Remove entries corresponding to valid parameters. */
+		hash_seq_init(&hash_seq, pgss_hash);
+		while ((entry = hash_seq_search(&hash_seq)) != NULL)
+		{
+			if ((!userid || entry->key.userid == userid) &&
+				(!dbid || entry->key.dbid == dbid) &&
+				(!queryid || entry->key.queryid == queryid))
+			{
+				hash_search(pgss_hash, &entry->key, HASH_REMOVE, NULL);
+				num_remove++;
+			}
+		}
+	}
+	else
+	{
+		/* Remove all entries. */
+		hash_seq_init(&hash_seq, pgss_hash);
+		while ((entry = hash_seq_search(&hash_seq)) != NULL)
+		{
+			hash_search(pgss_hash, &entry->key, HASH_REMOVE, NULL);
+			num_remove++;
+		}
 	}
 
+	/* All entries are removed? */
+	if (num_entries != num_remove)
+		goto release_lock;
+
 	/*
 	 * Write new empty query file, perhaps even creating a new one to recover
 	 * if the file was missing.
@@ -2274,6 +2345,7 @@ done:
 	/* This counts as a query text garbage collection for our purposes */
 	record_gc_qtexts();
 
+release_lock:
 	LWLockRelease(pgss->lock);
 }
 
diff --git a/contrib/pg_stat_statements/pg_stat_statements.control b/contrib/pg_stat_statements/pg_stat_statements.control
index 617038b4c0..14cb422354 100644
--- a/contrib/pg_stat_statements/pg_stat_statements.control
+++ b/contrib/pg_stat_statements/pg_stat_statements.control
@@ -1,5 +1,5 @@
 # pg_stat_statements extension
 comment = 'track execution statistics of all SQL statements executed'
-default_version = '1.6'
+default_version = '1.7'
 module_pathname = '$libdir/pg_stat_statements'
 relocatable = true
diff --git a/contrib/pg_stat_statements/sql/pg_stat_statements.sql b/contrib/pg_stat_statements/sql/pg_stat_statements.sql
index a8361fd1bf..d981c69a96 100644
--- a/contrib/pg_stat_statements/sql/pg_stat_statements.sql
+++ b/contrib/pg_stat_statements/sql/pg_stat_statements.sql
@@ -195,4 +195,57 @@ DROP FUNCTION PLUS_TWO(INTEGER);
 
 SELECT query, calls, rows FROM pg_stat_statements ORDER BY query COLLATE "C";
 
+--
+-- Track user activity and reset them
+--
+SELECT pg_stat_statements_reset();
+CREATE ROLE regress_stats_user1;
+CREATE ROLE regress_stats_user2;
+
+SET ROLE regress_stats_user1;
+
+SELECT 1 AS "ONE";
+SELECT 1+1 AS "TWO";
+
+RESET ROLE;
+SET ROLE regress_stats_user2;
+
+SELECT 1 AS "ONE";
+SELECT 1+1 AS "TWO";
+
+RESET ROLE;
+SELECT query, calls, rows FROM pg_stat_statements ORDER BY query COLLATE "C";
+
+--
+-- remove query ('SELECT $1 AS "ONE"') executed by two users
+--
+SELECT pg_stat_statements_reset(NULL,NULL,s.queryid) FROM pg_stat_statements AS s WHERE s.query = 'SELECT $1 AS "ONE"';
+SELECT query, calls, rows FROM pg_stat_statements ORDER BY query COLLATE "C";
+
+--
+-- remove query of a user (regress_stats_user1)
+--
+SELECT pg_stat_statements_reset(r.oid) FROM pg_roles AS r WHERE r.rolname = 'regress_stats_user1';
+SELECT query, calls, rows FROM pg_stat_statements ORDER BY query COLLATE "C";
+
+--
+-- remove query ('SELECT $1 + $2 AS "TWO"') executed by regress_stats_user2 in the current_database
+--
+SELECT pg_stat_statements_reset((SELECT r.oid FROM pg_roles AS r WHERE r.rolname = 'regress_stats_user2'),
+								(SELECT d.oid from pg_database As d where datname = current_database()),
+								(SELECT s.queryid FROM pg_stat_statements AS s WHERE s.query = 'SELECT $1 AS "ONE"'));
+SELECT query, calls, rows FROM pg_stat_statements ORDER BY query COLLATE "C";
+
+--
+-- reset all
+--
+SELECT pg_stat_statements_reset();
+SELECT query, calls, rows FROM pg_stat_statements ORDER BY query COLLATE "C";
+
+--
+-- cleanup
+--
+DROP ROLE regress_stats_user1;
+DROP ROLE regress_stats_user2;
+
 DROP EXTENSION pg_stat_statements;
diff --git a/doc/src/sgml/pgstatstatements.sgml b/doc/src/sgml/pgstatstatements.sgml
index c0217ed485..e9a9dafbf8 100644
--- a/doc/src/sgml/pgstatstatements.sgml
+++ b/doc/src/sgml/pgstatstatements.sgml
@@ -336,7 +336,7 @@
   <variablelist>
    <varlistentry>
     <term>
-     <function>pg_stat_statements_reset() returns void</function>
+     <function>pg_stat_statements_reset(userid Oid, dbid Oid, queryid bigint) returns void</function>
      <indexterm>
       <primary>pg_stat_statements_reset</primary>
      </indexterm>
@@ -344,9 +344,16 @@
 
     <listitem>
      <para>
-      <function>pg_stat_statements_reset</function> discards all statistics
-      gathered so far by <filename>pg_stat_statements</filename>.
-      By default, this function can only be executed by superusers.
+      <function>pg_stat_statements_reset</function> discards statistics
+      gathered so far by <filename>pg_stat_statements</filename> corresponding
+      to the specified <structfield>userid</structfield>, <structfield>dbid</structfield>
+      and <structfield>queryid</structfield>.  If any of the parameters are not
+      specified, the default value <literal>NULL</literal>(invalid) is used for
+      each of them and the statistics that match with other parameters will be
+      reset.  If no parameter is specified or all the specified parameters are
+      <literal>NULL</literal>(invalid), it will discard all statistics.  By
+      default, this function can only be executed by superusers.  Access may be
+      granted to others using <command>GRANT</command>.
      </para>
     </listitem>
    </varlistentry>
-- 
2.18.0.windows.1

