pg hooks for online table (#693)

* pg hooks for online table (#24)

Adding hooks for PG online table. 

- PreOnlineTableOp_hook
- PostOnlineTableOp_hook
- OnlineTableSecurityLabel_hook

The 3 hooks are all used to bypass the ownership check on an online
table.

The Pre/Post hooks are used to allow non-online table owner to perform a
few DDLs on the table, e.g., CREATE INDEX.
- Pre hook will check the DDL type and whether the referenced relation
is an online table. If the conditions are met, it will set the session
user as the table owner. The rest of the PG code will work since the
session is acting as the table owner now.
- Post hook will restore the current user to the session user. 

OnlineTableSecurityLabel_hook is used to allow customers to add who can
perform DDL on an online table. By default, the databricks_superuser
can. OnlineTableSecurityLabel_hook will check if the session user is one
of the roles attached to the online table. If it is, it will allow
updating the label to add / remove roles. Otherwise, it will throw the
permission error.

Why we add them?
1. These 3 hooks are very specific for online table. We don't intend to
use them for other purposes.
2. This is also the reason we are not overloading the public hook also.
I need to be very specific on the code point where we want to set as
table owner. standard_ProcessUtility contains too many stuff. I don't
want to make our change generalize to all PG utilities.

The extension change is inside
https://github.com/databricks-eng/hadron/pull/492

* First pass of review.

---------

Co-authored-by: Haoyu Huang <[email protected]>

Author: dimitri

src/backend/commands/seclabel.c

@@ -33,6 +33,10 @@ typedef struct
 
 static List *label_provider_list = NIL;
 
+/* BEGIN_PG_NEON */
+OnlineTableSecurityLabel_hook_type OnlineTableSecurityLabel_hook = NULL;
+/* END_PG_NEON */
+
 static bool
 SecLabelSupportsObjectType(ObjectType objtype)
 {
@@ -168,9 +172,18 @@ ExecSecLabelStmt(SecLabelStmt *stmt)
 	address = get_object_address(stmt->objtype, stmt->object,
 								 &relation, ShareUpdateExclusiveLock, false);
 
-	/* Require ownership of the target object. */
-	check_object_ownership(GetUserId(), stmt->objtype, address,
-						   stmt->object, relation);
+	/* BEGIN_PG_NEON */
+	if (OnlineTableSecurityLabel_hook != NULL)
+	{
+		OnlineTableSecurityLabel_hook(stmt, &address, relation);
+	}
+	else
+	{
+		/* Require ownership of the target object. */
+		check_object_ownership(GetUserId(), stmt->objtype, address,
+							   stmt->object, relation);
+	}
+	/* END_PG_NEON */
 
 	/* Perform other integrity checks as needed. */
 	switch (stmt->objtype)

src/backend/tcop/utility.c

@@ -69,6 +69,11 @@
 /* Hook for plugins to get control in ProcessUtility() */
 ProcessUtility_hook_type ProcessUtility_hook = NULL;
 
+/* BEGIN_PG_NEON */
+PreOnlineTableOp_hook_type PreOnlineTableOp_hook = NULL;
+PostOnlineTableOp_hook_type PostOnlineTableOp_hook = NULL;
+/* END_PG_NEON */
+
 /* local function declarations */
 static int	ClassifyUtilityCommandAsReadOnly(Node *parsetree);
 static void ProcessUtilitySlow(ParseState *pstate,
@@ -1112,6 +1117,13 @@ ProcessUtilitySlow(ParseState *pstate,
 		if (isCompleteQuery)
 			EventTriggerDDLCommandStart(parsetree);
 
+		/* BEGIN_PG_NEON */
+		if (PreOnlineTableOp_hook != NULL)
+		{
+			PreOnlineTableOp_hook(parsetree);
+		}
+		/* END_PG_NEON */
+
 		switch (nodeTag(parsetree))
 		{
 				/*
@@ -1932,6 +1944,13 @@ ProcessUtilitySlow(ParseState *pstate,
 	}
 	PG_FINALLY();
 	{
+		/* BEGIN_PG_NEON */
+		if (PostOnlineTableOp_hook != NULL)
+		{
+			PostOnlineTableOp_hook(parsetree);
+		}
+		/* END_PG_NEON */
+
 		if (needCleanup)
 			EventTriggerEndCompleteQuery();
 	}

src/include/commands/seclabel.h

@@ -31,4 +31,20 @@ typedef void (*check_object_relabel_type) (const ObjectAddress *object,
 extern void register_label_provider(const char *provider_name,
 									check_object_relabel_type hook);
 
+/* BEGIN_PG_NEON */
+
+/*
+ * OnlineTableSecurityLabel_hook is used to allow customers to perform DDL on
+ * an online table. By default, only databricks_superuser can.
+ * OnlineTableSecurityLabel_hook checks if the session user is one of the
+ * roles attached to the online table. If it is, it allows updating the label
+ * to add/remove roles. Otherwise, a permission error is thrown.
+ */
+typedef void (*OnlineTableSecurityLabel_hook_type) (SecLabelStmt *stmt,
+													const ObjectAddress *object,
+													Relation relation);
+
+extern PGDLLIMPORT OnlineTableSecurityLabel_hook_type OnlineTableSecurityLabel_hook;
+/* END_PG_NEON */
+
 #endif							/* SECLABEL_H */

src/include/tcop/utility.h

@@ -77,6 +77,29 @@ typedef void (*ProcessUtility_hook_type) (PlannedStmt *pstmt,
 										  DestReceiver *dest, QueryCompletion *qc);
 extern PGDLLIMPORT ProcessUtility_hook_type ProcessUtility_hook;
 
+/* BEGIN_PG_NEON */
+
+/*
+ * The Pre/Post hooks are used to allow roles that do not own an online table
+ * to still perform a few DDLs on the table, e.g., CREATE INDEX.
+ *
+ * - Pre hook will check the DDL type and whether the referenced relation is
+ *   an online table. If the conditions are met, it will set the session user
+ *   as the table owner.
+ *
+ *   Postgres code then (typically) executes as the session user, which is now
+ *   SET to the table owner.
+ *
+ * - Post hook will restore the current user to the session user.
+ */
+typedef void (*PreOnlineTableOp_hook_type) (Node *parsetree);
+extern PGDLLIMPORT PreOnlineTableOp_hook_type PreOnlineTableOp_hook;
+
+typedef void (*PostOnlineTableOp_hook_type) (Node *parsetree);
+extern PGDLLIMPORT PostOnlineTableOp_hook_type PostOnlineTableOp_hook;
+
+/* END_PG_NEON */
+
 extern void ProcessUtility(PlannedStmt *pstmt, const char *queryString,
 						   bool readOnlyTree,
 						   ProcessUtilityContext context, ParamListInfo params,