Skip to content

Commit cd8ce3a

Browse files
adunstanadunstan
committed
Add hooks for session start and session end
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
1 parent ebc189e commit cd8ce3a

File tree

13 files changed

+233
-0
lines changed

13 files changed

+233
-0
lines changed

‎src/backend/tcop/postgres.c

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

172+
/* Hook for plugins to get control at start of session */
173+
session_start_hook_type session_start_hook = NULL;
174+
172175
/* ----------------------------------------------------------------
173176
* decls for routines only used in this file
174177
* ----------------------------------------------------------------
@@ -3857,6 +3860,9 @@ PostgresMain(int argc, char *argv[],
38573860
if (!IsUnderPostmaster)
38583861
PgStartTime = GetCurrentTimestamp();
38593862

3863+
if (session_start_hook)
3864+
(*session_start_hook) ();
3865+
38603866
/*
38613867
* POSTGRES main processing loop begins here
38623868
*

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

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

79+
/* Hook for plugins to get control at end of session */
80+
session_end_hook_type session_end_hook = NULL;
7981

8082
/*** InitPostgres support ***/
8183

@@ -1154,6 +1156,10 @@ ShutdownPostgres(int code, Datum arg)
11541156
* them explicitly.
11551157
*/
11561158
LockReleaseAll(USER_LOCKMETHOD, true);
1159+
1160+
/* Hook at session end */
1161+
if (session_end_hook)
1162+
(*session_end_hook) ();
11571163
}
11581164

11591165

‎src/include/tcop/tcopprot.h

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,13 @@ extern PGDLLIMPORT const char *debug_query_string;
3535
extern int max_stack_depth;
3636
extern int PostAuthDelay;
3737

38+
/* Hook for plugins to get control at start and end of session */
39+
typedef void (*session_start_hook_type) (void);
40+
typedef void (*session_end_hook_type) (void);
41+
42+
extern PGDLLIMPORT session_start_hook_type session_start_hook;
43+
extern PGDLLIMPORT session_end_hook_type session_end_hook;
44+
3845
/* GUC-configurable parameters */
3946

4047
typedef enum

‎src/test/modules/Makefile

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ SUBDIRS = \
1515
test_pg_dump \
1616
test_rbtree \
1717
test_rls_hooks \
18+
test_session_hooks \
1819
test_shm_mq \
1920
worker_spi
2021

‎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: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
# src/test/modules/test_session_hooks/Makefile
2+
3+
MODULES = test_session_hooks
4+
PGFILEDESC = "test_session_hooks - Test session hooks with an extension"
5+
6+
EXTENSION = test_session_hooks
7+
DATA = test_session_hooks--1.0.sql
8+
9+
REGRESS = test_session_hooks
10+
REGRESS_OPTS = --temp-config=$(top_srcdir)/src/test/modules/test_session_hooks/session_hooks.conf
11+
12+
ifdef USE_PGXS
13+
PG_CONFIG = pg_config
14+
PGXS := $(shell $(PG_CONFIG) --pgxs)
15+
include $(PGXS)
16+
else
17+
subdir = src/test/modules/test_session_hooks
18+
top_builddir = ../../../..
19+
include $(top_builddir)/src/Makefile.global
20+
include $(top_srcdir)/contrib/contrib-global.mk
21+
endif

‎src/test/modules/test_session_hooks/README

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
test_session_hooks is an example of how to use session start and end
2+
hooks.

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

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
1+
CREATE ROLE regress_sess_hook_usr1 SUPERUSER LOGIN;
2+
CREATE ROLE regress_sess_hook_usr2 SUPERUSER LOGIN;
3+
\set prevdb :DBNAME
4+
\set prevusr :USER
5+
CREATE TABLE session_hook_log(id SERIAL, dbname TEXT, username TEXT, hook_at TEXT);
6+
SELECT * FROM session_hook_log ORDER BY id;
7+
id | dbname | username | hook_at
8+
----+--------+----------+---------
9+
(0 rows)
10+
11+
\c :prevdb regress_sess_hook_usr1
12+
SELECT * FROM session_hook_log ORDER BY id;
13+
id | dbname | username | hook_at
14+
----+--------+----------+---------
15+
(0 rows)
16+
17+
\c :prevdb regress_sess_hook_usr2
18+
SELECT * FROM session_hook_log ORDER BY id;
19+
id | dbname | username | hook_at
20+
----+--------------------+------------------------+---------
21+
1 | contrib_regression | regress_sess_hook_usr2 | START
22+
(1 row)
23+
24+
\c :prevdb :prevusr
25+
SELECT * FROM session_hook_log ORDER BY id;
26+
id | dbname | username | hook_at
27+
----+--------------------+------------------------+---------
28+
1 | contrib_regression | regress_sess_hook_usr2 | START
29+
2 | contrib_regression | regress_sess_hook_usr2 | END
30+
(2 rows)
31+

‎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: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
CREATE ROLE regress_sess_hook_usr1 SUPERUSER LOGIN;
2+
CREATE ROLE regress_sess_hook_usr2 SUPERUSER LOGIN;
3+
\set prevdb :DBNAME
4+
\set prevusr :USER
5+
CREATE TABLE session_hook_log(id SERIAL, dbname TEXT, username TEXT, hook_at TEXT);
6+
SELECT * FROM session_hook_log ORDER BY id;
7+
\c :prevdb regress_sess_hook_usr1
8+
SELECT * FROM session_hook_log ORDER BY id;
9+
\c :prevdb regress_sess_hook_usr2
10+
SELECT * FROM session_hook_log ORDER BY id;
11+
\c :prevdb :prevusr
12+
SELECT * FROM session_hook_log ORDER BY id;

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
/* src/test/modules/test_hook_session/test_hook_session--1.0.sql */
2+
3+
-- complain if script is sourced in psql, rather than via CREATE EXTENSION
4+
\echo Use "CREATE EXTENSION test_hook_session" to load this file. \quit

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

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

‎src/test/modules/test_session_hooks/test_session_hooks.control

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
comment = 'Test start/end hook session with an extension'
2+
default_version = '1.0'
3+
relocatable = true

0 commit comments

Comments
 (0)