From e788bd924c19e296bd34316e30e3ba1b68354e64 Mon Sep 17 00:00:00 2001 From: Michael Paquier Date: Tue, 1 Oct 2019 12:15:25 +0900 Subject: [PATCH] Add hooks for session start and session end, take two MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit These hooks can be used in loadable modules. A simple test module is included. The first attempt was done with cd8ce3a but we lacked handling for NO_INSTALLCHECK in the MSVC scripts (problem solved afterwards by 431f1599) so the buildfarm got angry. This also fixes a couple of issues noticed upon review compared to the first attempt, so the code has slightly changed, resulting in a more simple test module. Author: Fabrízio de Royes Mello, Yugo Nagata Reviewed-by: Andrew Dunstan, Michael Paquier, Aleksandr Parfenov Discussion: https://postgr.es/m/20170720204733.40f2b7eb.nagata@sraoss.co.jp Discussion: https://postgr.es/m/20190823042602.GB5275@paquier.xyz --- src/backend/tcop/postgres.c | 6 + src/backend/utils/init/postinit.c | 6 + src/include/tcop/tcopprot.h | 7 + src/test/modules/Makefile | 1 + .../modules/test_session_hooks/.gitignore | 4 + src/test/modules/test_session_hooks/Makefile | 23 +++ src/test/modules/test_session_hooks/README | 11 ++ .../expected/test_session_hooks.out | 37 +++++ .../test_session_hooks/session_hooks.conf | 2 + .../sql/test_session_hooks.sql | 19 +++ .../test_session_hooks/test_session_hooks.c | 146 ++++++++++++++++++ 11 files changed, 262 insertions(+) create mode 100644 src/test/modules/test_session_hooks/.gitignore create mode 100644 src/test/modules/test_session_hooks/Makefile create mode 100644 src/test/modules/test_session_hooks/README create mode 100644 src/test/modules/test_session_hooks/expected/test_session_hooks.out create mode 100644 src/test/modules/test_session_hooks/session_hooks.conf create mode 100644 src/test/modules/test_session_hooks/sql/test_session_hooks.sql create mode 100644 src/test/modules/test_session_hooks/test_session_hooks.c diff --git a/src/backend/tcop/postgres.c b/src/backend/tcop/postgres.c index e8d8e6f828..6d80cc2d64 100644 --- a/src/backend/tcop/postgres.c +++ b/src/backend/tcop/postgres.c @@ -171,6 +171,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 * ---------------------------------------------------------------- @@ -3968,6 +3971,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 index 29c5ec7b58..151703a437 100644 --- a/src/backend/utils/init/postinit.c +++ b/src/backend/utils/init/postinit.c @@ -78,6 +78,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 ***/ @@ -1195,6 +1197,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 index ec21f7e45c..63581048b2 100644 --- a/src/include/tcop/tcopprot.h +++ b/src/include/tcop/tcopprot.h @@ -30,6 +30,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 index b2eaef3bff..abb3203903 100644 --- a/src/test/modules/Makefile +++ b/src/test/modules/Makefile @@ -21,6 +21,7 @@ SUBDIRS = \ test_predtest \ test_rbtree \ test_rls_hooks \ + test_session_hooks \ test_shm_mq \ unsafe_tests \ worker_spi diff --git a/src/test/modules/test_session_hooks/.gitignore b/src/test/modules/test_session_hooks/.gitignore new file mode 100644 index 0000000000..5dcb3ff972 --- /dev/null +++ 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 new file mode 100644 index 0000000000..e8ba6990f6 --- /dev/null +++ b/src/test/modules/test_session_hooks/Makefile @@ -0,0 +1,23 @@ +# src/test/modules/test_session_hooks/Makefile + +MODULE_big = test_session_hooks +OBJS = test_session_hooks.o $(WIN32RES) +PGFILEDESC = "test_session_hooks - tests for start and end session hooks" + +REGRESS = test_session_hooks +REGRESS_OPTS = --temp-config=$(top_srcdir)/src/test/modules/test_session_hooks/session_hooks.conf +# Disabled because these tests require extra configuration with +# "shared_preload_libraries=test_session_hooks", which typical +# installcheck users do not have (e.g. buildfarm clients). +NO_INSTALLCHECK = 1 + +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 new file mode 100644 index 0000000000..0fe3a986dc --- /dev/null +++ b/src/test/modules/test_session_hooks/README @@ -0,0 +1,11 @@ +test_session_hooks +================== + +test_session_hooks is an example of how to use session start and end +hooks. + +This module will insert into a pre-existing table called "session_hook_log" +a log activity which happens at session start and end. It is possible +to control which user information is logged when using the configuration +parameter "test_session_hooks.username". If set, the hooks will log only +information of the session user matching the parameter value. 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 new file mode 100644 index 0000000000..120dcc9a11 --- /dev/null +++ b/src/test/modules/test_session_hooks/expected/test_session_hooks.out @@ -0,0 +1,37 @@ +-- +-- Tests for start and end session hooks +-- +-- Only activity from role regress_sess_hook_usr2 is logged. +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) + +DROP ROLE regress_sess_hook_usr1; +DROP ROLE regress_sess_hook_usr2; diff --git a/src/test/modules/test_session_hooks/session_hooks.conf b/src/test/modules/test_session_hooks/session_hooks.conf new file mode 100644 index 0000000000..fc62b4adef --- /dev/null +++ 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 new file mode 100644 index 0000000000..1cc06f613b --- /dev/null +++ b/src/test/modules/test_session_hooks/sql/test_session_hooks.sql @@ -0,0 +1,19 @@ +-- +-- Tests for start and end session hooks +-- + +-- Only activity from role regress_sess_hook_usr2 is logged. +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; +DROP ROLE regress_sess_hook_usr1; +DROP ROLE regress_sess_hook_usr2; diff --git a/src/test/modules/test_session_hooks/test_session_hooks.c b/src/test/modules/test_session_hooks/test_session_hooks.c new file mode 100644 index 0000000000..d047c5d219 --- /dev/null +++ b/src/test/modules/test_session_hooks/test_session_hooks.c @@ -0,0 +1,146 @@ +/* ------------------------------------------------------------------------- + * + * test_session_hooks.c + * Code for testing start and end session hooks. + * + * Portions Copyright (c) 1996-2019, PostgreSQL Global Development Group + * Portions Copyright (c) 1994, Regents of the University of California + * + * 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"; + +/* Previous hooks on stack */ +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()); + + /* Check the current user validity */ + username = GetUserNameFromId(GetUserId(), false); + + /* Register log just for configured username */ + if (strcmp(username, session_hook_username) == 0) + { + 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);", + quote_literal_cstr(dbname), + quote_literal_cstr(username), + quote_literal_cstr(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(void) +{ + if (prev_session_start_hook) + prev_session_start_hook(); + + /* consider only normal backends */ + if (MyBackendId == InvalidBackendId) + return; + + /* consider backends connected to a database */ + if (!OidIsValid(MyDatabaseId)) + return; + + register_session_hook("START"); +} + +/* sample session end hook function */ +static void +sample_session_end_hook(void) +{ + if (prev_session_end_hook) + prev_session_end_hook(); + + /* consider only normal backends */ + if (MyBackendId == InvalidBackendId) + return; + + /* consider backends connected to a database */ + if (!OidIsValid(MyDatabaseId)) + return; + + register_session_hook("END"); +} + +/* + * Module load callback + */ +void +_PG_init(void) +{ + /* Save previous hooks */ + 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; +}