From ab1c798a1f72dff884819c238055251b65cc9d2b Mon Sep 17 00:00:00 2001
From: Tristan Partin <tristan.partin@databricks.com>
Date: Wed, 26 Nov 2025 10:23:34 -0600
Subject: [PATCH] Move NEON_AUTH_TOKEN to a builtin GUC

This environment variable is used as the password to connect to another
postgres instance as the walreceiver. The purpose of moving to a GUC is
so that we can reload the storage auth token periodically.

Signed-off-by: Tristan Partin <tristan.partin@databricks.com>
---
 .../libpqwalreceiver/libpqwalreceiver.c         | 14 ++++++++------
 src/backend/replication/walreceiver.c           | 17 +++++++++++++++++
 src/backend/utils/misc/guc.c                    | 11 +++++++++++
 src/include/replication/walreceiver.h           |  3 +++
 4 files changed, 39 insertions(+), 6 deletions(-)

diff --git a/src/backend/replication/libpqwalreceiver/libpqwalreceiver.c b/src/backend/replication/libpqwalreceiver/libpqwalreceiver.c
index d64a95ded97..cccd0615484 100644
--- a/src/backend/replication/libpqwalreceiver/libpqwalreceiver.c
+++ b/src/backend/replication/libpqwalreceiver/libpqwalreceiver.c
@@ -131,7 +131,6 @@ libpqrcv_connect(const char *conninfo, bool logical, const char *appname,
 /* BEGIN_NEON */
 	const char *keys[7];
 	const char *vals[7];
-	char * neon_auth_token = NULL;
 /* END_NEON */
 	int			i = 0;
 
@@ -143,18 +142,21 @@ libpqrcv_connect(const char *conninfo, bool logical, const char *appname,
 	vals[i] = conninfo;
 
 /* BEGIN_NEON */
+	/*
+	 * We use neon_storage_token for the password because conninfo strings are
+	 * limited to MAXCONNINFO in length. Our tokens encode Unity Catalog
+	 * permissions, so they can be quite lengthy.
+	 */
 	if (pg_strcasecmp(appname, "walreceiver") == 0)
 	{
-		neon_auth_token = getenv("NEON_AUTH_TOKEN");
-		if (neon_auth_token != NULL)
+		if (neon_storage_token[0] != '\0')
 		{
-			elog(LOG, "Use NEON_AUTH_TOKEN to connect");
 			keys[++i] = "password";
-			vals[i] = neon_auth_token;
+			vals[i] = neon_storage_token;
 		}
 		else
 		{
-			elog(LOG, "NEON_AUTH_TOKEN is undefined in the environment");
+			elog(LOG, "no storage token set");
 		}
 	}
 /* END_NEON */
diff --git a/src/backend/replication/walreceiver.c b/src/backend/replication/walreceiver.c
index 5c53a3c1086..a51f717917d 100644
--- a/src/backend/replication/walreceiver.c
+++ b/src/backend/replication/walreceiver.c
@@ -90,6 +90,7 @@
 int			wal_receiver_status_interval;
 int			wal_receiver_timeout;
 bool		hot_standby_feedback;
+char	   *neon_storage_token;
 
 /* libpqwalreceiver connection */
 static WalReceiverConn *wrconn = NULL;
@@ -1339,6 +1340,22 @@ WalRcvGetStateString(WalRcvState state)
 	return "UNKNOWN";
 }
 
+/*
+ * We currently grant the privileged role pg_monitor, which implies
+ * pg_read_all_settings. Until we fix that, let's just redact the content unless
+ * the user requesting the value is a superuser.
+ *
+ * See: https://databricks.atlassian.net/browse/LKB-7128
+ */
+const char *
+show_neon_storage_token(void)
+{
+	if (superuser())
+		return neon_storage_token;
+
+	return "**********";
+}
+
 /*
  * Returns activity of WAL receiver, including pid, state and xlog locations
  * received from the WAL sender of another server.
diff --git a/src/backend/utils/misc/guc.c b/src/backend/utils/misc/guc.c
index c6ee811ba26..dd7d129bce0 100644
--- a/src/backend/utils/misc/guc.c
+++ b/src/backend/utils/misc/guc.c
@@ -4801,6 +4801,17 @@ static struct config_string ConfigureNamesString[] =
 		check_restrict_nonsystem_relation_kind, assign_restrict_nonsystem_relation_kind, NULL
 	},
 
+	{
+		{"neon_storage_token", PGC_POSTMASTER, REPLICATION_STANDBY,
+			"Authentication token for Neon storage",
+			NULL,
+			GUC_NO_SHOW_ALL | GUC_NO_RESET_ALL | GUC_NOT_IN_SAMPLE | GUC_SUPERUSER_ONLY
+		},
+		&neon_storage_token,
+		"",
+		NULL, NULL, show_neon_storage_token
+	},
+
 	/* End-of-list marker */
 	{
 		{NULL, 0, 0, NULL, NULL}, NULL, NULL, NULL, NULL, NULL
diff --git a/src/include/replication/walreceiver.h b/src/include/replication/walreceiver.h
index 81184aa92f3..2b473ff0f64 100644
--- a/src/include/replication/walreceiver.h
+++ b/src/include/replication/walreceiver.h
@@ -28,6 +28,7 @@
 extern PGDLLIMPORT int wal_receiver_status_interval;
 extern PGDLLIMPORT int wal_receiver_timeout;
 extern PGDLLIMPORT bool hot_standby_feedback;
+extern PGDLLIMPORT char *neon_storage_token;
 
 /*
  * MAXCONNINFO: maximum size of a connection string.
@@ -454,6 +455,8 @@ walrcv_clear_result(WalRcvExecResult *walres)
 extern void WalReceiverMain(void) pg_attribute_noreturn();
 extern void ProcessWalRcvInterrupts(void);
 
+extern const char *show_neon_storage_token(void);
+
 /* prototypes for functions in walreceiverfuncs.c */
 extern Size WalRcvShmemSize(void);
 extern void WalRcvShmemInit(void);