From a0058023bff151f2b96ccbc592b13da91a4f11d3 Mon Sep 17 00:00:00 2001
From: Andres Freund <andres@anarazel.de>
Date: Thu, 25 Feb 2016 18:40:11 -0800
Subject: [PATCH] Add CitusExtensionOwner(), to execute some priviledged
 operations under.

There exist some operations we have to execute with elevated
privileges. The most expedient user for that is the user owning the
citusdb extension.
---
 .../distributed/utils/metadata_cache.c        | 68 +++++++++++++++++++
 src/include/distributed/metadata_cache.h      |  1 +
 2 files changed, 69 insertions(+)

diff --git a/src/backend/distributed/utils/metadata_cache.c b/src/backend/distributed/utils/metadata_cache.c
index 605a07264..d798ec4a0 100644
--- a/src/backend/distributed/utils/metadata_cache.c
+++ b/src/backend/distributed/utils/metadata_cache.c
@@ -16,7 +16,9 @@
 #include "access/htup_details.h"
 #include "access/nbtree.h"
 #include "access/xact.h"
+#include "access/sysattr.h"
 #include "catalog/indexing.h"
+#include "catalog/pg_extension.h"
 #include "catalog/pg_namespace.h"
 #include "catalog/pg_type.h"
 #include "commands/extension.h"
@@ -651,6 +653,72 @@ CitusExtraDataContainerFuncId(void)
 }
 
 
+/*
+ * CitusExtensionOwner() returns the owner of the 'citus' extension. That user
+ * is, amongst others, used to perform actions a normal user might not be
+ * allowed to perform.
+ */
+extern Oid
+CitusExtensionOwner(void)
+{
+	Relation relation = NULL;
+	SysScanDesc scandesc;
+	ScanKeyData entry[1];
+	HeapTuple extensionTuple = NULL;
+	Form_pg_extension extensionForm = NULL;
+	static Oid extensionOwner = InvalidOid;
+
+	if (extensionOwner != InvalidOid)
+	{
+		return extensionOwner;
+	}
+
+	relation = heap_open(ExtensionRelationId, AccessShareLock);
+
+	ScanKeyInit(&entry[0],
+				Anum_pg_extension_extname,
+				BTEqualStrategyNumber, F_NAMEEQ,
+				CStringGetDatum("citus"));
+
+	scandesc = systable_beginscan(relation, ExtensionNameIndexId, true,
+								  NULL, 1, entry);
+
+	extensionTuple = systable_getnext(scandesc);
+
+	/* We assume that there can be at most one matching tuple */
+	if (HeapTupleIsValid(extensionTuple))
+	{
+		extensionForm = (Form_pg_extension) GETSTRUCT(extensionTuple);
+
+		/*
+		 * For some operations Citus requires superuser permissions; we use
+		 * the extension owner for that. The extension owner is guaranteed to
+		 * be a superuser (otherwise C functions can't be created), but it'd
+		 * be possible to change the owner. So check that this still a
+		 * superuser.
+		 */
+		if (!superuser_arg(extensionForm->extowner))
+		{
+			ereport(ERROR, (errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
+							errmsg("citus extension needs to be owned by superuser")));
+		}
+		extensionOwner = extensionForm->extowner;
+		Assert(OidIsValid(extensionOwner));
+	}
+	else
+	{
+		ereport(ERROR, (errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
+						errmsg("citus extension not loaded")));
+	}
+
+	systable_endscan(scandesc);
+
+	heap_close(relation, AccessShareLock);
+
+	return extensionOwner;
+}
+
+
 /* return the  username of the currently active role */
 char *
 CurrentUserName(void)
diff --git a/src/include/distributed/metadata_cache.h b/src/include/distributed/metadata_cache.h
index 8df923734..d66268a98 100644
--- a/src/include/distributed/metadata_cache.h
+++ b/src/include/distributed/metadata_cache.h
@@ -70,5 +70,6 @@ extern Oid DistShardPlacementShardidIndexId(void);
 extern Oid CitusExtraDataContainerFuncId(void);
 
 /* user related functions */
+extern Oid CitusExtensionOwner(void);
 extern char * CurrentUserName(void);
 #endif /* METADATA_CACHE_H */