From 0280035860362f3ce2dde1d1d2ba4cee4f6d2aae Mon Sep 17 00:00:00 2001
From: "mykyta.oleksiienko" <mykyta.oleksiienko@gmail.com>
Date: Wed, 24 Jul 2024 11:23:58 +0300
Subject: [PATCH] Adding support system_traces.coordinator_port and source_port

---
 session.go | 40 +++++++++++++++++++++-------------------
 1 file changed, 21 insertions(+), 19 deletions(-)

diff --git a/session.go b/session.go
index 0eac4cf0e..1149bd9a2 100644
--- a/session.go
+++ b/session.go
@@ -2111,14 +2111,15 @@ func NewTraceWriter(session *Session, w io.Writer) Tracer {
 
 func (t *traceWriter) Trace(traceId []byte) {
 	var (
-		coordinator string
-		duration    int
+		coordinator     string
+		duration        int
+		coordinatorPort int
 	)
-	iter := t.session.control.query(`SELECT coordinator, duration
-			FROM system_traces.sessions
-			WHERE session_id = ?`, traceId)
+	iter := t.session.control.query(`SELECT coordinator, duration, coordinator_port
+      FROM system_traces.sessions
+      WHERE session_id = ?`, traceId)
 
-	iter.Scan(&coordinator, &duration)
+	iter.Scan(&coordinator, &duration, &coordinatorPort)
 	if err := iter.Close(); err != nil {
 		t.mu.Lock()
 		fmt.Fprintln(t.w, "Error:", err)
@@ -2127,26 +2128,27 @@ func (t *traceWriter) Trace(traceId []byte) {
 	}
 
 	var (
-		timestamp time.Time
-		activity  string
-		source    string
-		elapsed   int
-		thread    string
+		timestamp  time.Time
+		activity   string
+		source     string
+		elapsed    int
+		thread     string
+		sourcePort int
 	)
 
 	t.mu.Lock()
 	defer t.mu.Unlock()
 
-	fmt.Fprintf(t.w, "Tracing session %016x (coordinator: %s, duration: %v):\n",
-		traceId, coordinator, time.Duration(duration)*time.Microsecond)
+	fmt.Fprintf(t.w, "Tracing session %016x (coordinator: %s, duration: %v, coordinator_port: %d):\n",
+		traceId, coordinator, time.Duration(duration)*time.Microsecond, coordinatorPort)
 
-	iter = t.session.control.query(`SELECT event_id, activity, source, source_elapsed, thread
-			FROM system_traces.events
-			WHERE session_id = ?`, traceId)
+	iter = t.session.control.query(`SELECT event_id, activity, source, source_elapsed, source_port, thread
+      FROM system_traces.events
+      WHERE session_id = ?`, traceId)
 
-	for iter.Scan(&timestamp, &activity, &source, &elapsed, &thread) {
-		fmt.Fprintf(t.w, "%s: %s [%s] (source: %s, elapsed: %d)\n",
-			timestamp.Format("2006/01/02 15:04:05.999999"), activity, thread, source, elapsed)
+	for iter.Scan(&timestamp, &activity, &source, &elapsed, &sourcePort, &thread) {
+		fmt.Fprintf(t.w, "%s: %s [%s] (source: %s, elapsed: %d, source_port: %d)\n",
+			timestamp.Format("2006/01/02 15:04:05.999999"), activity, thread, source, elapsed, sourcePort)
 	}
 
 	if err := iter.Close(); err != nil {