aboutsummaryrefslogtreecommitdiff
path: root/src/test/modules/test_session_hooks/test_session_hooks.c
blob: d047c5d219ab3d555efe3868f6432d237865845b (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
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;
}