Merge pull request #5748 from citusdata/improve_visibility_of_citus_stats

Improve visibility rules for non-priviledge roles on citus_stat_activity

src/backend/distributed/transaction/backend_data.c

@@ -83,6 +83,7 @@ typedef struct BackendManagementShmemData
 
 static void StoreAllActiveTransactions(Tuplestorestate *tupleStore, TupleDesc
 									   tupleDescriptor);
+static bool UserHasPermissionToViewStatsOf(Oid currentUserId, Oid backendOwnedId);
 static uint64 GenerateGlobalPID(void);
 
 static shmem_startup_hook_type prev_shmem_startup_hook = NULL;
@@ -114,8 +115,6 @@ assign_distributed_transaction_id(PG_FUNCTION_ARGS)
 {
 	CheckCitusVersion(ERROR);
 
-	Oid userId = GetUserId();
-
 	/* prepare data before acquiring spinlock to protect against errors */
 	int32 initiatorNodeIdentifier = PG_GETARG_INT32(0);
 	uint64 transactionNumber = PG_GETARG_INT64(1);
@@ -144,9 +143,6 @@ assign_distributed_transaction_id(PG_FUNCTION_ARGS)
 							   "transaction id")));
 	}
 
-	MyBackendData->databaseId = MyDatabaseId;
-	MyBackendData->userId = userId;
-
 	MyBackendData->transactionId.initiatorNodeIdentifier = initiatorNodeIdentifier;
 	MyBackendData->transactionId.transactionNumber = transactionNumber;
 	MyBackendData->transactionId.timestamp = timestamp;
