From 67635f896c52cb0d9be8b2411a794e927e6da0dd Mon Sep 17 00:00:00 2001
From: Jack Christensen <jack@jackchristensen.com>
Date: Sat, 21 May 2022 17:30:47 -0500
Subject: [PATCH] Fix output to include message size and add some docs

---
 pgproto3/doc.go        |  7 +++++
 pgproto3/trace.go      |  3 ++
 pgproto3/trace_test.go | 64 +++++++++++++++++++++---------------------
 3 files changed, 42 insertions(+), 32 deletions(-)

diff --git a/pgproto3/doc.go b/pgproto3/doc.go
index 8226dc98..e0e1cf87 100644
--- a/pgproto3/doc.go
+++ b/pgproto3/doc.go
@@ -1,4 +1,11 @@
 // Package pgproto3 is a encoder and decoder of the PostgreSQL wire protocol version 3.
 //
+// The primary interfaces are Frontend and Backend. They correspond to a client and server respectively. Messages are
+// sent with Send (or a specialized Send variant). Messages are automatically bufferred to minimize small writes. Call
+// Flush to ensure a message has actually been sent.
+//
+// The Trace method of Frontend and Backend can be used to examine the wire-level message traffic. It outputs in a
+// similar format to the PQtrace function in libpq.
+//
 // See https://www.postgresql.org/docs/current/protocol-message-formats.html for meanings of the different messages.
 package pgproto3
diff --git a/pgproto3/trace.go b/pgproto3/trace.go
index 704b2ee7..d3edc4aa 100644
--- a/pgproto3/trace.go
+++ b/pgproto3/trace.go
@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"fmt"
 	"io"
+	"strconv"
 	"strings"
 	"time"
 )
@@ -399,6 +400,8 @@ func (t *tracer) beginTrace(sender byte, encodedLen int32, msgType string) {
 	t.buf.WriteByte(sender)
 	t.buf.WriteByte('\t')
 	t.buf.WriteString(msgType)
+	t.buf.WriteByte('\t')
+	t.buf.WriteString(strconv.FormatInt(int64(encodedLen), 10))
 }
 
 func (t *tracer) finishTrace() {
diff --git a/pgproto3/trace_test.go b/pgproto3/trace_test.go
index a4057008..0ace057b 100644
--- a/pgproto3/trace_test.go
+++ b/pgproto3/trace_test.go
@@ -40,38 +40,38 @@ func TestTrace(t *testing.T) {
 	result := conn.ExecParams(ctx, "select n from generate_series(1,5) n", nil, nil, nil, nil).Read()
 	require.NoError(t, result.Err)
 
-	expected := `F	StartupMessage
-B	AuthenticationOk
-B	ParameterStatus	 "application_name" ""
-B	ParameterStatus	 "client_encoding" "UTF8"
-B	ParameterStatus	 "DateStyle" "ISO, MDY"
-B	ParameterStatus	 "default_transaction_read_only" "off"
-B	ParameterStatus	 "in_hot_standby" "off"
-B	ParameterStatus	 "integer_datetimes" "on"
-B	ParameterStatus	 "IntervalStyle" "postgres"
-B	ParameterStatus	 "is_superuser" "on"
-B	ParameterStatus	 "server_encoding" "UTF8"
-B	ParameterStatus	 "server_version" "14.3"
-B	ParameterStatus	 "session_authorization" "jack"
-B	ParameterStatus	 "standard_conforming_strings" "on"
-B	ParameterStatus	 "TimeZone" "America/Chicago"
-B	BackendKeyData	 NNNN NNNN
-B	ReadyForQuery	 I
-F	Parse	 "" "select n from generate_series(1,5) n" 0
-F	Bind	 "" "" 0 0 0
-F	Describe	 P ""
-F	Execute	 "" 0
-F	Sync
-B	ParseComplete
-B	BindComplete
-B	RowDescription	 1 "n" 0 0 23 4 -1 0
-B	DataRow	 1 1 '1'
-B	DataRow	 1 1 '2'
-B	DataRow	 1 1 '3'
-B	DataRow	 1 1 '4'
-B	DataRow	 1 1 '5'
-B	CommandComplete	 "SELECT 5"
-B	ReadyForQuery	 I
+	expected := `F	StartupMessage	37
+B	AuthenticationOk	9
+B	ParameterStatus	23	 "application_name" ""
+B	ParameterStatus	26	 "client_encoding" "UTF8"
+B	ParameterStatus	24	 "DateStyle" "ISO, MDY"
+B	ParameterStatus	39	 "default_transaction_read_only" "off"
+B	ParameterStatus	24	 "in_hot_standby" "off"
+B	ParameterStatus	26	 "integer_datetimes" "on"
+B	ParameterStatus	28	 "IntervalStyle" "postgres"
+B	ParameterStatus	21	 "is_superuser" "on"
+B	ParameterStatus	26	 "server_encoding" "UTF8"
+B	ParameterStatus	25	 "server_version" "14.3"
+B	ParameterStatus	32	 "session_authorization" "jack"
+B	ParameterStatus	36	 "standard_conforming_strings" "on"
+B	ParameterStatus	30	 "TimeZone" "America/Chicago"
+B	BackendKeyData	13	 NNNN NNNN
+B	ReadyForQuery	6	 I
+F	Parse	45	 "" "select n from generate_series(1,5) n" 0
+F	Bind	13	 "" "" 0 0 0
+F	Describe	7	 P ""
+F	Execute	10	 "" 0
+F	Sync	5
+B	ParseComplete	5
+B	BindComplete	5
+B	RowDescription	27	 1 "n" 0 0 23 4 -1 0
+B	DataRow	12	 1 1 '1'
+B	DataRow	12	 1 1 '2'
+B	DataRow	12	 1 1 '3'
+B	DataRow	12	 1 1 '4'
+B	DataRow	12	 1 1 '5'
+B	CommandComplete	14	 "SELECT 5"
+B	ReadyForQuery	6	 I
 `
 
 	require.Equal(t, expected, traceOutput.String())