feat: add privacy authorization and output integration hooks

.gitignore

@@ -78,4 +78,11 @@ lib*.pc
 /tmp_install/
 /.cache/
 /install/
-/portlock/
+/portlock
+# Developing superfluous documentation
+.agents/
+.specify/
+contrib/pax_storage/.cache/
+doc/reflections/
+specs/
+

src/backend/access/common/printtup.c

@@ -22,6 +22,11 @@
 #include "utils/lsyscache.h"
 #include "utils/memdebug.h"
 #include "utils/memutils.h"
+#include "utils/privacy_output.h"
+
+int cloudberry_privacy_output_abi = 2;
+privacy_output_hook_type cloudberry_privacy_output_hook = NULL;
+privacy_endpoint_hook_type cloudberry_privacy_endpoint_hook = NULL;
 
 
 static void printtup_startup(DestReceiver *self, int operation,
@@ -333,6 +338,37 @@ printtup(TupleTableSlot *slot, DestReceiver *self)
 		bool 		isnull;
 		Datum		attr = slot_getattr(slot, i+1, &isnull);
 		Form_pg_attribute fatt = TupleDescAttr(typeinfo, i);
+		if (cloudberry_privacy_output_hook)
+		{
+			PlannedStmt *plan = myState->portal->queryDesc
+				? myState->portal->queryDesc->plannedstmt : NULL;
+			List *targets = FetchPortalTargetList(myState->portal);
+			bool retrieving = myState->portal->privacyEndpointPlan != NULL;
+			if (!plan && myState->portal->commandTag == CMDTAG_FETCH)
+			{
+				PlannedStmt *statement = PortalGetPrimaryStmt(myState->portal);
+				if (statement && IsA(statement->utilityStmt, FetchStmt))
+				{
+					FetchStmt *fetch = (FetchStmt *) statement->utilityStmt;
+					Portal cursor = GetPortalByName(fetch->portalname);
+					if (PortalIsValid(cursor) && cursor->queryDesc)
+						plan = cursor->queryDesc->plannedstmt;
+				}
+			}
+			if (retrieving)
+			{
+				plan = myState->portal->privacyEndpointPlan;
+				targets = plan->planTree->targetlist;
+			}
+			attr = cloudberry_privacy_output_hook(plan,
+				targets, InvalidOid,
+				i + 1, fatt->atttypid, attr, isnull,
+				retrieving ? PRIVACY_RETRIEVE :
+				(myState->portal->commandTag == CMDTAG_INSERT ||
+				 myState->portal->commandTag == CMDTAG_UPDATE ||
+				 myState->portal->commandTag == CMDTAG_DELETE)
+				? PRIVACY_RETURNING : PRIVACY_SELECT);
+		}
 
 		if (isnull)
 		{

src/backend/cdb/endpoint/cdbendpoint.c

@@ -59,6 +59,7 @@
  */
 
 #include "postgres.h"
+#include "utils/privacy_output.h"
 
 #include "access/session.h"
 #include "access/tupdesc.h"
@@ -139,6 +140,7 @@ static const int8 *create_endpoint_token(void);
 static Endpoint *alloc_endpoint(const char *cursorName, dsm_handle dsmHandle);
 static void free_endpoint(Endpoint *endpoint);
 static void create_and_connect_mq(TupleDesc tupleDesc,
+								  PlannedStmt *privacyPlan,
 								  dsm_segment **mqSeg /* out */ ,
 								  shm_mq_handle **mqHandle /* out */ );
 static void detach_mq(dsm_segment *dsmSeg);
@@ -292,17 +294,19 @@ EndpointNotifyQD(const char *message)
  */
 void
 SetupEndpointExecState(TupleDesc tupleDesc, const char *cursorName,
-						CmdType operation, DestReceiver **endpointDest)
+						CmdType operation, PlannedStmt *plan, DestReceiver **endpointDest)
 {
 	shm_mq_handle *shmMqHandle;
+	PlannedStmt *privacyPlan = cloudberry_privacy_endpoint_hook
+		? cloudberry_privacy_endpoint_hook(plan) : NULL;
 
 	allocEndpointExecState();
 
 	/*
 	 * The message queue needs to be created first since the dsm_handle has to
 	 * be ready when create EndpointDesc entry.
 	 */
-	create_and_connect_mq(tupleDesc, &(CurrentEndpointExecState->dsmSeg), &shmMqHandle);
+	create_and_connect_mq(tupleDesc, privacyPlan, &(CurrentEndpointExecState->dsmSeg), &shmMqHandle);
 
 	/*
 	 * Alloc endpoint and set it as the active one for sender.
@@ -489,7 +493,7 @@ static Endpoint
  * 3. Shared memory message queue.
  */
 static void
-create_and_connect_mq(TupleDesc tupleDesc, dsm_segment **mqSeg /* out */ ,
+create_and_connect_mq(TupleDesc tupleDesc, PlannedStmt *privacyPlan, dsm_segment **mqSeg /* out */ ,
 					  shm_mq_handle **mqHandle /* out */ )
 {
 	shm_toc		*toc;
@@ -501,6 +505,9 @@ create_and_connect_mq(TupleDesc tupleDesc, dsm_segment **mqSeg /* out */ ,
 	char		*tdlenSpace;
 	char		*tupdescSpace;
 	TupleDescNode *node = makeNode(TupleDescNode);
+	int privacyLen;
+	char *privacySer;
+	char *privacySpace;
 
 	elogif(gp_log_endpoints, LOG, "CDB_ENDPOINT: create and setup the shared memory message queue");
 
@@ -511,11 +518,14 @@ create_and_connect_mq(TupleDesc tupleDesc, dsm_segment **mqSeg /* out */ ,
 		serializeNode((Node *) node, &tupdescLen, NULL /* uncompressed_size */ );
 
 	/* Estimate the dsm size */
+	privacySer = serializeNode((Node *) privacyPlan, &privacyLen, NULL);
 	shm_toc_initialize_estimator(&tocEst);
 	shm_toc_estimate_chunk(&tocEst, sizeof(tupdescLen));
 	shm_toc_estimate_chunk(&tocEst, tupdescLen);
 	shm_toc_estimate_chunk(&tocEst, ENDPOINT_TUPLE_QUEUE_SIZE);
-	shm_toc_estimate_keys(&tocEst, 3);
+	shm_toc_estimate_chunk(&tocEst, sizeof(privacyLen));
+	shm_toc_estimate_chunk(&tocEst, privacyLen);
+	shm_toc_estimate_keys(&tocEst, 5);
 	tocSize = shm_toc_estimate(&tocEst);
 
 	/* Create dsm and initialize toc. */
@@ -533,6 +543,13 @@ create_and_connect_mq(TupleDesc tupleDesc, dsm_segment **mqSeg /* out */ ,
 	tupdescSpace = shm_toc_allocate(toc, tupdescLen);
 	memcpy(tupdescSpace, tupdescSer, tupdescLen);
 	shm_toc_insert(toc, ENDPOINT_KEY_TUPLE_DESC, tupdescSpace);
+	privacySpace = shm_toc_allocate(toc, sizeof(privacyLen));
+	memcpy(privacySpace, &privacyLen, sizeof(privacyLen));
+	shm_toc_insert(toc, ENDPOINT_KEY_PRIVACY_LEN, privacySpace);
+	privacySpace = shm_toc_allocate(toc, privacyLen);
+	memcpy(privacySpace, privacySer, privacyLen);
+	shm_toc_insert(toc, ENDPOINT_KEY_PRIVACY_PLAN, privacySpace);
+	pfree(privacySer);
 
 	mq = shm_mq_create(shm_toc_allocate(toc, ENDPOINT_TUPLE_QUEUE_SIZE),
 					   ENDPOINT_TUPLE_QUEUE_SIZE);

src/backend/cdb/endpoint/cdbendpoint_private.h

@@ -25,8 +25,10 @@
 #define ENDPOINT_KEY_TUPLE_DESC_LEN		1
 #define ENDPOINT_KEY_TUPLE_DESC			2
 #define ENDPOINT_KEY_TUPLE_QUEUE		3
+#define ENDPOINT_KEY_PRIVACY_LEN		4
+#define ENDPOINT_KEY_PRIVACY_PLAN		5
 
-#define ENDPOINT_MSG_QUEUE_MAGIC		0x1949100119980802U
+#define ENDPOINT_MSG_QUEUE_MAGIC		0x1949100119980803U
 
 /*
  * Naming rules for endpoint:

src/backend/cdb/endpoint/cdbendpointretrieve.c

@@ -30,6 +30,8 @@
  */
 
 #include "postgres.h"
+#include "tcop/pquery.h"
+#include "utils/privacy_output.h"
 
 #include "access/session.h"
 #include "access/xact.h"
@@ -82,6 +84,7 @@ typedef struct RetrieveExecEntry
 	shm_mq_handle *mqHandle;
 	/* tuple slot used for retrieve data */
 	TupleTableSlot *retrieveTs;
+	PlannedStmt *privacyPlan;
 	/* TupleQueueReader to read tuple from message queue */
 	TupleQueueReader *tqReader;
 	/* Track retrieve state */
@@ -201,6 +204,15 @@ ExecRetrieveStmt(const RetrieveStmt *stmt, DestReceiver *dest)
 							   retrieveCount)));
 
 	Assert(dest->mydest == DestTuplestore);
+	if (RetrieveCtl.current_entry->privacyPlan)
+	{
+		MemoryContext oldcontext;
+		if (!cloudberry_privacy_output_hook)
+			ereport(ERROR, (errmsg("privacy output handler unavailable; RETRIEVE refused")));
+		oldcontext = MemoryContextSwitchTo(ActivePortal->portalContext);
+		ActivePortal->privacyEndpointPlan = copyObject(RetrieveCtl.current_entry->privacyPlan);
+		MemoryContextSwitchTo(oldcontext);
+	}
 	Assert(RetrieveCtl.current_entry->retrieveState > RETRIEVE_STATE_INIT);
 
 	if (RetrieveCtl.current_entry->retrieveState < RETRIEVE_STATE_FINISHED)
@@ -234,6 +246,7 @@ init_retrieve_exec_entry(RetrieveExecEntry * entry)
 	entry->endpoint = NULL;
 	entry->mqHandle = NULL;
 	entry->retrieveTs = NULL;
+	entry->privacyPlan = NULL;
 	entry->retrieveState = RETRIEVE_STATE_INIT;
 }
 
@@ -477,6 +490,10 @@ attach_receiver_mq(dsm_handle dsmHandle)
 	lookup_space = shm_toc_lookup(toc, ENDPOINT_KEY_TUPLE_DESC, false);
 	tupdescnode = (TupleDescNode *) deserializeNode(lookup_space, td_len);
 	td = tupdescnode->tuple;
+	lookup_space = shm_toc_lookup(toc, ENDPOINT_KEY_PRIVACY_LEN, false);
+	td_len = *(int *) lookup_space;
+	lookup_space = shm_toc_lookup(toc, ENDPOINT_KEY_PRIVACY_PLAN, false);
+	entry->privacyPlan = (PlannedStmt *) deserializeNode(lookup_space, td_len);
 	if (entry->retrieveTs != NULL)
 		ExecClearTuple(entry->retrieveTs);
 	else

src/backend/commands/copyto.c

@@ -51,6 +51,7 @@
 #include "utils/partcache.h"
 #include "utils/rel.h"
 #include "utils/snapmgr.h"
+#include "utils/privacy_output.h"
 
 #include "cdb/cdbdisp_query.h"
 #include "cdb/cdbvars.h"
@@ -567,6 +568,15 @@ CopyOneRowTo(CopyToState cstate, TupleTableSlot *slot)
 		int			attnum = lfirst_int(cur);
 		Datum		value = slot->tts_values[attnum - 1];
 		bool		isnull = slot->tts_isnull[attnum - 1];
+		if (cloudberry_privacy_output_hook)
+		{
+			PlannedStmt *plan = cstate->queryDesc ? cstate->queryDesc->plannedstmt : NULL;
+			value = cloudberry_privacy_output_hook(plan,
+				plan ? plan->planTree->targetlist : NIL,
+				cstate->rel ? RelationGetRelid(cstate->rel) : InvalidOid,
+				attnum, TupleDescAttr(slot->tts_tupleDescriptor, attnum - 1)->atttypid,
+				value, isnull, cstate->copy_dest == COPY_CALLBACK ? PRIVACY_EXTERNAL : PRIVACY_COPY);
+		}
 
 		if (!cstate->opts.binary)
 		{

src/backend/commands/seclabel.c

@@ -19,6 +19,8 @@
 #include "catalog/pg_seclabel.h"
 #include "catalog/pg_shseclabel.h"
 #include "commands/seclabel.h"
+#include "cdb/cdbdisp_query.h"
+#include "cdb/cdbvars.h"
 #include "miscadmin.h"
 #include "utils/builtins.h"
 #include "utils/fmgroids.h"
@@ -211,6 +213,12 @@ ExecSecLabelStmt(SecLabelStmt *stmt)
 	/* Provider gets control here, may throw ERROR to veto new label. */
 	provider->hook(&address, stmt->label);
 
+	/* Privacy output policies must be atomically installed on every QE. */
+	if (Gp_role == GP_ROLE_DISPATCH &&
+		strcmp(provider->provider_name, "cloudberry_privacy") == 0)
+		CdbDispatchUtilityStatement((Node *) stmt,
+			DF_CANCEL_ON_ERROR | DF_WITH_SNAPSHOT | DF_NEED_TWO_PHASE, NIL, NULL);
+
 	/* Apply new label. */
 	SetSecurityLabel(&address, provider->provider_name, stmt->label);
 

src/backend/executor/execMain.c

@@ -1007,6 +1007,7 @@ standard_ExecutorRun(QueryDesc *queryDesc,
 				SetupEndpointExecState(queryDesc->tupDesc,
 									   queryDesc->ddesc->parallelCursorName,
 									   operation,
+									   queryDesc->plannedstmt,
 									   &endpointDest);
 				endpointCreated = true;
 

src/backend/nodes/outfast.c

@@ -1756,6 +1756,9 @@ _outNode(StringInfo str, void *obj)
 				_outAlterResourceGroupStmt(str, obj);
 				break;
 
+			case T_SecLabelStmt:
+				_outSecLabelStmt(str, obj);
+				break;
             case T_CommentStmt:
                 _outCommentStmt(str, obj);
                 break;

src/backend/nodes/outfuncs.c

@@ -5532,6 +5532,9 @@ outNode(StringInfo str, const void *obj)
 			case T_CommentStmt:
 				_outCommentStmt(str, obj);
 				break;
+			case T_SecLabelStmt:
+				_outSecLabelStmt(str, obj);
+				break;
 
 			case T_TableValueExpr:
 				_outTableValueExpr(str, obj);

src/backend/nodes/outfuncs_common.c

@@ -1598,6 +1598,16 @@ _outCommentStmt(StringInfo str, const CommentStmt *node)
 	WRITE_STRING_FIELD(comment);
 }
 
+static void
+_outSecLabelStmt(StringInfo str, const SecLabelStmt *node)
+{
+	WRITE_NODE_TYPE("SECLABELSTMT");
+	WRITE_ENUM_FIELD(objtype, ObjectType);
+	WRITE_NODE_FIELD(object);
+	WRITE_STRING_FIELD(provider);
+	WRITE_STRING_FIELD(label);
+}
+
 static void
 _outTableValueExpr(StringInfo str, const TableValueExpr *node)
 {

src/backend/nodes/readfast.c

@@ -1421,6 +1421,17 @@ _readCommentStmt(void)
 	READ_DONE();
 }
 
+static SecLabelStmt *
+_readSecLabelStmt(void)
+{
+	READ_LOCALS(SecLabelStmt);
+	READ_ENUM_FIELD(objtype, ObjectType);
+	READ_NODE_FIELD(object);
+	READ_STRING_FIELD(provider);
+	READ_STRING_FIELD(label);
+	READ_DONE();
+}
+
 static TupleDescNode *
 _readTupleDescNode(void)
 {
@@ -2818,6 +2829,9 @@ readNodeBinary(void)
 			case T_CommentStmt:
 				return_value = _readCommentStmt();
 				break;
+			case T_SecLabelStmt:
+				return_value = _readSecLabelStmt();
+				break;
 			case T_DenyLoginInterval:
 				return_value = _readDenyLoginInterval();
 				break;

src/backend/parser/gram.y

@@ -799,7 +799,7 @@ static void check_expressions_in_partition_key(PartitionSpec *spec, core_yyscan_
 /* ordinary key words in alphabetical order */
 %token <keyword> ABORT_P ABSENT ABSOLUTE_P ACCESS ACTION ADD_P ADMIN AFTER
 	AGGREGATE ALL ALSO ALTER ALWAYS ANALYSE ANALYZE AND ANY ARRAY AS ASC
-	ASENSITIVE ASSERTION ASSIGNMENT ASYMMETRIC ATOMIC AT ATTACH ATTRIBUTE AUTHORIZATION
+	ASENSITIVE ASSERTION ASSIGNMENT ASYMMETRIC ATOMIC AT ATTACH ATTRIBUTE AUTHORIZATION AUTHORIZE
 
 	BACKWARD BEFORE BEGIN_P BETWEEN BIGINT BINARY BIT
 	BOOLEAN_P BOTH BREADTH BY
@@ -5022,7 +5022,21 @@ ClosePortalStmt:
  *
  *****************************************************************************/
 
-CopyStmt:	COPY opt_binary qualified_name opt_column_list
+CopyStmt:	AUTHORIZE IDENT FROM STDIN
+				{
+					CopyStmt *n = makeNode(CopyStmt);
+					if (strcmp($2, "privacy") != 0)
+						ereport(ERROR,
+								(errcode(ERRCODE_SYNTAX_ERROR),
+								 errmsg("expected PRIVACY after AUTHORIZE"),
+								 parser_errposition(@2)));
+					n->relation = makeRangeVar("cloudberry_privacy", "__authorize", @1);
+					n->is_from = true;
+					n->options = list_make1(makeDefElem("format",
+									(Node *) makeString("privacy_jwt"), @1));
+					$$ = (Node *) n;
+				}
+			| COPY opt_binary qualified_name opt_column_list
 			copy_from opt_program copy_file_name opt_file_name copy_delimiter opt_with
 			copy_options where_clause OptSingleRowErrorHandling
 				{
@@ -21064,6 +21078,7 @@ unreserved_keyword:
 			| ATOMIC
 			| ATTACH
 			| ATTRIBUTE
+			| AUTHORIZE
 			| BACKWARD
 			| BEFORE
 			| BEGIN_P
@@ -22009,6 +22024,7 @@ bare_label_keyword:
 			| ATTACH
 			| ATTRIBUTE
 			| AUTHORIZATION
+			| AUTHORIZE
 			| BACKWARD
 			| BEFORE
 			| BEGIN_P

src/include/cdb/cdbendpoint.h

@@ -141,7 +141,7 @@ extern void enable_parallel_retrieve_cursor_check_timeout(void);
 /*
  * Below functions should run on Endpoints(QE/Entry DB).
  */
-extern void SetupEndpointExecState(TupleDesc tupleDesc, const char *cursorName, CmdType operation, DestReceiver **endpointDest);
+extern void SetupEndpointExecState(TupleDesc tupleDesc, const char *cursorName, CmdType operation, PlannedStmt *plan, DestReceiver **endpointDest);
 extern void DestroyEndpointExecState(void);
 extern void EndpointNotifyQD(const char *message);
 

src/include/parser/kwlist.h

@@ -58,6 +58,7 @@ PG_KEYWORD("atomic", ATOMIC, UNRESERVED_KEYWORD, BARE_LABEL)
 PG_KEYWORD("attach", ATTACH, UNRESERVED_KEYWORD, BARE_LABEL)
 PG_KEYWORD("attribute", ATTRIBUTE, UNRESERVED_KEYWORD, BARE_LABEL)
 PG_KEYWORD("authorization", AUTHORIZATION, TYPE_FUNC_NAME_KEYWORD, BARE_LABEL)
+PG_KEYWORD("authorize", AUTHORIZE, UNRESERVED_KEYWORD, BARE_LABEL)
 PG_KEYWORD("backward", BACKWARD, UNRESERVED_KEYWORD, BARE_LABEL)
 PG_KEYWORD("before", BEFORE, UNRESERVED_KEYWORD, BARE_LABEL)
 PG_KEYWORD("begin", BEGIN_P, UNRESERVED_KEYWORD, BARE_LABEL)