Location via proxy:   [ UP ]  
[Report a bug]   [Manage cookies]                
Skip to content

Commit e788bd9

Browse files
michaelpqmichaelpq
committed
Add hooks for session start and session end, take two
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 431f159) 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
1 parent 41a6de4 commit e788bd9

File tree

11 files changed

+262
-0
lines changed

11 files changed

+262
-0
lines changed

‎src/backend/tcop/postgres.c

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -171,6 +171,9 @@ static ProcSignalReason RecoveryConflictReason;
171171
static MemoryContext row_description_context = NULL;
172172
static StringInfoData row_description_buf;
173173

174+
/* Hook for plugins to get control at start of session */
175+
session_start_hook_type session_start_hook = NULL;
176+
174177
/* ----------------------------------------------------------------
175178
* decls for routines only used in this file
176179
* ----------------------------------------------------------------
@@ -3968,6 +3971,9 @@ PostgresMain(int argc, char *argv[],
39683971
if (!IsUnderPostmaster)
39693972
PgStartTime = GetCurrentTimestamp();
39703973

3974+
if (session_start_hook)
3975+
(*session_start_hook) ();
3976+
39713977
/*
39723978
* POSTGRES main processing loop begins here
39733979
*

‎src/backend/utils/init/postinit.c

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -78,6 +78,8 @@ static bool ThereIsAtLeastOneRole(void);
7878
static void process_startup_options(Port *port, bool am_superuser);
7979
static void process_settings(Oid databaseid, Oid roleid);
8080

81+
/* Hook for plugins to get control at end of session */
82+
session_end_hook_type session_end_hook = NULL;
8183

8284
/*** InitPostgres support ***/
8385

@@ -1195,6 +1197,10 @@ ShutdownPostgres(int code, Datum arg)
11951197
* them explicitly.
11961198
*/
11971199
LockReleaseAll(USER_LOCKMETHOD, true);
1200+
1201+
/* Hook at session end */
1202+
if (session_end_hook)
1203+
(*session_end_hook) ();
11981204
}
11991205

12001206

‎src/include/tcop/tcopprot.h

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,13 @@ extern PGDLLIMPORT const char *debug_query_string;
3030
extern int max_stack_depth;
3131
extern int PostAuthDelay;
3232

33+
/* Hook for plugins to get control at start and end of session */
34+
typedef void (*session_start_hook_type) (void);
35+
typedef void (*session_end_hook_type) (void);
36+
37+
extern PGDLLIMPORT session_start_hook_type session_start_hook;
38+
extern PGDLLIMPORT session_end_hook_type session_end_hook;
39+
3340
/* GUC-configurable parameters */
3441

3542
typedef enum

‎src/test/modules/Makefile

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ SUBDIRS = \
2121
test_predtest \
2222
test_rbtree \
2323
test_rls_hooks \
24+
test_session_hooks \
2425
test_shm_mq \
2526
unsafe_tests \
2627
worker_spi

‎src/test/modules/test_session_hooks/.gitignore

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
# Generated subdirectories
2+
/log/
3+
/results/
4+
/tmp_check/

‎src/test/modules/test_session_hooks/Makefile

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,23 @@
1+
# src/test/modules/test_session_hooks/Makefile
2+
3+
MODULE_big = test_session_hooks
4+
OBJS = test_session_hooks.o $(WIN32RES)
5+
PGFILEDESC = "test_session_hooks - tests for start and end session hooks"
6+
7+
REGRESS = test_session_hooks
8+
REGRESS_OPTS = --temp-config=$(top_srcdir)/src/test/modules/test_session_hooks/session_hooks.conf
9+
# Disabled because these tests require extra configuration with
10+
# "shared_preload_libraries=test_session_hooks", which typical
11+
# installcheck users do not have (e.g. buildfarm clients).
12+
NO_INSTALLCHECK = 1
13+
14+
ifdef USE_PGXS
15+
PG_CONFIG = pg_config
16+
PGXS := $(shell $(PG_CONFIG) --pgxs)
17+
include $(PGXS)
18+
else
19+
subdir = src/test/modules/test_session_hooks
20+
top_builddir = ../../../..
21+
include $(top_builddir)/src/Makefile.global
22+
include $(top_srcdir)/contrib/contrib-global.mk
23+
endif

