From e769a72d6a6f615e07d171ab0d4f145f56924965 Mon Sep 17 00:00:00 2001
From: Andres Taylor <andres@planetscale.com>
Date: Fri, 22 Nov 2024 10:11:52 +0100
Subject: [PATCH] minor fixes

Signed-off-by: Andres Taylor <andres@planetscale.com>
---
 go/transactions/transaction_signature.go |  4 ++++
 go/transactions/transactions.go          | 24 +++++++++++++++---------
 2 files changed, 19 insertions(+), 9 deletions(-)

diff --git a/go/transactions/transaction_signature.go b/go/transactions/transaction_signature.go
index db29ffb..1aa63c6 100644
--- a/go/transactions/transaction_signature.go
+++ b/go/transactions/transaction_signature.go
@@ -174,5 +174,9 @@ func (m *txSignatureMap) MarshalJSON() ([]byte, error) {
 		signatures = append(signatures, bucket...)
 	}
 
+	sort.Slice(signatures, func(i, j int) bool {
+		return signatures[i].Count > signatures[j].Count
+	})
+
 	return json.Marshal(signatures)
 }
diff --git a/go/transactions/transactions.go b/go/transactions/transactions.go
index d6e3bdb..02a2803 100644
--- a/go/transactions/transactions.go
+++ b/go/transactions/transactions.go
@@ -88,7 +88,14 @@ func (s *state) parse(q string) sqlparser.Statement {
 
 func (s *state) startProducing(loader data.IteratorLoader, defaultAutocommit bool, ch chan<- []sqlparser.Statement) {
 	connections := map[int]*Connection{}
-
+	getConn := func(id int) *Connection {
+		connection, ok := connections[id]
+		if !ok {
+			connection = &Connection{Autocommit: defaultAutocommit}
+			connections[id] = connection
+		}
+		return connection
+	}
 	_ = data.ForeachSQLQuery(loader, func(query data.Query) error {
 		stmt := s.parse(query.Query)
 		if stmt == nil {
@@ -98,21 +105,20 @@ func (s *state) startProducing(loader data.IteratorLoader, defaultAutocommit boo
 		case *sqlparser.Begin:
 		case *sqlparser.Commit:
 			// Commit seen, so we can yield the queries in the transaction
-			connection := connections[query.ConnectionID]
+			connection := getConn(query.ConnectionID)
+			if connection.Transaction == nil {
+				return nil
+			}
 			ch <- connection.Transaction
 			connection.Transaction = nil
 		case *sqlparser.Set:
-			connection := connections[query.ConnectionID]
-			connection.Autocommit = getAutocommitStatus(stmt, connection.Autocommit)
+			conn := getConn(query.ConnectionID)
+			conn.Autocommit = getAutocommitStatus(stmt, defaultAutocommit)
 		default:
 			if !sqlparser.IsDMLStatement(stmt) {
 				return nil
 			}
-			connection, ok := connections[query.ConnectionID]
-			if !ok {
-				connection = &Connection{Autocommit: defaultAutocommit}
-				connections[query.ConnectionID] = connection
-			}
+			connection := getConn(query.ConnectionID)
 			if connection.Autocommit {
 				ch <- []sqlparser.Statement{stmt}
 			} else {