@@ -357,49 +353,44 @@ StoreAllActiveTransactions(Tuplestorestate *tupleStore, TupleDesc tupleDescripto
 {
 	Datum values[ACTIVE_TRANSACTION_COLUMN_COUNT];
 	bool isNulls[ACTIVE_TRANSACTION_COLUMN_COUNT];
-	bool showAllTransactions = superuser();
+	bool showAllBackends = superuser();
 	const Oid userId = GetUserId();
 
-	/*
+	if (!showAllBackends && is_member_of_role(userId, ROLE_PG_MONITOR))
-	 * We don't want to initialize memory while spinlock is held so we
-	 * prefer to do it here. This initialization is done only for the first
-	 * row.
-	 */
-	memset(values, 0, sizeof(values));
-	memset(isNulls, false, sizeof(isNulls));
-
-	if (is_member_of_role(userId, ROLE_PG_MONITOR))
 	{
-		showAllTransactions = true;
+		showAllBackends = true;
 	}
 
 	/* we're reading all distributed transactions, prevent new backends */
 	LockBackendSharedMemory(LW_SHARED);
 
-	for (int backendIndex = 0; backendIndex < MaxBackends; ++backendIndex)
+	for (int backendIndex = 0; backendIndex < TotalProcCount(); ++backendIndex)
 	{
+		bool showCurrentBackendDetails = showAllBackends;
 		BackendData *currentBackend =
 			&backendManagementShmemData->backends[backendIndex];
+		PGPROC *currentProc = &ProcGlobal->allProcs[backendIndex];
 
 		/* to work on data after releasing g spinlock to protect against errors */
 		uint64 transactionNumber = 0;
 
 		SpinLockAcquire(&currentBackend->mutex);
 
-		if (currentBackend->globalPID == INVALID_CITUS_INTERNAL_BACKEND_GPID)
+		if (currentProc->pid == 0)
 		{
+			/* unused PGPROC slot */
 			SpinLockRelease(&currentBackend->mutex);
 			continue;
 		}
 
 		/*
 		 * Unless the user has a role that allows seeing all transactions (superuser,
-		 * pg_monitor), skip over transactions belonging to other users.
+		 * pg_monitor), we only follow pg_stat_statements owner checks.
 		 */
-		if (!showAllTransactions && currentBackend->userId != userId)
+		if (!showCurrentBackendDetails &&
+			UserHasPermissionToViewStatsOf(userId, currentProc->roleId))
 		{
-			SpinLockRelease(&currentBackend->mutex);
+			showCurrentBackendDetails = true;
-			continue;
 		}
 
 		Oid databaseId = currentBackend->databaseId;
@@ -418,16 +409,42 @@ StoreAllActiveTransactions(Tuplestorestate *tupleStore, TupleDesc tupleDescripto
 
 		SpinLockRelease(&currentBackend->mutex);
 
+		memset(values, 0, sizeof(values));
+		memset(isNulls, false, sizeof(isNulls));
+
+		/*
+		 * We imitate pg_stat_activity such that if a user doesn't have enough
+		 * privileges, we only show the minimal information including the pid,
+		 * global pid and distributedCommandOriginator.
+		 *
+		 * pid is already can be found in pg_stat_activity for any process, and
+		 * the rest doesn't reveal anything critial for under priviledge users
+		 * but still could be useful for monitoring purposes of Citus.
+		 */
+		if (showCurrentBackendDetails)
+		{
 			bool missingOk = true;
-		int nodeId = ExtractNodeIdFromGlobalPID(currentBackend->globalPID, missingOk);
+			int initiatorNodeId =
+				ExtractNodeIdFromGlobalPID(currentBackend->globalPID, missingOk);
 
 			values[0] = ObjectIdGetDatum(databaseId);
 			values[1] = Int32GetDatum(backendPid);
-		values[2] = Int32GetDatum(nodeId);
+			values[2] = Int32GetDatum(initiatorNodeId);
 			values[3] = !distributedCommandOriginator;
 			values[4] = UInt64GetDatum(transactionNumber);
 			values[5] = TimestampTzGetDatum(transactionIdTimestamp);
 			values[6] = UInt64GetDatum(currentBackend->globalPID);
+		}
+		else
+		{
+			isNulls[0] = true;
+			values[1] = Int32GetDatum(backendPid);
+			isNulls[2] = true;
+			values[3] = !distributedCommandOriginator;
+			isNulls[4] = true;
+			isNulls[5] = true;
+			values[6] = UInt64GetDatum(currentBackend->globalPID);
+		}
 
 		tuplestore_putvalues(tupleStore, tupleDescriptor, values, isNulls);
 
@@ -444,6 +461,35 @@ StoreAllActiveTransactions(Tuplestorestate *tupleStore, TupleDesc tupleDescripto
 }
 
 
+/*
+ * UserHasPermissionToViewStatsOf returns true if currentUserId can
+ * see backends of backendOwnedId.
+ *
+ * We follow the same approach with pg_stat_activity.
+ */
+static
+bool
+UserHasPermissionToViewStatsOf(Oid currentUserId, Oid backendOwnedId)
+{
+	if (has_privs_of_role(currentUserId, backendOwnedId))
+	{
+		return true;
+	}
+
+	if (is_member_of_role(currentUserId,
+#if PG_VERSION_NUM >= PG_VERSION_14
+						  ROLE_PG_READ_ALL_STATS))
+#else
+						  DEFAULT_ROLE_READ_ALL_STATS))
+#endif
+	{
+		return true;
+	}
+
+	return false;
+}
+
+
 /*
  * InitializeBackendManagement requests the necessary shared memory
  * from Postgres and sets up the shared memory startup hook.
@@ -649,8 +695,6 @@ UnSetDistributedTransactionId(void)
 	{
 		SpinLockAcquire(&MyBackendData->mutex);
 
-		MyBackendData->databaseId = 0;
-		MyBackendData->userId = 0;
 		MyBackendData->cancelledDueToDeadlock = false;
 		MyBackendData->transactionId.initiatorNodeIdentifier = 0;
 		MyBackendData->transactionId.transactionOriginator = false;
@@ -674,6 +718,8 @@ UnSetGlobalPID(void)
 		SpinLockAcquire(&MyBackendData->mutex);
 
 		MyBackendData->globalPID = 0;
+		MyBackendData->databaseId = 0;
+		MyBackendData->userId = 0;
 
 		SpinLockRelease(&MyBackendData->mutex);
 	}
@@ -755,13 +801,9 @@ AssignDistributedTransactionId(void)
 	uint64 nextTransactionNumber = pg_atomic_fetch_add_u64(transactionNumberSequence, 1);
 	int32 localGroupId = GetLocalGroupId();
 	TimestampTz currentTimestamp = GetCurrentTimestamp();
-	Oid userId = GetUserId();
 
 	SpinLockAcquire(&MyBackendData->mutex);
 
-	MyBackendData->databaseId = MyDatabaseId;
-	MyBackendData->userId = userId;
-
 	MyBackendData->transactionId.initiatorNodeIdentifier = localGroupId;
 	MyBackendData->transactionId.transactionOriginator = true;
 	MyBackendData->transactionId.transactionNumber = nextTransactionNumber;
@@ -793,9 +835,15 @@ AssignGlobalPID(void)
 		globalPID = ExtractGlobalPID(application_name);
 	}
 
+	Oid userId = GetUserId();
+
 	SpinLockAcquire(&MyBackendData->mutex);
+
 	MyBackendData->globalPID = globalPID;
 	MyBackendData->distributedCommandOriginator = distributedCommandOriginator;
+	MyBackendData->databaseId = MyDatabaseId;
+	MyBackendData->userId = userId;
+
 	SpinLockRelease(&MyBackendData->mutex);
 }
 

src/test/regress/expected/isolation_citus_dist_activity.out

@@ -42,7 +42,7 @@ query                                         |query_hostname  |query_hostport|d
 (1 row)
 
 step s3-view-worker:
-	SELECT query, query_hostname, query_hostport, distributed_query_host_name, distributed_query_host_port, state, wait_event_type, wait_event, usename, datname FROM citus_worker_stat_activity WHERE query NOT ILIKE '%pg_prepared_xacts%' AND query NOT ILIKE '%COMMIT%' ORDER BY query DESC;
+	SELECT query, query_hostname, query_hostport, distributed_query_host_name, distributed_query_host_port, state, wait_event_type, wait_event, usename, datname FROM citus_worker_stat_activity WHERE query NOT ILIKE '%pg_prepared_xacts%' AND query NOT ILIKE '%COMMIT%' AND backend_type = 'client backend' ORDER BY query DESC;
 
 query                                                                                                      |query_hostname|query_hostport|distributed_query_host_name|distributed_query_host_port|state              |wait_event_type|wait_event|usename |datname
 ---------------------------------------------------------------------
@@ -112,7 +112,7 @@ query                                        |query_hostname  |query_hostport|di
 (1 row)
 
 step s3-view-worker:
-	SELECT query, query_hostname, query_hostport, distributed_query_host_name, distributed_query_host_port, state, wait_event_type, wait_event, usename, datname FROM citus_worker_stat_activity WHERE query NOT ILIKE '%pg_prepared_xacts%' AND query NOT ILIKE '%COMMIT%' ORDER BY query DESC;
+	SELECT query, query_hostname, query_hostport, distributed_query_host_name, distributed_query_host_port, state, wait_event_type, wait_event, usename, datname FROM citus_worker_stat_activity WHERE query NOT ILIKE '%pg_prepared_xacts%' AND query NOT ILIKE '%COMMIT%' AND backend_type = 'client backend' ORDER BY query DESC;
 
 query                                                                     |query_hostname|query_hostport|distributed_query_host_name|distributed_query_host_port|state              |wait_event_type|wait_event|usename |datname
 ---------------------------------------------------------------------
@@ -176,7 +176,7 @@ query                                |query_hostname  |query_hostport|distribute
 (1 row)
 
 step s3-view-worker:
-	SELECT query, query_hostname, query_hostport, distributed_query_host_name, distributed_query_host_port, state, wait_event_type, wait_event, usename, datname FROM citus_worker_stat_activity WHERE query NOT ILIKE '%pg_prepared_xacts%' AND query NOT ILIKE '%COMMIT%' ORDER BY query DESC;
+	SELECT query, query_hostname, query_hostport, distributed_query_host_name, distributed_query_host_port, state, wait_event_type, wait_event, usename, datname FROM citus_worker_stat_activity WHERE query NOT ILIKE '%pg_prepared_xacts%' AND query NOT ILIKE '%COMMIT%' AND backend_type = 'client backend' ORDER BY query DESC;
 
 query                                                                        |query_hostname|query_hostport|distributed_query_host_name|distributed_query_host_port|state              |wait_event_type|wait_event|usename |datname
 ---------------------------------------------------------------------
@@ -243,7 +243,7 @@ query                                                   |query_hostname  |query_
 (1 row)
 
 step s3-view-worker:
-	SELECT query, query_hostname, query_hostport, distributed_query_host_name, distributed_query_host_port, state, wait_event_type, wait_event, usename, datname FROM citus_worker_stat_activity WHERE query NOT ILIKE '%pg_prepared_xacts%' AND query NOT ILIKE '%COMMIT%' ORDER BY query DESC;
+	SELECT query, query_hostname, query_hostport, distributed_query_host_name, distributed_query_host_port, state, wait_event_type, wait_event, usename, datname FROM citus_worker_stat_activity WHERE query NOT ILIKE '%pg_prepared_xacts%' AND query NOT ILIKE '%COMMIT%' AND backend_type = 'client backend' ORDER BY query DESC;
 
 query                                                                                                       |query_hostname|query_hostport|distributed_query_host_name|distributed_query_host_port|state              |wait_event_type|wait_event|usename |datname
 ---------------------------------------------------------------------

src/test/regress/expected/isolation_get_all_active_transactions.out

@@ -49,23 +49,32 @@ count
 (1 row)
 
 step s3-as-user-1:
- -- User should only be able to see its own transactions
+	-- Eventhough we change the user via SET ROLE, the backends' (e.g., s1/2-begin-insert)
+	-- userId (e.g., PG_PROC->userId) does not change, and hence none of the
+	-- transactions show up because here we are using test_user_1. This is a
+	-- limitation of isolation tester, we should be able to re-connect with
+	-- test_user_1 on s1/2-begin-insert to show that test_user_1 sees only its own processes
 	SET ROLE test_user_1;
 	SELECT count(*) FROM get_all_active_transactions() WHERE transaction_number != 0;
 	SELECT count(*) FROM get_global_active_transactions() WHERE transaction_number != 0;
 
 count
 ---------------------------------------------------------------------
-    1
+    0
 (1 row)
 
 count
 ---------------------------------------------------------------------
-    2
+    1
 (1 row)
 
 step s3-as-readonly:
- -- Other user should not see transactions
+	-- Eventhough we change the user via SET ROLE, the backends' (e.g., s1/2-begin-insert)
+	-- userId (e.g., PG_PROC->userId) does not change, and hence none of the
+	-- transactions show up because here we are using test_readonly. This is a
+	-- limitation of isolation tester, we should be able to re-connect with
+	-- test_readonly on s1/2-begin-insert to show that test_readonly sees only
+	-- its own processes
 	SET ROLE test_readonly;
 	SELECT count(*) FROM get_all_active_transactions() WHERE transaction_number != 0;
 	SELECT count(*) FROM get_global_active_transactions() WHERE transaction_number != 0;

src/test/regress/expected/isolation_replicate_reference_tables_to_coordinator.out

@@ -119,7 +119,8 @@ step s2-view-worker:
           query NOT ILIKE '%COMMIT%' AND
           query NOT ILIKE '%dump_local_%' AND
           query NOT ILIKE '%citus_internal_local_blocked_processes%' AND
-          query NOT ILIKE '%add_node%'
+          query NOT ILIKE '%add_node%' AND
+          backend_type = 'client backend'
     ORDER BY query, query_hostport DESC;
 
 query                                                                         |query_hostname|query_hostport|distributed_query_host_name|distributed_query_host_port|state              |wait_event_type|wait_event|usename |datname

src/test/regress/spec/isolation_citus_dist_activity.spec

@@ -89,7 +89,7 @@ step "s3-rollback"
 
 step "s3-view-worker"
 {
-	SELECT query, query_hostname, query_hostport, distributed_query_host_name, distributed_query_host_port, state, wait_event_type, wait_event, usename, datname FROM citus_worker_stat_activity WHERE query NOT ILIKE '%pg_prepared_xacts%' AND query NOT ILIKE '%COMMIT%' ORDER BY query DESC;
+	SELECT query, query_hostname, query_hostport, distributed_query_host_name, distributed_query_host_port, state, wait_event_type, wait_event, usename, datname FROM citus_worker_stat_activity WHERE query NOT ILIKE '%pg_prepared_xacts%' AND query NOT ILIKE '%COMMIT%' AND backend_type = 'client backend' ORDER BY query DESC;
 }
 
 // we prefer to sleep before "s2-view-dist" so that we can ensure

src/test/regress/spec/isolation_get_all_active_transactions.spec

@@ -77,7 +77,11 @@ step "s3-as-admin"
 
 step "s3-as-user-1"
 {
-	-- User should only be able to see its own transactions
+	-- Even though we change the user via SET ROLE, the backends' (e.g., s1/2-begin-insert)
+	-- userId (e.g., PG_PROC->userId) does not change, and hence none of the
+	-- transactions show up because here we are using test_user_1. This is a
+	-- limitation of isolation tester, we should be able to re-connect with
+	-- test_user_1 on s1/2-begin-insert to show that test_user_1 sees only its own processes
 	SET ROLE test_user_1;
 	SELECT count(*) FROM get_all_active_transactions() WHERE transaction_number != 0;
 	SELECT count(*) FROM get_global_active_transactions() WHERE transaction_number != 0;
@@ -85,7 +89,12 @@ step "s3-as-user-1"
 
 step "s3-as-readonly"
 {
-	-- Other user should not see transactions
+	-- Even though we change the user via SET ROLE, the backends' (e.g., s1/2-begin-insert)
+	-- userId (e.g., PG_PROC->userId) does not change, and hence none of the
+	-- transactions show up because here we are using test_readonly. This is a
+	-- limitation of isolation tester, we should be able to re-connect with
+	-- test_readonly on s1/2-begin-insert to show that test_readonly sees only
+	-- its own processes
 	SET ROLE test_readonly;
 	SELECT count(*) FROM get_all_active_transactions() WHERE transaction_number != 0;
 	SELECT count(*) FROM get_global_active_transactions() WHERE transaction_number != 0;

src/test/regress/spec/isolation_replicate_reference_tables_to_coordinator.spec

@@ -93,7 +93,8 @@ step "s2-view-worker"
           query NOT ILIKE '%COMMIT%' AND
           query NOT ILIKE '%dump_local_%' AND
           query NOT ILIKE '%citus_internal_local_blocked_processes%' AND
-          query NOT ILIKE '%add_node%'
+          query NOT ILIKE '%add_node%' AND
+          backend_type = 'client backend'
     ORDER BY query, query_hostport DESC;
 }