‎src/test/modules/test_session_hooks/README

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,11 @@
1+
test_session_hooks
2+
==================
3+
4+
test_session_hooks is an example of how to use session start and end
5+
hooks.
6+
7+
This module will insert into a pre-existing table called "session_hook_log"
8+
a log activity which happens at session start and end. It is possible
9+
to control which user information is logged when using the configuration
10+
parameter "test_session_hooks.username". If set, the hooks will log only
11+
information of the session user matching the parameter value.

‎src/test/modules/test_session_hooks/expected/test_session_hooks.out

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
--
2+
-- Tests for start and end session hooks
3+
--
4+
-- Only activity from role regress_sess_hook_usr2 is logged.
5+
CREATE ROLE regress_sess_hook_usr1 SUPERUSER LOGIN;
6+
CREATE ROLE regress_sess_hook_usr2 SUPERUSER LOGIN;
7+
\set prevdb :DBNAME
8+
\set prevusr :USER
9+
CREATE TABLE session_hook_log(id SERIAL, dbname TEXT, username TEXT, hook_at TEXT);
10+
SELECT * FROM session_hook_log ORDER BY id;
11+
id | dbname | username | hook_at
12+
----+--------+----------+---------
13+
(0 rows)
14+
15+
\c :prevdb regress_sess_hook_usr1
16+
SELECT * FROM session_hook_log ORDER BY id;
17+
id | dbname | username | hook_at
18+
----+--------+----------+---------
19+
(0 rows)
20+
21+
\c :prevdb regress_sess_hook_usr2
22+
SELECT * FROM session_hook_log ORDER BY id;
23+
id | dbname | username | hook_at
24+
----+--------------------+------------------------+---------
25+
1 | contrib_regression | regress_sess_hook_usr2 | START
26+
(1 row)
27+
28+
\c :prevdb :prevusr
29+
SELECT * FROM session_hook_log ORDER BY id;
30+
id | dbname | username | hook_at
31+
----+--------------------+------------------------+---------
32+
1 | contrib_regression | regress_sess_hook_usr2 | START
33+
2 | contrib_regression | regress_sess_hook_usr2 | END
34+
(2 rows)
35+
36+
DROP ROLE regress_sess_hook_usr1;
37+
DROP ROLE regress_sess_hook_usr2;

‎src/test/modules/test_session_hooks/session_hooks.conf

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
shared_preload_libraries = 'test_session_hooks'
2+
test_session_hooks.username = regress_sess_hook_usr2

‎src/test/modules/test_session_hooks/sql/test_session_hooks.sql

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
--
2+
-- Tests for start and end session hooks
3+
--
4+
5+
-- Only activity from role regress_sess_hook_usr2 is logged.
6+
CREATE ROLE regress_sess_hook_usr1 SUPERUSER LOGIN;
7+
CREATE ROLE regress_sess_hook_usr2 SUPERUSER LOGIN;
8+
\set prevdb :DBNAME
9+
\set prevusr :USER
10+
CREATE TABLE session_hook_log(id SERIAL, dbname TEXT, username TEXT, hook_at TEXT);
11+
SELECT * FROM session_hook_log ORDER BY id;
12+
\c :prevdb regress_sess_hook_usr1
13+
SELECT * FROM session_hook_log ORDER BY id;
14+
\c :prevdb regress_sess_hook_usr2
15+
SELECT * FROM session_hook_log ORDER BY id;
16+
\c :prevdb :prevusr
17+
SELECT * FROM session_hook_log ORDER BY id;
18+
DROP ROLE regress_sess_hook_usr1;
19+
DROP ROLE regress_sess_hook_usr2;

‎src/test/modules/test_session_hooks/test_session_hooks.c

