From db52f5949884e57769b7ba7fc56be64c47ef1fa1 Mon Sep 17 00:00:00 2001
From: "Chao Li (Evan)" <lic@highgo.com>
Date: Wed, 20 May 2026 14:57:25 +0800
Subject: [PATCH v1] Set notice receiver before libpq connection startup

Some callers of libpqsrv_connect() and libpqsrv_connect_params()
install libpqsrv_notice_receiver after the connection has been
established. However, notices may also be received while
libpqsrv_connect_internal() is still processing the asynchronous
connection startup.

Add variants of these helpers that install a notice receiver immediately
after PQconnectStart() or PQconnectStartParams(), before entering
libpqsrv_connect_internal(). Use them in dblink, postgres_fdw, and
libpqwalreceiver, so notices emitted during connection startup are
handled in the same way as notices received after the connection is
established.

Author: Chao Li <lic@highgo.com>
Reviewed-by:
Discussion: https://postgr.es/m/
---
 contrib/dblink/dblink.c                       | 15 ++++---
 contrib/postgres_fdw/connection.c             | 11 +++--
 .../libpqwalreceiver/libpqwalreceiver.c       | 11 +++--
 src/include/libpq/libpq-be-fe-helpers.h       | 42 ++++++++++++++-----
 4 files changed, 49 insertions(+), 30 deletions(-)

diff --git a/contrib/dblink/dblink.c b/contrib/dblink/dblink.c
index d843eee7e97..2093aef990f 100644
--- a/contrib/dblink/dblink.c
+++ b/contrib/dblink/dblink.c
@@ -222,7 +222,10 @@ dblink_get_conn(char *conname_or_str,
 			dblink_we_get_conn = WaitEventExtensionNew("DblinkGetConnect");
 
 		/* OK to make connection */
-		conn = libpqsrv_connect(connstr, dblink_we_get_conn);
+		conn = libpqsrv_connect_with_notice_receiver(connstr,
+													 dblink_we_get_conn,
+													 libpqsrv_notice_receiver,
+													 "received message via remote connection");
 
 		if (PQstatus(conn) == CONNECTION_BAD)
 		{
@@ -235,9 +238,6 @@ dblink_get_conn(char *conname_or_str,
 					 errdetail_internal("%s", msg)));
 		}
 
-		PQsetNoticeReceiver(conn, libpqsrv_notice_receiver,
-							"received message via remote connection");
-
 		dblink_security_check(conn, NULL, connstr);
 		if (PQclientEncoding(conn) != GetDatabaseEncoding())
 			PQsetClientEncoding(conn, GetDatabaseEncodingName());
@@ -321,7 +321,9 @@ dblink_connect(PG_FUNCTION_ARGS)
 	}
 
 	/* OK to make connection */
-	conn = libpqsrv_connect(connstr, dblink_we_connect);
+	conn = libpqsrv_connect_with_notice_receiver(connstr, dblink_we_connect,
+												 libpqsrv_notice_receiver,
+												 "received message via remote connection");
 
 	if (PQstatus(conn) == CONNECTION_BAD)
 	{
@@ -336,9 +338,6 @@ dblink_connect(PG_FUNCTION_ARGS)
 				 errdetail_internal("%s", msg)));
 	}
 
-	PQsetNoticeReceiver(conn, libpqsrv_notice_receiver,
-						"received message via remote connection");
-
 	/* check password actually used if not superuser */
 	dblink_security_check(conn, connname, connstr);
 
diff --git a/contrib/postgres_fdw/connection.c b/contrib/postgres_fdw/connection.c
index 3d2a8d0519d..0278f4a7cea 100644
--- a/contrib/postgres_fdw/connection.c
+++ b/contrib/postgres_fdw/connection.c
@@ -646,9 +646,11 @@ connect_pg_server(ForeignServer *server, UserMapping *user)
 			pgfdw_we_connect = WaitEventExtensionNew("PostgresFdwConnect");
 
 		/* OK to make connection */
-		conn = libpqsrv_connect_params(keywords, values,
-									   false,	/* expand_dbname */
-									   pgfdw_we_connect);
+		conn = libpqsrv_connect_params_with_notice_receiver(keywords, values,
+															false,	/* expand_dbname */
+															pgfdw_we_connect,
+															libpqsrv_notice_receiver,
+															"received message via remote connection");
 
 		if (!conn || PQstatus(conn) != CONNECTION_OK)
 			ereport(ERROR,
@@ -657,9 +659,6 @@ connect_pg_server(ForeignServer *server, UserMapping *user)
 							server->servername),
 					 errdetail_internal("%s", pchomp(PQerrorMessage(conn)))));
 
