Add hooks for session start and session end

adunstan · adunstan · commit cd8ce3a22c0b · 2017-11-15T10:16:34.000-05:00
These hooks can be used in loadable modules. A simple test module is included. Discussion: https://postgr.es/m/20170720204733.40f2b7eb.nagata@sraoss.co.jp Fabrízio de Royes Mello and Yugo Nagata Reviewed by Michael Paquier and Aleksandr Parfenov
diff --git a/src/backend/tcop/postgres.c b/src/backend/tcop/postgres.c
@@ -169,6 +169,9 @@ static ProcSignalReason RecoveryConflictReason;
 static MemoryContext row_description_context = NULL;
 static StringInfoData row_description_buf;
 
+/* Hook for plugins to get control at start of session */
+session_start_hook_type session_start_hook = NULL;
+
 /* ----------------------------------------------------------------
  *		decls for routines only used in this file
  * ----------------------------------------------------------------
@@ -3857,6 +3860,9 @@ PostgresMain(int argc, char *argv[],
 	if (!IsUnderPostmaster)
 		PgStartTime = GetCurrentTimestamp();
 
+	if (session_start_hook)
+		(*session_start_hook) ();
+
 	/*
 	 * POSTGRES main processing loop begins here
 	 *
diff --git a/src/backend/utils/init/postinit.c b/src/backend/utils/init/postinit.c
@@ -76,6 +76,8 @@ static bool ThereIsAtLeastOneRole(void);
 static void process_startup_options(Port *port, bool am_superuser);
 static void process_settings(Oid databaseid, Oid roleid);
 
+/* Hook for plugins to get control at end of session */
+session_end_hook_type session_end_hook = NULL;
 
 /*** InitPostgres support ***/
 
@@ -1154,6 +1156,10 @@ ShutdownPostgres(int code, Datum arg)
 	 * them explicitly.
 	 */
 	LockReleaseAll(USER_LOCKMETHOD, true);
+
+	/* Hook at session end */
+	if (session_end_hook)
+		(*session_end_hook) ();
 }
 
 
diff --git a/src/include/tcop/tcopprot.h b/src/include/tcop/tcopprot.h
@@ -35,6 +35,13 @@ extern PGDLLIMPORT const char *debug_query_string;
 extern int	max_stack_depth;
 extern int	PostAuthDelay;
 
+/* Hook for plugins to get control at start and end of session */
+typedef void (*session_start_hook_type) (void);
+typedef void (*session_end_hook_type) (void);
+
+extern PGDLLIMPORT session_start_hook_type session_start_hook;
+extern PGDLLIMPORT session_end_hook_type session_end_hook;
+
 /* GUC-configurable parameters */
 
 typedef enum
diff --git a/src/test/modules/Makefile b/src/test/modules/Makefile
@@ -15,6 +15,7 @@ SUBDIRS = \
 		  test_pg_dump \
 		  test_rbtree \
 		  test_rls_hooks \
+		  test_session_hooks \
 		  test_shm_mq \
 		  worker_spi
 
