From 9461806a8b223293114382d2bc4d137735774a1f Mon Sep 17 00:00:00 2001
From: Hari Babu <kommi.haribabu@gmail.com>
Date: Wed, 20 Jun 2018 16:12:17 +1000
Subject: [PATCH] Function to reset statistics of a specific statement

New function pg_stat_statements_reset_query() is added to
support reset of a specific query instead of resetting all the
query statistics. This is useful to get the fresh new statistics
of a specific query to observation.
---
 contrib/pg_stat_statements/Makefile                |  8 ++--
 .../expected/pg_stat_statements.out                | 47 ++++++++++++++++++++++
 .../pg_stat_statements--1.5--1.6.sql               | 14 +++++++
 contrib/pg_stat_statements/pg_stat_statements.c    | 31 ++++++++++++++
 .../pg_stat_statements/pg_stat_statements.control  |  2 +-
 .../pg_stat_statements/sql/pg_stat_statements.sql  | 14 +++++++
 6 files changed, 111 insertions(+), 5 deletions(-)
 create mode 100644 contrib/pg_stat_statements/pg_stat_statements--1.5--1.6.sql

diff --git a/contrib/pg_stat_statements/Makefile b/contrib/pg_stat_statements/Makefile
index 39b368b70e..f0942f829e 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.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
+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 pg_stat_statements--unpackaged--1.0.sql
 PGFILEDESC = "pg_stat_statements - execution statistics of SQL statements"
 
 LDFLAGS_SL += $(filter -lm, $(LIBS))
diff --git a/contrib/pg_stat_statements/expected/pg_stat_statements.out b/contrib/pg_stat_statements/expected/pg_stat_statements.out
index 5318c3550c..0904020bd8 100644
--- a/contrib/pg_stat_statements/expected/pg_stat_statements.out
+++ b/contrib/pg_stat_statements/expected/pg_stat_statements.out
@@ -395,4 +395,51 @@ SELECT query, calls, rows FROM pg_stat_statements ORDER BY query COLLATE "C";
  SELECT pg_stat_statements_reset()         |     1 |    1
 (8 rows)
 
+--
+-- reset statistics of a specific query 'DROP TABLE test'
+--
+SELECT pg_stat_statements_reset_query(s.queryid) FROM pg_stat_statements AS s WHERE s.query = 'DROP TABLE test';
+ pg_stat_statements_reset_query 
+--------------------------------
+ 
+(1 row)
+
+SELECT query, calls, rows FROM pg_stat_statements ORDER BY query COLLATE "C";
+                                              query                                               | calls | rows 
+--------------------------------------------------------------------------------------------------+-------+------
+ CREATE INDEX test_b ON test(b)                                                                   |     1 |    0
+ DROP FUNCTION IF EXISTS PLUS_ONE(INTEGER)                                                        |     1 |    0
+ DROP FUNCTION PLUS_ONE(INTEGER)                                                                  |     1 |    0
+ DROP FUNCTION PLUS_TWO(INTEGER)                                                                  |     1 |    0
+ DROP TABLE IF EXISTS test                                                                        |     3 |    0
+ SELECT $1                                                                                        |     1 |    1
+ SELECT pg_stat_statements_reset()                                                                |     1 |    1
+ SELECT pg_stat_statements_reset_query(s.queryid) FROM pg_stat_statements AS s WHERE s.query = $1 |     1 |    1
+ SELECT query, calls, rows FROM pg_stat_statements ORDER BY query COLLATE "C"                     |     1 |    8
+(9 rows)
+
+--
+-- try to remove unknown queryid
+--
+SELECT pg_stat_statements_reset_query(0);
+ pg_stat_statements_reset_query 
+--------------------------------
+ 
+(1 row)
+
+SELECT query, calls, rows FROM pg_stat_statements ORDER BY query COLLATE "C";
+                                              query                                               | calls | rows 
+--------------------------------------------------------------------------------------------------+-------+------
+ CREATE INDEX test_b ON test(b)                                                                   |     1 |    0
+ DROP FUNCTION IF EXISTS PLUS_ONE(INTEGER)                                                        |     1 |    0
+ DROP FUNCTION PLUS_ONE(INTEGER)                                                                  |     1 |    0
+ DROP FUNCTION PLUS_TWO(INTEGER)                                                                  |     1 |    0
+ DROP TABLE IF EXISTS test                                                                        |     3 |    0
+ SELECT $1                                                                                        |     1 |    1
+ SELECT pg_stat_statements_reset()                                                                |     1 |    1
+ SELECT pg_stat_statements_reset_query($1)                                                        |     1 |    1
+ SELECT pg_stat_statements_reset_query(s.queryid) FROM pg_stat_statements AS s WHERE s.query = $1 |     1 |    1
+ SELECT query, calls, rows FROM pg_stat_statements ORDER BY query COLLATE "C"                     |     2 |   17
+(10 rows)
+
 DROP EXTENSION pg_stat_statements;
