From 34df853bdae23c862350160fcd0edaa0c49d1da4 Mon Sep 17 00:00:00 2001
From: Jelte Fennema <jelte.fennema@microsoft.com>
Date: Tue, 10 Jan 2023 16:21:57 +0100
Subject: [PATCH] Fix bug introduced by #6412 (#6590)

In #6412 I made a change to not re-assign the global PID if it was
already set. This inadvertently introduced a regression where `userId`
and `databaseId` would not be set on the backend data when the global
PID was assigned in the authentication hook.

This fixes it by doing two things:
1. Removing `userId` from `BackendData`, since it's not used anywhere
   anyway.
2. Move assignment of `databaseId` to dedicated
   `SetBackendDataDatabaseId` function, that isn't a no-op when global
   pid is already set.

Since #6412 is not released yet this does not need a description.
---
 src/backend/distributed/shared_library_init.c |  2 ++
 .../distributed/transaction/backend_data.c    | 21 ++++++++++++++-----
 src/include/distributed/backend_data.h        |  2 +-
 3 files changed, 19 insertions(+), 6 deletions(-)

diff --git a/src/backend/distributed/shared_library_init.c b/src/backend/distributed/shared_library_init.c
index c15268056..f39954cd6 100644
--- a/src/backend/distributed/shared_library_init.c
+++ b/src/backend/distributed/shared_library_init.c
@@ -678,6 +678,8 @@ StartupCitusBackend(void)
 	 * the gpid from the application_name.
 	 */
 	AssignGlobalPID();
+
+	SetBackendDataDatabaseId();
 	RegisterConnectionCleanup();
 }
 
diff --git a/src/backend/distributed/transaction/backend_data.c b/src/backend/distributed/transaction/backend_data.c
index 3daebc6e8..a93d3ab51 100644
--- a/src/backend/distributed/transaction/backend_data.c
+++ b/src/backend/distributed/transaction/backend_data.c
@@ -765,7 +765,6 @@ UnSetGlobalPID(void)
 
 		MyBackendData->globalPID = 0;
 		MyBackendData->databaseId = 0;
-		MyBackendData->userId = 0;
 		MyBackendData->distributedCommandOriginator = false;
 
 		SpinLockRelease(&MyBackendData->mutex);
@@ -923,19 +922,31 @@ 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);
 }
 
 
+/*
+ * SetBackendDataDatabaseId sets the databaseId in the backend data.
+ *
+ * NOTE: this needs to be run after the auth hook, because in the auth hook the
+ * database is not known yet.
+ */
+void
+SetBackendDataDatabaseId(void)
+{
+	Assert(MyDatabaseId != InvalidOid);
+	SpinLockAcquire(&MyBackendData->mutex);
+	MyBackendData->databaseId = MyDatabaseId;
+	SpinLockRelease(&MyBackendData->mutex);
+}
+
+
 /*
  * SetBackendDataDistributedCommandOriginator is used to set the distributedCommandOriginator
  * field on MyBackendData.
diff --git a/src/include/distributed/backend_data.h b/src/include/distributed/backend_data.h
index b0846c8a7..28b53482f 100644
--- a/src/include/distributed/backend_data.h
+++ b/src/include/distributed/backend_data.h
@@ -36,7 +36,6 @@
 typedef struct BackendData
 {
 	Oid databaseId;
-	Oid userId;
 	slock_t mutex;
 	bool cancelledDueToDeadlock;
 	uint64 globalPID;
@@ -60,6 +59,7 @@ extern void AssignDistributedTransactionId(void);
 extern void AssignGlobalPID(void);
 extern void SetBackendDataGlobalPID(uint64 globalPID);
 extern uint64 GetGlobalPID(void);
+extern void SetBackendDataDatabaseId(void);
 extern void SetBackendDataDistributedCommandOriginator(bool
 													   distributedCommandOriginator);
 extern uint64 ExtractGlobalPID(const char *applicationName);