From 061ab76370e7d050f107f85b344ac7095d7252b9 Mon Sep 17 00:00:00 2001
From: Anthonin Bonnefoy <anthonin.bonnefoy@datadoghq.com>
Date: Fri, 26 Jun 2026 11:06:16 +0200
Subject: Add PGTRACE env to enable protocol tracing in libpq

Add 2 env vars, PGTRACE and PGTRACEFLAGS. When PGTRACE is set, PQtrace
will be set when connection options are processed. This allows to trace
protocol messages like startup messages and BackendKeyData. This also
allows enabling protocol tracing on any application using libpq.

When a cancel connection is created with PQcancelCreate, we copy tracing
fd and trace flags to allow the cancel requests to be traced. Since
writes to the trace_file are now done line by line, we can share a fd
between the source connection and the cancel connection.
---
 doc/src/sgml/libpq.sgml           | 48 +++++++++++++++++++++++
 src/interfaces/libpq/fe-cancel.c  |  6 +++
 src/interfaces/libpq/fe-connect.c | 65 ++++++++++++++++++++++++++++++-
 src/interfaces/libpq/libpq-int.h  |  4 ++
 4 files changed, 122 insertions(+), 1 deletion(-)

diff --git a/doc/src/sgml/libpq.sgml b/doc/src/sgml/libpq.sgml
index 7d3c3bb66d8..b2157860a40 100644
--- a/doc/src/sgml/libpq.sgml
+++ b/doc/src/sgml/libpq.sgml
@@ -2424,6 +2424,34 @@ postgresql://%2Fvar%2Flib%2Fpostgresql/dbname
       </listitem>
      </varlistentry>
 
+     <varlistentry id="libpq-connect-trace-file" xreflabel="trace_file">
+      <term><literal>trace_file</literal></term>
+      <listitem>
+       <para>
+        If set, enables tracing of the client/server protocol exchange to
+        the named file. The special value <literal>-</literal> sends the
+        trace to <literal>stdout</literal>. Multiple connections using
+        the same trace file is currently not supported.
+       </para>
+      </listitem>
+     </varlistentry>
+
+     <varlistentry id="libpq-connect-trace-flags" xreflabel="trace_flags">
+      <term><literal>trace_flags</literal></term>
+      <listitem>
+       <para>
+        Specifies the trace flags applied when
+        <xref linkend="libpq-connect-trace-file"/> is set.  The value is
+        an integer that is the bitwise OR of the flag values accepted by
+        <xref linkend="libpq-PQsetTraceFlags"/>:
+        <literal>1</literal> for <literal>PQTRACE_SUPPRESS_TIMESTAMPS</literal>,
+        <literal>2</literal> for <literal>PQTRACE_REGRESS_MODE</literal>.
+        The default is <literal>0</literal>.  This option has no effect
+        unless <literal>trace_file</literal> is also set.
+       </para>
+      </listitem>
+     </varlistentry>
+
      <varlistentry id="libpq-connect-load-balance-hosts" xreflabel="load_balance_hosts">
       <term><literal>load_balance_hosts</literal></term>
       <listitem>
@@ -9454,6 +9482,26 @@ myEventProc(PGEventId evtId, void *evtInfo, void *passThrough)
       linkend="libpq-connect-oauth-ca-file"/> connection parameter.
      </para>
     </listitem>
+
+    <listitem>
+     <para>
+      <indexterm>
+       <primary><envar>PGTRACE</envar></primary>
+      </indexterm>
+      <envar>PGTRACE</envar> behaves the same as the <xref
+      linkend="libpq-connect-trace-file"/> connection parameter.
+     </para>
+    </listitem>
+
+    <listitem>
+     <para>
+      <indexterm>
+       <primary><envar>PGTRACEFLAGS</envar></primary>
+      </indexterm>
+      <envar>PGTRACEFLAGS</envar> behaves the same as the <xref
+      linkend="libpq-connect-trace-flags"/> connection parameter.
+     </para>
+    </listitem>
    </itemizedlist>
   </para>
 
diff --git a/src/interfaces/libpq/fe-cancel.c b/src/interfaces/libpq/fe-cancel.c
index 4b5945979c4..c213ac4c08f 100644
--- a/src/interfaces/libpq/fe-cancel.c
+++ b/src/interfaces/libpq/fe-cancel.c
@@ -98,6 +98,12 @@ PQcancelCreate(PGconn *conn)
 	 */
 	cancelConn->cancelRequest = true;
 
+	/*
+	 * Copy Pfdebug and trace flags from the original connection.
+	 */
+	cancelConn->Pfdebug = conn->Pfdebug;
+	cancelConn->traceFlags = conn->traceFlags;
+
 	if (!pqCopyPGconn(conn, cancelConn))
 		return (PGcancelConn *) cancelConn;
 
