From c7157eb7042273c5c2e3e60cf559322a3d9bf2d0 Mon Sep 17 00:00:00 2001
From: Mike Palmiotto <mike.palmiotto@crunchydata.com>
Date: Fri, 28 Sep 2018 12:51:16 -0400
Subject: [PATCH] Add set_user hooks queue

Previously, each subsequent hook would overwrite the last. Add a
queueing mechanism, so we can register multiple hooks with set_user.
---
 set_user.c | 24 ++++++++++++++----------
 set_user.h | 32 +++++++++++++++-----------------
 2 files changed, 29 insertions(+), 27 deletions(-)

diff --git a/set_user.c b/set_user.c
index 97ea8fc..1e77ba0 100644
--- a/set_user.c
+++ b/set_user.c
@@ -630,19 +630,23 @@ PU_hook(Node *parsetree, const char *queryString,
 void
 PostSetUserHook(bool is_reset, const char *username)
 {
-	/* Inter-extension hooks */
-	SetUserHooks **post_hooks;
+	List		  **hooks_queue;
+	ListCell	   *hooks_entry = NULL;
 
-	post_hooks = (SetUserHooks **) find_rendezvous_variable(SET_USER_HOOKS_KEY);
-	if (*post_hooks)
+	hooks_queue = (List **) find_rendezvous_variable(SET_USER_HOOKS_KEY);
+	foreach (hooks_entry, *hooks_queue)
 	{
-		if (!is_reset && (*post_hooks)->post_set_user)
+		SetUserHooks	**post_hooks = (SetUserHooks **) lfirst(hooks_entry);
+		if (post_hooks)
 		{
-			(*post_hooks)->post_set_user(username);
-		}
-		else if ((*post_hooks)->post_reset_user)
-		{
-			(*post_hooks)->post_reset_user();
+			if (!is_reset && (*post_hooks)->post_set_user)
+			{
+				(*post_hooks)->post_set_user(username);
+			}
+			else if ((*post_hooks)->post_reset_user)
+			{
+				(*post_hooks)->post_reset_user();
+			}
 		}
 	}
 }
diff --git a/set_user.h b/set_user.h
index 1aef515..f2c1419 100644
--- a/set_user.h
+++ b/set_user.h
@@ -16,27 +16,25 @@ typedef struct SetUserHooks
  * set_user hooks.
  *
  * Takes in two function pointers, which should be defined in the extension.
+ * Each subsequent call to this function adds a new hook to the queue.
  */
 static inline void register_set_user_hooks(void *set_user_hook, void *reset_user_hook)
 {
-	static SetUserHooks **set_user_hooks = NULL;
-	MemoryContext oldcontext;
-
-	/*
-	 * Grab the set user hooks from the rendezvous hash. This should always be
-	 * NULL, unless some other extension has implemented set_user hooks.
-	 */
-	set_user_hooks = (SetUserHooks **) find_rendezvous_variable(SET_USER_HOOKS_KEY);
-
-	/*
-	 * Populate the hash entry for set_user hooks with extension function
-	 * pointers. These hooks should really only be implemented by one
-	 * extension, so we'll zero out the entry and override, even if non-NULL.
-	 */
+	static List			  **HooksQueue;
+	static SetUserHooks	   *next_hook_entry = NULL;
+	MemoryContext			oldcontext;
+
 	oldcontext = MemoryContextSwitchTo(TopMemoryContext);
-	*set_user_hooks = palloc0(sizeof(SetUserHooks));
-	(*set_user_hooks)->post_set_user = set_user_hook;
-	(*set_user_hooks)->post_reset_user = reset_user_hook;
+
+	/* Grab the SetUserHooks queue from the rendezvous hash */
+	HooksQueue = (List **) find_rendezvous_variable(SET_USER_HOOKS_KEY);
+
+	/* Populate a new hooks entry and append it to the queue */
+	next_hook_entry = palloc0(sizeof(SetUserHooks));
+	next_hook_entry->post_set_user = set_user_hook;
+	next_hook_entry->post_reset_user = reset_user_hook;
+
+	*HooksQueue = lappend(*HooksQueue, &next_hook_entry);
 	MemoryContextSwitchTo(oldcontext);
 }