Changes to Trilogy adapter for new error classes

lib/semian/trilogy_adapter.rb

@@ -7,9 +7,9 @@
 module ActiveRecord
   module ConnectionAdapters
     class TrilogyAdapter
-      StatementInvalid.include(::Semian::AdapterError)
+      ActiveRecord::ActiveRecordError.include(::Semian::AdapterError)
 
-      class SemianError < StatementInvalid
+      class SemianError < ActiveRecordError
         def initialize(semian_identifier, *args)
           super(*args)
           @semian_identifier = semian_identifier
@@ -30,6 +30,7 @@ module TrilogyAdapter
 
     def initialize(*options)
       *, config = options
+      @read_timeout = config[:read_timeout] || 0
       @raw_semian_options = config.delete(:semian)
       @semian_identifier = begin
         name = semian_options && semian_options[:name]
@@ -54,41 +55,37 @@ def execute(sql, name = nil, async: false, allow_retry: false)
     end
 
     def with_resource_timeout(temp_timeout)
+      connection_was_nil = false
+
       if connection.nil?
-        prev_read_timeout = @config[:read_timeout] || 0 # We're assuming defaulting the timeout to 0 won't change in Trilogy
-        @config[:read_timeout] = temp_timeout # Create new client with config
+        prev_read_timeout = @read_timeout # Use read_timeout from when we first connected
+        connection_was_nil = true
       else
         prev_read_timeout = connection.read_timeout
         connection.read_timeout = temp_timeout
       end
+
       yield
+
+      connection.read_timeout = temp_timeout if connection_was_nil
     ensure
-      @config[:read_timeout] = prev_read_timeout
       connection&.read_timeout = prev_read_timeout
     end
 
     private
 
-    def resource_exceptions
-      []
-    end
-
     def acquire_semian_resource(**)
       super
-    # We're going to need to rescue ConnectionNotEstablished here
-    # and fix this upstream in the Trilogy adapter -- right now, #new_client raises
-    # raw ECONNREFUSED errors
-    # Also, we shouldn't be wrapping TIMEDOUT and ECONNREFUSED in StatementInvalid => use
-    # more appropriate error classes
-    # We see ECONNREFUSED wrapped as an AR::StatementInvalid exception instead of the
-    # raw one when we go through #execute, because it gets translated in #with_raw_connection
-    # I think #new_client should just be more nuanced
     rescue ActiveRecord::StatementInvalid => error
-      if error.cause.is_a?(Errno::ETIMEDOUT) || error.cause.is_a?(Errno::ECONNREFUSED)
+      if error.cause.is_a?(Trilogy::TimeoutError)
         semian_resource.mark_failed(error)
         error.semian_identifier = semian_identifier
       end
       raise
+     end
+
+    def resource_exceptions
+      [ActiveRecord::ConnectionNotEstablished]
     end
 
     # TODO: share this with Mysql2

test/adapters/trilogy_adapter_test.rb

@@ -69,7 +69,7 @@ def test_unconfigured
       def test_connection_errors_open_the_circuit
         @proxy.downstream(:latency, latency: 2200).apply do
           ERROR_THRESHOLD.times do
-            assert_raises(ActiveRecord::StatementInvalid) do
+            assert_raises(ActiveRecord::ConnectionNotEstablished) do
               @adapter.execute("SELECT 1;")
             end
           end
@@ -104,7 +104,7 @@ def test_read_timeout_error_opens_the_circuit
         end
 
         # After TrilogyAdapter::CircuitOpenError check regular queries are working fine.
-        result = Timecop.travel(ERROR_TIMEOUT + 1) do
+        result = time_travel(ERROR_TIMEOUT + 1) do
           @adapter.execute("SELECT 1 + 1;")
         end
 
@@ -151,7 +151,7 @@ def test_query_instrumentation
 
       def test_network_errors_are_tagged_with_the_resource_identifier
         @proxy.down do
-          error = assert_raises(ActiveRecord::StatementInvalid) do
+          error = assert_raises(ActiveRecord::ConnectionNotEstablished) do
             @adapter.execute("SELECT 1 + 1;")
           end
           assert_equal(@adapter.semian_identifier, error.semian_identifier)
@@ -213,7 +213,7 @@ def test_circuit_breaker_on_connect
           trilogy_adapter.connect!
         end
 
-        Timecop.travel(ERROR_TIMEOUT + 1) do
+        time_travel(ERROR_TIMEOUT + 1) do
           trilogy_adapter.connect!
         end
       end
@@ -249,7 +249,7 @@ def test_circuit_breaker_on_query
           @adapter.execute("SELECT 1 + 1;")
         end
 
-        Timecop.travel(ERROR_TIMEOUT + 1) do
+        time_travel(ERROR_TIMEOUT + 1) do
           assert_equal(2, @adapter.execute("SELECT 1 + 1;").to_a.flatten.first)
         end
       end
@@ -309,8 +309,8 @@ def test_changes_timeout_when_half_open_and_configured
         adapter = trilogy_adapter(semian: SEMIAN_OPTIONS.merge(half_open_resource_timeout: 1))
 
         @proxy.downstream(:latency, latency: 3000).apply do
-          (ERROR_THRESHOLD * 2).times do
-            assert_raises(ActiveRecord::StatementInvalid) do
+          ERROR_THRESHOLD.times do
+            assert_raises(ActiveRecord::ConnectionNotEstablished) do
               adapter.execute("SELECT 1 + 1;")
             end
           end
@@ -320,15 +320,16 @@ def test_changes_timeout_when_half_open_and_configured
           adapter.execute("SELECT 1 + 1;")
         end
 
-        Timecop.travel(ERROR_TIMEOUT + 1) do
+        # Circuit moves to half-open state, so 1500 of latency should result in error
+        time_travel(ERROR_TIMEOUT + 1) do
           @proxy.downstream(:latency, latency: 1500).apply do
-            assert_raises(ActiveRecord::StatementInvalid) do
+            assert_raises(ActiveRecord::ConnectionNotEstablished) do
               adapter.execute("SELECT 1 + 1;")
             end
           end
         end
 
-        Timecop.travel(ERROR_TIMEOUT * 2 + 1) do
+        time_travel(ERROR_TIMEOUT * 2 + 1) do
           adapter.execute("SELECT 1 + 1;")
           adapter.execute("SELECT 1 + 1;")
 
@@ -345,16 +346,15 @@ def test_changes_timeout_when_half_open_and_configured
 
       def test_circuit_open_errors_do_not_trigger_the_circuit_breaker
         @proxy.down do
-          err = assert_raises(ActiveRecord::StatementInvalid) do
+          err = assert_raises(ActiveRecord::ConnectionNotEstablished) do
             @adapter.execute("SELECT 1;")
           end
           2.times do
             assert_raises(TrilogyAdapter::CircuitOpenError) do
               @adapter.execute("SELECT 1;")
             end
             error = Semian[:trilogy_adapter_testing].circuit_breaker.last_error
-            assert_equal(ActiveRecord::StatementInvalid, error.class)
-            assert_equal(Errno::ECONNREFUSED, error.cause.class)
+            assert_equal(ActiveRecord::ConnectionNotEstablished, error.class)
           end
         end
       end