diff --git a/src/interfaces/libpq/fe-connect.c b/src/interfaces/libpq/fe-connect.c
index 38422becc48..4c1f8597ddf 100644
--- a/src/interfaces/libpq/fe-connect.c
+++ b/src/interfaces/libpq/fe-connect.c
@@ -385,6 +385,15 @@ static const internalPQconninfoOption PQconninfoOptions[] = {
 		"Target-Session-Attrs", "", 15, /* sizeof("prefer-standby") = 15 */
 	offsetof(struct pg_conn, target_session_attrs)},
 
+	{"trace_file", "PGTRACE",
+		NULL, NULL,
+		"Trace-File", "D", 64,
+	offsetof(struct pg_conn, trace_file_str)},
+
+	{"trace_flags", "PGTRACEFLAGS", NULL, NULL,
+		"Trace-Flags", "D", 1,
+	offsetof(struct pg_conn, trace_flags_str)},
+
 	{"load_balance_hosts", "PGLOADBALANCEHOSTS",
 		DefaultLoadBalanceHosts, NULL,
 		"Load-Balance-Hosts", "", 8,	/* sizeof("disable") = 8 */
@@ -1257,6 +1266,7 @@ bool
 pqConnectOptions2(PGconn *conn)
 {
 	int			i;
+	int			trace_flags = 0;
 
 	/*
 	 * Allocate memory for details about each host to which we might possibly
@@ -2057,6 +2067,44 @@ pqConnectOptions2(PGconn *conn)
 		conn->scram_client_key_len = len;
 	}
 
+	/*
+	 * If the connection already has a Pfdebug opened, use it
+	 */
+	if (conn->trace_file_str && !conn->Pfdebug)
+	{
+		if (strcmp(conn->trace_file_str, "-") == 0)
+		{
+			conn->trace_file = stdout;
+		}
+
+		/*
+		 * If the connection is created for a cancel request, and trace_file
+		 * is not stdout, don't try to trace it as the source connection has
+		 * already an opened fd on this trace file
+		 */
+		else if (!conn->cancelRequest)
+		{
+			conn->trace_file = fopen(conn->trace_file_str, "w");
+			if (conn->trace_file == NULL)
+			{
+				libpq_append_conn_error(conn, "could not open trace file \"%s\": %m", conn->trace_file_str);
+				return false;
+			}
+			/* Make it line-buffered */
+			setvbuf(conn->trace_file, NULL, PG_IOLBF, 0);
+		}
+		PQtrace(conn, conn->trace_file);
+
+		if (conn->trace_flags_str)
+		{
+			if (!pqParseIntParam(conn->trace_flags_str, &trace_flags, conn,
+								 "trace_flags"))
+				return false;
+		}
+
+		PQsetTraceFlags(conn, trace_flags);
+	}
+
 	if (conn->scram_server_key)
 	{
 		int			len;
@@ -5043,6 +5091,7 @@ pqMakeEmptyPGconn(void)
 	conn->sock = PGINVALID_SOCKET;
 	conn->altsock = PGINVALID_SOCKET;
 	conn->Pfdebug = NULL;
+	conn->trace_file = NULL;
 
 	/*
 	 * We try to send at least 8K at a time, which is the usual size of pipe
@@ -5164,7 +5213,21 @@ freePGconn(PGconn *conn)
 	free(conn->oauth_client_secret);
 	free(conn->oauth_ca_file);
 	free(conn->oauth_scope);
-	/* Note that conn->Pfdebug is not ours to close or free */
+	free(conn->trace_file_str);
+	free(conn->trace_flags_str);
+
+	/* Untrace the connection to flush Pfdebug */
+	PQuntrace(conn);
+
+	/*
+	 * Note that conn->Pfdebug is not ours to close or free if provided
+	 * externally. However, if tracing was enabled through PGTRACE, we need to
+	 * close conn->trace_file.
+	 */
+	if (conn->trace_file && conn->trace_file != stdout)
+	{
+		fclose(conn->trace_file);
+	}
 	free(conn->events);
 	pqReleaseConnHosts(conn);
 	free(conn->connip);
diff --git a/src/interfaces/libpq/libpq-int.h b/src/interfaces/libpq/libpq-int.h
index 3f921207a14..8c926483b43 100644
--- a/src/interfaces/libpq/libpq-int.h
+++ b/src/interfaces/libpq/libpq-int.h
@@ -449,6 +449,10 @@ struct pg_conn
 	char	   *oauth_ca_file;	/* CA file path */
 	bool		oauth_want_retry;	/* should we retry on failure? */
 
+	FILE	   *trace_file;
+	char	   *trace_file_str;
+	char	   *trace_flags_str;
+
 	/* Optional file to write trace info to */
 	FILE	   *Pfdebug;
 	int			traceFlags;
-- 
2.54.0