diff --git a/contrib/pg_stat_statements/pg_stat_statements--1.5--1.6.sql b/contrib/pg_stat_statements/pg_stat_statements--1.5--1.6.sql
new file mode 100644
index 0000000000..c5fbd47627
--- /dev/null
+++ b/contrib/pg_stat_statements/pg_stat_statements--1.5--1.6.sql
@@ -0,0 +1,14 @@
+/* contrib/pg_stat_statements/pg_stat_statements--1.5--1.6.sql */
+
+-- complain if script is sourced in psql, rather than via ALTER EXTENSION
+\echo Use "ALTER EXTENSION pg_stat_statements UPDATE TO '1.6'" to load this file. \quit
+
+-- Register functions.
+CREATE FUNCTION pg_stat_statements_reset_query(IN queryid bigint)
+RETURNS void
+AS 'MODULE_PATHNAME'
+LANGUAGE C PARALLEL SAFE;
+
+-- Don't want this to be available to non-superusers.
+REVOKE ALL ON FUNCTION pg_stat_statements_reset_query(bigint) FROM PUBLIC;
+GRANT EXECUTE ON FUNCTION pg_stat_statements_reset_query(bigint) TO pg_read_all_stats;
diff --git a/contrib/pg_stat_statements/pg_stat_statements.c b/contrib/pg_stat_statements/pg_stat_statements.c
index cc9efab243..9b24bc200f 100644
--- a/contrib/pg_stat_statements/pg_stat_statements.c
+++ b/contrib/pg_stat_statements/pg_stat_statements.c
@@ -293,6 +293,7 @@ PG_FUNCTION_INFO_V1(pg_stat_statements_reset);
 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);
+PG_FUNCTION_INFO_V1(pg_stat_statements_reset_query);
 
 static void pgss_shmem_startup(void);
 static void pgss_shmem_shutdown(int code, Datum arg);
@@ -1306,6 +1307,36 @@ pg_stat_statements_reset(PG_FUNCTION_ARGS)
 	PG_RETURN_VOID();
 }
 
+/*
+ * Reset specific statement statistics.
+ */
+Datum
+pg_stat_statements_reset_query(PG_FUNCTION_ARGS)
+{
+	pgssHashKey key;
+	pgssEntry  *entry;
+
+	if (!pgss || !pgss_hash)
+		ereport(ERROR,
+				(errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
+				 errmsg("pg_stat_statements must be loaded via shared_preload_libraries")));
+
+	/* Set up key for hashtable search */
+	key.userid = GetUserId();
+	key.dbid = MyDatabaseId;
+	key.queryid = PG_GETARG_INT64(0);
+
+	/* Remove the key out of hash if available */
+	LWLockAcquire(pgss->lock, LW_EXCLUSIVE);
+	entry = (pgssEntry *) hash_search(pgss_hash, &key, HASH_REMOVE, NULL);
+
+	if (!entry || key.queryid != entry->key.queryid)	/* not found */
+		ereport(LOG, (errmsg("queryid %ld does not exist", key.queryid)));
+
+	LWLockRelease(pgss->lock);
+	PG_RETURN_VOID();
+}
+
 /* Number of output arguments (columns) for various API versions */
 #define PG_STAT_STATEMENTS_COLS_V1_0	14
 #define PG_STAT_STATEMENTS_COLS_V1_1	18
diff --git a/contrib/pg_stat_statements/pg_stat_statements.control b/contrib/pg_stat_statements/pg_stat_statements.control
index 193fcdfafa..617038b4c0 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.5'
+default_version = '1.6'
 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..1888047a22 100644
--- a/contrib/pg_stat_statements/sql/pg_stat_statements.sql
+++ b/contrib/pg_stat_statements/sql/pg_stat_statements.sql
@@ -195,4 +195,18 @@ DROP FUNCTION PLUS_TWO(INTEGER);
 
 SELECT query, calls, rows FROM pg_stat_statements ORDER BY query COLLATE "C";
 
+--
+-- reset statistics of a specific query 'DROP TABLE test'
+--
+SELECT pg_stat_statements_reset_query(s.queryid) FROM pg_stat_statements AS s WHERE s.query = 'DROP TABLE test';
+
+SELECT query, calls, rows FROM pg_stat_statements ORDER BY query COLLATE "C";
+
+--
+-- try to remove unknown queryid
+--
+SELECT pg_stat_statements_reset_query(0);
+
+SELECT query, calls, rows FROM pg_stat_statements ORDER BY query COLLATE "C";
+
 DROP EXTENSION pg_stat_statements;
-- 
2.16.1.windows.4