-		PQsetNoticeReceiver(conn, libpqsrv_notice_receiver,
-							"received message via remote connection");
-
 		/* Perform post-connection security checks. */
 		pgfdw_security_check(keywords, values, user, conn);
 
diff --git a/src/backend/replication/libpqwalreceiver/libpqwalreceiver.c b/src/backend/replication/libpqwalreceiver/libpqwalreceiver.c
index 9f04c9ed25d..cc5f77d53f7 100644
--- a/src/backend/replication/libpqwalreceiver/libpqwalreceiver.c
+++ b/src/backend/replication/libpqwalreceiver/libpqwalreceiver.c
@@ -223,9 +223,11 @@ libpqrcv_connect(const char *conninfo, bool replication, bool logical,
 
 	conn = palloc0_object(WalReceiverConn);
 	conn->streamConn =
-		libpqsrv_connect_params(keys, vals,
-								 /* expand_dbname = */ true,
-								WAIT_EVENT_LIBPQWALRECEIVER_CONNECT);
+		libpqsrv_connect_params_with_notice_receiver(keys, vals,
+													  /* expand_dbname = */ true,
+													 WAIT_EVENT_LIBPQWALRECEIVER_CONNECT,
+													 libpqsrv_notice_receiver,
+													 "received message via replication");
 
 	if (options_val != NULL)
 		pfree(options_val);
@@ -245,9 +247,6 @@ libpqrcv_connect(const char *conninfo, bool replication, bool logical,
 				 errhint("Target server's authentication method must be changed, or set password_required=false in the subscription parameters.")));
 	}
 
-	PQsetNoticeReceiver(conn->streamConn, libpqsrv_notice_receiver,
-						"received message via replication");
-
 	/*
 	 * Set always-secure search path for the cases where the connection is
 	 * used to run SQL queries, so malicious users can't get control.
diff --git a/src/include/libpq/libpq-be-fe-helpers.h b/src/include/libpq/libpq-be-fe-helpers.h
index 85d8b63f019..0571bb1fe50 100644
--- a/src/include/libpq/libpq-be-fe-helpers.h
+++ b/src/include/libpq/libpq-be-fe-helpers.h
@@ -43,6 +43,23 @@ static inline void libpqsrv_connect_internal(PGconn *conn, uint32 wait_event_inf
 static inline PGresult *libpqsrv_get_result_last(PGconn *conn, uint32 wait_event_info);
 static inline PGresult *libpqsrv_get_result(PGconn *conn, uint32 wait_event_info);
 
+static inline PGconn *
+libpqsrv_connect_with_notice_receiver(const char *conninfo,
+									  uint32 wait_event_info,
+									  PQnoticeReceiver proc, void *arg)
+{
+	PGconn	   *conn = NULL;
+
+	libpqsrv_connect_prepare();
+
+	conn = PQconnectStart(conninfo);
+	if (conn != NULL && proc != NULL)
+		PQsetNoticeReceiver(conn, proc, arg);
+
+	libpqsrv_connect_internal(conn, wait_event_info);
+
+	return conn;
+}
 
 /*
  * PQconnectdb() wrapper that reserves a file descriptor and processes
@@ -54,12 +71,24 @@ static inline PGresult *libpqsrv_get_result(PGconn *conn, uint32 wait_event_info
  */
 static inline PGconn *
 libpqsrv_connect(const char *conninfo, uint32 wait_event_info)
+{
+	return libpqsrv_connect_with_notice_receiver(conninfo, wait_event_info, NULL, NULL);
+}
+
+static inline PGconn *
+libpqsrv_connect_params_with_notice_receiver(const char *const *keywords,
+											 const char *const *values,
+											 int expand_dbname,
+											 uint32 wait_event_info,
+											 PQnoticeReceiver proc, void *arg)
 {
 	PGconn	   *conn = NULL;
 
 	libpqsrv_connect_prepare();
 
-	conn = PQconnectStart(conninfo);
+	conn = PQconnectStartParams(keywords, values, expand_dbname);
+	if (conn != NULL && proc != NULL)
+		PQsetNoticeReceiver(conn, proc, arg);
 
 	libpqsrv_connect_internal(conn, wait_event_info);
 
@@ -76,15 +105,8 @@ libpqsrv_connect_params(const char *const *keywords,
 						int expand_dbname,
 						uint32 wait_event_info)
 {
-	PGconn	   *conn = NULL;
-
-	libpqsrv_connect_prepare();
-
-	conn = PQconnectStartParams(keywords, values, expand_dbname);
-
-	libpqsrv_connect_internal(conn, wait_event_info);
-
-	return conn;
+	return libpqsrv_connect_params_with_notice_receiver(keywords, values, expand_dbname,
+														wait_event_info, NULL, NULL);
 }
 
 /*
-- 
2.50.1 (Apple Git-155)

