From f2ad5bd1fd86d2e61ee2523ced386aa777b8f5e5 Mon Sep 17 00:00:00 2001
From: Adrianna Chang <adrianna.chang@shopify.com>
Date: Tue, 6 Dec 2022 15:51:14 -0500
Subject: [PATCH] Fix ProtocolError => ConnectionError and add test

---
 contrib/ruby/lib/trilogy.rb      | 55 +++++++++++++++++++-------------
 contrib/ruby/test/client_test.rb |  8 +++++
 2 files changed, 40 insertions(+), 23 deletions(-)

diff --git a/contrib/ruby/lib/trilogy.rb b/contrib/ruby/lib/trilogy.rb
index cb0b38f0..c0b898c9 100644
--- a/contrib/ruby/lib/trilogy.rb
+++ b/contrib/ruby/lib/trilogy.rb
@@ -14,6 +14,18 @@ module ConnectionError
 
   class BaseError < StandardError
     include Error
+
+    attr_reader :error_code
+
+    def initialize(error_message = nil, error_code = nil)
+      message = error_code ? "#{error_code}: #{error_message}" : error_message
+      super(message)
+      @error_code = error_code
+    end
+  end
+
+  class BaseConnectionError < BaseError
+    include ConnectionError
   end
 
   # Trilogy::ClientError is the base error type for invalid queries or parameters
@@ -28,42 +40,39 @@ class QueryError < ClientError
   class CastError < ClientError
   end
 
-
   class TimeoutError < Errno::ETIMEDOUT
     include ConnectionError
+
+    attr_reader :error_code
+
+    def initialize(error_message = nil, error_code = nil)
+      super
+      @error_code = error_code
+    end
   end
 
   class ProtocolError < BaseError
     ERROR_CODES = {
       1205 => TimeoutError, # ER_LOCK_WAIT_TIMEOUT
-      1044 => ConnectionError, # ER_DBACCESS_DENIED_ERROR
-      1045 => ConnectionError, # ER_ACCESS_DENIED_ERROR
-      1152 => ConnectionError, # ER_ABORTING_CONNECTION
-      1153 => ConnectionError, # ER_NET_PACKET_TOO_LARGE
-      1154 => ConnectionError, # ER_NET_READ_ERROR_FROM_PIPE
-      1155 => ConnectionError, # ER_NET_FCNTL_ERROR
-      1156 => ConnectionError, # ER_NET_PACKETS_OUT_OF_ORDER
-      1157 => ConnectionError, # ER_NET_UNCOMPRESS_ERROR
-      1158 => ConnectionError, # ER_NET_READ_ERROR
-      1159 => ConnectionError, # ER_NET_READ_INTERRUPTED
-      1160 => ConnectionError, # ER_NET_ERROR_ON_WRITE
-      1161 => ConnectionError, # ER_NET_WRITE_INTERRUPTED
-      1927 => ConnectionError, # ER_CONNECTION_KILLED
+      1044 => BaseConnectionError, # ER_DBACCESS_DENIED_ERROR
+      1045 => BaseConnectionError, # ER_ACCESS_DENIED_ERROR
+      1152 => BaseConnectionError, # ER_ABORTING_CONNECTION
+      1153 => BaseConnectionError, # ER_NET_PACKET_TOO_LARGE
+      1154 => BaseConnectionError, # ER_NET_READ_ERROR_FROM_PIPE
+      1155 => BaseConnectionError, # ER_NET_FCNTL_ERROR
+      1156 => BaseConnectionError, # ER_NET_PACKETS_OUT_OF_ORDER
+      1157 => BaseConnectionError, # ER_NET_UNCOMPRESS_ERROR
+      1158 => BaseConnectionError, # ER_NET_READ_ERROR
+      1159 => BaseConnectionError, # ER_NET_READ_INTERRUPTED
+      1160 => BaseConnectionError, # ER_NET_ERROR_ON_WRITE
+      1161 => BaseConnectionError, # ER_NET_WRITE_INTERRUPTED
+      1927 => BaseConnectionError, # ER_CONNECTION_KILLED
     }
-
-    attr_reader :error_code, :error_message
-
     class << self
       def from_code(message, code)
         ERROR_CODES.fetch(code, self).new(message, code)
       end
     end
-
-    def initialize(error_message, error_code)
-      super("#{error_code}: #{error_message}")
-      @error_code = error_code
-      @error_message = error_message
-    end
   end
 
   class SSLError < BaseError
diff --git a/contrib/ruby/test/client_test.rb b/contrib/ruby/test/client_test.rb
index 6d283e84..6eca9a6f 100644
--- a/contrib/ruby/test/client_test.rb
+++ b/contrib/ruby/test/client_test.rb
@@ -528,6 +528,14 @@ def test_timeout_error
     end
   end
 
+  def test_connection_error
+    err = assert_raises Trilogy::BaseConnectionError do
+      new_tcp_client(username: "foo")
+    end
+
+    assert_includes err.message, "Access denied for user 'foo'"
+  end
+
   def test_database_error
     client = new_tcp_client