From 1b16a9fe0222eca9d064712fd1a9550179e44798 Mon Sep 17 00:00:00 2001
From: Sait Talha Nisanci <s.talhanisanci@gmail.com>
Date: Wed, 30 Dec 2020 14:19:03 +0300
Subject: [PATCH] wip

---
 src/backend/distributed/cimv/create.c           |  8 --------
 src/backend/distributed/cimv/refresh.c          |  9 ---------
 src/backend/distributed/commands/utility_hook.c | 15 ++++++++-------
 3 files changed, 8 insertions(+), 24 deletions(-)

diff --git a/src/backend/distributed/cimv/create.c b/src/backend/distributed/cimv/create.c
index e0c7560e4..d0e5e9dde 100644
--- a/src/backend/distributed/cimv/create.c
+++ b/src/backend/distributed/cimv/create.c
@@ -149,13 +149,6 @@ CreateCimv(CimvCreate *cimvCreate)
 		elog(ERROR, "SPI_connect failed");
 	}
 
-	Oid savedUserId = InvalidOid;
-	int savedSecurityContext = 0;
-
-	/* make sure we have write access */
-	GetUserIdAndSecContext(&savedUserId, &savedSecurityContext);
-	SetUserIdAndSecContext(CitusExtensionOwner(), SECURITY_LOCAL_USERID_CHANGE);
-
 	CreateMatTable(cimvCreate, false);
 
 	if (cimvCreate->createOptions->schedule != NULL)
@@ -180,7 +173,6 @@ CreateCimv(CimvCreate *cimvCreate)
 		RefreshCimv(cimvCreate->formCimv, cimvCreate->stmt->into->skipData, true);
 	}
 
-	SetUserIdAndSecContext(savedUserId, savedSecurityContext);
 }
 
 
diff --git a/src/backend/distributed/cimv/refresh.c b/src/backend/distributed/cimv/refresh.c
index d035a7079..0d4eae3fd 100644
--- a/src/backend/distributed/cimv/refresh.c
+++ b/src/backend/distributed/cimv/refresh.c
@@ -50,13 +50,6 @@ RefreshCimv(Form_pg_cimv formCimv, bool skipData, bool isCreate)
 		elog(ERROR, "SPI_connect failed");
 	}
 
-	Oid savedUserId = InvalidOid;
-	int savedSecurityContext = 0;
-
-	/* make sure we have write access */
-	GetUserIdAndSecContext(&savedUserId, &savedSecurityContext);
-	SetUserIdAndSecContext(CitusExtensionOwner(), SECURITY_LOCAL_USERID_CHANGE);
-
 	const char *matTableSchemaName = get_namespace_name(get_rel_namespace(
 															formCimv->mattable));
 	const char *matTableName = get_rel_name(formCimv->mattable);
@@ -156,8 +149,6 @@ RefreshCimv(Form_pg_cimv formCimv, bool skipData, bool isCreate)
 		}
 	}
 
-	SetUserIdAndSecContext(savedUserId, savedSecurityContext);
-
 	/* Close SPI context. */
 	if (SPI_finish() != SPI_OK_FINISH)
 	{
diff --git a/src/backend/distributed/commands/utility_hook.c b/src/backend/distributed/commands/utility_hook.c
index c00aff181..f623ca75b 100644
--- a/src/backend/distributed/commands/utility_hook.c
+++ b/src/backend/distributed/commands/utility_hook.c
@@ -56,6 +56,7 @@
 #include "distributed/multi_explain.h"
 #include "distributed/multi_physical_planner.h"
 #include "distributed/resource_lock.h"
+#include "distributed/security_utils.h"
 #include "distributed/transmit.h"
 #include "distributed/version_compat.h"
 #include "distributed/worker_transaction.h"
@@ -460,29 +461,27 @@ multi_ProcessUtility(PlannedStmt *pstmt,
 	bool continueProcessing = true;
 	if (IsA(parsetree, CreateTableAsStmt))
 	{
-		Oid savedUserId = InvalidOid;
-		int savedSecurityContext = 0;
-
-		/* make sure we have write access */
-		GetUserIdAndSecContext(&savedUserId, &savedSecurityContext);
-		SetUserIdAndSecContext(CitusExtensionOwner(), SECURITY_LOCAL_USERID_CHANGE);
+		PushCitusSecurityContext();
 		continueProcessing = !ProcessCreateMaterializedViewStmt((const
 																 CreateTableAsStmt *)
 																parsetree, queryString,
 																pstmt);
-		SetUserIdAndSecContext(savedUserId, savedSecurityContext);														
+		PopCitusSecurityContext();														
 	}
 
 	if (IsA(parsetree, RefreshMatViewStmt))
 	{
+		PushCitusSecurityContext();
 		continueProcessing = !ProcessRefreshMaterializedViewStmt(
 			(RefreshMatViewStmt *) parsetree);
+		PopCitusSecurityContext();	
 	}
 
 	if (IsA(parsetree, DropStmt))
 	{
 		DropStmt *dropStatement = (DropStmt *) parsetree;
 
+		PushCitusSecurityContext();
 		if (dropStatement->removeType == OBJECT_MATVIEW)
 		{
 			ProcessDropMaterializedViewStmt(dropStatement);
@@ -491,6 +490,8 @@ multi_ProcessUtility(PlannedStmt *pstmt,
 		{
 			ProcessDropViewStmt(dropStatement);
 		}
+		PopCitusSecurityContext();
+
 	}
 
 	if (IsDropCitusExtensionStmt(parsetree))