Lines changed: 146 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,146 @@
1+
/* -------------------------------------------------------------------------
2+
*
3+
* test_session_hooks.c
4+
* Code for testing start and end session hooks.
5+
*
6+
* Portions Copyright (c) 1996-2019, PostgreSQL Global Development Group
7+
* Portions Copyright (c) 1994, Regents of the University of California
8+
*
9+
* IDENTIFICATION
10+
* src/test/modules/test_session_hooks/test_session_hooks.c
11+
*
12+
* -------------------------------------------------------------------------
13+
*/
14+
#include "postgres.h"
15+
16+
#include "access/xact.h"
17+
#include "commands/dbcommands.h"
18+
#include "executor/spi.h"
19+
#include "lib/stringinfo.h"
20+
#include "miscadmin.h"
21+
#include "tcop/tcopprot.h"
22+
#include "utils/snapmgr.h"
23+
#include "utils/builtins.h"
24+
25+
PG_MODULE_MAGIC;
26+
27+
/* Entry point of library loading/unloading */
28+
void _PG_init(void);
29+
void _PG_fini(void);
30+
31+
/* GUC variables */
32+
static char *session_hook_username = "postgres";
33+
34+
/* Previous hooks on stack */
35+
static session_start_hook_type prev_session_start_hook = NULL;
36+
static session_end_hook_type prev_session_end_hook = NULL;
37+
38+
static void
39+
register_session_hook(const char *hook_at)
40+
{
41+
const char *username;
42+
43+
StartTransactionCommand();
44+
SPI_connect();
45+
PushActiveSnapshot(GetTransactionSnapshot());
46+
47+
/* Check the current user validity */
48+
username = GetUserNameFromId(GetUserId(), false);
49+
50+
/* Register log just for configured username */
51+
if (strcmp(username, session_hook_username) == 0)
52+
{
53+
const char *dbname;
54+
int ret;
55+
StringInfoData buf;
56+
57+
dbname = get_database_name(MyDatabaseId);
58+
59+
initStringInfo(&buf);
60+
61+
appendStringInfo(&buf, "INSERT INTO session_hook_log (dbname, username, hook_at) ");
62+
appendStringInfo(&buf, "VALUES (%s, %s, %s);",
63+
quote_literal_cstr(dbname),
64+
quote_literal_cstr(username),
65+
quote_literal_cstr(hook_at));
66+
67+
ret = SPI_exec(buf.data, 0);
68+
if (ret != SPI_OK_INSERT)
69+
elog(ERROR, "SPI_execute failed: error code %d", ret);
70+
}
71+
72+
SPI_finish();
73+
PopActiveSnapshot();
74+
CommitTransactionCommand();
75+
}
76+
77+
/* sample session start hook function */
78+
static void
79+
sample_session_start_hook(void)
80+
{
81+
if (prev_session_start_hook)
82+
prev_session_start_hook();
83+
84+
/* consider only normal backends */
85+
if (MyBackendId == InvalidBackendId)
86+
return;
87+
88+
/* consider backends connected to a database */
89+
if (!OidIsValid(MyDatabaseId))
90+
return;
91+
92+
register_session_hook("START");
93+
}
94+
95+
/* sample session end hook function */
96+
static void
97+
sample_session_end_hook(void)
98+
{
99+
if (prev_session_end_hook)
100+
prev_session_end_hook();
101+
102+
/* consider only normal backends */
103+
if (MyBackendId == InvalidBackendId)
104+
return;
105+
106+
/* consider backends connected to a database */
107+
if (!OidIsValid(MyDatabaseId))
108+
return;
109+
110+
register_session_hook("END");
111+
}
112+
113+
/*
114+
* Module load callback
115+
*/
116+
void
117+
_PG_init(void)
118+
{
119+
/* Save previous hooks */
120+
prev_session_start_hook = session_start_hook;
121+
prev_session_end_hook = session_end_hook;
122+
123+
/* Set new hooks */
124+
session_start_hook = sample_session_start_hook;
125+
session_end_hook = sample_session_end_hook;
126+
127+
/* Load GUCs */
128+
DefineCustomStringVariable("test_session_hooks.username",
129+
"Username to register log on session start or end",
130+
NULL,
131+
&session_hook_username,
132+
"postgres",
133+
PGC_SIGHUP,
134+
0, NULL, NULL, NULL);
135+
}
136+
137+
/*
138+
* Module unload callback
139+
*/
140+
void
141+
_PG_fini(void)
142+
{
143+
/* Uninstall hooks */
144+
session_start_hook = prev_session_start_hook;
145+
session_end_hook = prev_session_end_hook;
146+
}

0 commit comments

Comments
 (0)