diff --git a/src/test/modules/test_session_hooks/.gitignore b/src/test/modules/test_session_hooks/.gitignore
@@ -0,0 +1,4 @@
+# Generated subdirectories
+/log/
+/results/
+/tmp_check/
diff --git a/src/test/modules/test_session_hooks/Makefile b/src/test/modules/test_session_hooks/Makefile
@@ -0,0 +1,21 @@
+# src/test/modules/test_session_hooks/Makefile
+
+MODULES = test_session_hooks
+PGFILEDESC = "test_session_hooks - Test session hooks with an extension"
+
+EXTENSION = test_session_hooks
+DATA = test_session_hooks--1.0.sql
+
+REGRESS = test_session_hooks
+REGRESS_OPTS = --temp-config=$(top_srcdir)/src/test/modules/test_session_hooks/session_hooks.conf
+
+ifdef USE_PGXS
+PG_CONFIG = pg_config
+PGXS := $(shell $(PG_CONFIG) --pgxs)
+include $(PGXS)
+else
+subdir = src/test/modules/test_session_hooks
+top_builddir = ../../../..
+include $(top_builddir)/src/Makefile.global
+include $(top_srcdir)/contrib/contrib-global.mk
+endif
diff --git a/src/test/modules/test_session_hooks/README b/src/test/modules/test_session_hooks/README
@@ -0,0 +1,2 @@
+test_session_hooks is an example of how to use session start and end
+hooks.
diff --git a/src/test/modules/test_session_hooks/expected/test_session_hooks.out b/src/test/modules/test_session_hooks/expected/test_session_hooks.out
@@ -0,0 +1,31 @@
+CREATE ROLE regress_sess_hook_usr1 SUPERUSER LOGIN;
+CREATE ROLE regress_sess_hook_usr2 SUPERUSER LOGIN;
+\set prevdb :DBNAME
+\set prevusr :USER
+CREATE TABLE session_hook_log(id SERIAL, dbname TEXT, username TEXT, hook_at TEXT);
+SELECT * FROM session_hook_log ORDER BY id;
+ id | dbname | username | hook_at 
+----+--------+----------+---------
+(0 rows)
+
+\c :prevdb regress_sess_hook_usr1
+SELECT * FROM session_hook_log ORDER BY id;
+ id | dbname | username | hook_at 
+----+--------+----------+---------
+(0 rows)
+
+\c :prevdb regress_sess_hook_usr2
+SELECT * FROM session_hook_log ORDER BY id;
+ id |       dbname       |        username        | hook_at 
+----+--------------------+------------------------+---------
+  1 | contrib_regression | regress_sess_hook_usr2 | START
+(1 row)
+
+\c :prevdb :prevusr
+SELECT * FROM session_hook_log ORDER BY id;
+ id |       dbname       |        username        | hook_at 
+----+--------------------+------------------------+---------
+  1 | contrib_regression | regress_sess_hook_usr2 | START
+  2 | contrib_regression | regress_sess_hook_usr2 | END
+(2 rows)
+
diff --git a/src/test/modules/test_session_hooks/session_hooks.conf b/src/test/modules/test_session_hooks/session_hooks.conf
@@ -0,0 +1,2 @@
+shared_preload_libraries = 'test_session_hooks'
+test_session_hooks.username = regress_sess_hook_usr2
diff --git a/src/test/modules/test_session_hooks/sql/test_session_hooks.sql b/src/test/modules/test_session_hooks/sql/test_session_hooks.sql
@@ -0,0 +1,12 @@
+CREATE ROLE regress_sess_hook_usr1 SUPERUSER LOGIN;
+CREATE ROLE regress_sess_hook_usr2 SUPERUSER LOGIN;
+\set prevdb :DBNAME
+\set prevusr :USER
+CREATE TABLE session_hook_log(id SERIAL, dbname TEXT, username TEXT, hook_at TEXT);
+SELECT * FROM session_hook_log ORDER BY id;
+\c :prevdb regress_sess_hook_usr1
+SELECT * FROM session_hook_log ORDER BY id;
+\c :prevdb regress_sess_hook_usr2
+SELECT * FROM session_hook_log ORDER BY id;
+\c :prevdb :prevusr
+SELECT * FROM session_hook_log ORDER BY id;
diff --git a/src/test/modules/test_session_hooks/test_session_hooks--1.0.sql b/src/test/modules/test_session_hooks/test_session_hooks--1.0.sql
@@ -0,0 +1,4 @@
+/* src/test/modules/test_hook_session/test_hook_session--1.0.sql */
+
+-- complain if script is sourced in psql, rather than via CREATE EXTENSION
+\echo Use "CREATE EXTENSION test_hook_session" to load this file. \quit
diff --git a/src/test/modules/test_session_hooks/test_session_hooks.c b/src/test/modules/test_session_hooks/test_session_hooks.c
@@ -0,0 +1,134 @@
+/* -------------------------------------------------------------------------
+ *
+ * test_session_hooks.c
+ * 		Code for testing SESSION hooks.
+ *
+ * Copyright (c) 2010-2017, PostgreSQL Global Development Group
+ *
+ * IDENTIFICATION
+ *		src/test/modules/test_session_hooks/test_session_hooks.c
+ *
+ * -------------------------------------------------------------------------
+ */
+#include "postgres.h"
+
+#include "access/xact.h"
+#include "commands/dbcommands.h"
+#include "executor/spi.h"
+#include "lib/stringinfo.h"
+#include "miscadmin.h"
+#include "tcop/tcopprot.h"
+#include "utils/snapmgr.h"
+#include "utils/builtins.h"
+
+PG_MODULE_MAGIC;
+
+/* Entry point of library loading/unloading */
+void		_PG_init(void);
+void		_PG_fini(void);
+
+/* GUC variables */
+static char *session_hook_username = "postgres";
+
+/* Original Hook */
+static session_start_hook_type prev_session_start_hook = NULL;
+static session_end_hook_type prev_session_end_hook = NULL;
+
+static void
+register_session_hook(const char *hook_at)
+{
+	const char *username;
+
+	StartTransactionCommand();
+	SPI_connect();
+	PushActiveSnapshot(GetTransactionSnapshot());
+
+	username = GetUserNameFromId(GetUserId(), false);
+
+	/* Register log just for configured username */
+	if (!strcmp(username, session_hook_username))
+	{
+		const char *dbname;
+		int ret;
+		StringInfoData	buf;
+
+		dbname = get_database_name(MyDatabaseId);
+
+		initStringInfo(&buf);
+
+		appendStringInfo(&buf, "INSERT INTO session_hook_log (dbname, username, hook_at) ");
+		appendStringInfo(&buf, "VALUES ('%s', '%s', '%s');",
+			dbname, username, hook_at);
+
+		ret = SPI_exec(buf.data, 0);
+		if (ret != SPI_OK_INSERT)
+			elog(ERROR, "SPI_execute failed: error code %d", ret);
+	}
+
+	SPI_finish();
+	PopActiveSnapshot();
+	CommitTransactionCommand();
+}
+
+/* sample session start hook function */
+static void
+sample_session_start_hook()
+{
+	/* Hook just normal backends */
+	if (MyBackendId != InvalidBackendId)
+	{
+		(void) register_session_hook("START");
+
+		if (prev_session_start_hook)
+			prev_session_start_hook();
+	}
+}
+
+/* sample session end hook function */
+static void
+sample_session_end_hook()
+{
+	/* Hook just normal backends */
+	if (MyBackendId != InvalidBackendId)
+	{
+		if (prev_session_end_hook)
+			prev_session_end_hook();
+
+		(void) register_session_hook("END");
+	}
+}
+
+/*
+ * Module Load Callback
+ */
+void
+_PG_init(void)
+{
+	/* Save Hooks for Unload */
+	prev_session_start_hook = session_start_hook;
+	prev_session_end_hook = session_end_hook;
+
+	/* Set New Hooks */
+	session_start_hook = sample_session_start_hook;
+	session_end_hook = sample_session_end_hook;
+
+	/* Load GUCs */
+	DefineCustomStringVariable("test_session_hooks.username",
+							   "Username to register log on session start or end",
+							   NULL,
+							   &session_hook_username,
+							   "postgres",
+							   PGC_SIGHUP,
+							   0, NULL, NULL, NULL);
+}
+
+/*
+ * Module Unload Callback
+ */
+void
+_PG_fini(void)
+{
+	/* Uninstall Hooks */
+	session_start_hook = prev_session_start_hook;
+	session_end_hook = prev_session_end_hook;
+}
diff --git a/src/test/modules/test_session_hooks/test_session_hooks.control b/src/test/modules/test_session_hooks/test_session_hooks.control
@@ -0,0 +1,3 @@
+comment = 'Test start/end hook session with an extension'
+default_version = '1.0'
+relocatable = true

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,2 @@`
	`1`	`+test_session_hooks is an example of how to use session start and end`
	`2`	`+hooks.`
Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,2 @@`
	`1`	`+shared_preload_libraries = 'test_session_hooks'`
	`2`	`+test_session_hooks.username = regress_sess_hook_usr2`