Refactor 'Raise ActiveRecord::ConnectionNotEstablished on calls to execute with a disconnected connection'

lib/active_record/connection_adapters/sqlserver/database_statements.rb

@@ -167,9 +167,7 @@ def execute_procedure(proc_name, *variables)
           log(sql, name) do
             case @connection_options[:mode]
             when :dblib
-              raise ActiveRecord::ConnectionNotEstablished if @connection.nil?
-
-              result = @connection.execute(sql)
+              result = ensure_established_connection! { @connection.execute(sql) }
               options = { as: :hash, cache_rows: true, timezone: ActiveRecord::Base.default_timezone || :utc }
               result.each(options) do |row|
                 r = row.with_indifferent_access
@@ -359,9 +357,7 @@ def sp_executesql_sql(sql, types, params, name)
         def raw_connection_do(sql)
           case @connection_options[:mode]
           when :dblib
-            raise ActiveRecord::ConnectionNotEstablished if @connection.nil?
-
-            result = @connection.execute(sql)
+            result = ensure_established_connection! { @connection.execute(sql) }
 
             # TinyTDS returns false instead of raising an exception if connection fails.
             # Getting around this by raising an exception ourselves while this PR
@@ -432,9 +428,7 @@ def _raw_select(sql, options = {})
         def raw_connection_run(sql)
           case @connection_options[:mode]
           when :dblib
-            raise ActiveRecord::ConnectionNotEstablished if @connection.nil?
-
-            @connection.execute(sql)
+            ensure_established_connection! { @connection.execute(sql) }
           end
         end
 
@@ -468,6 +462,12 @@ def finish_statement_handle(handle)
           end
           handle
         end
+
+        def ensure_established_connection!
+          raise ActiveRecord::ConnectionNotEstablished, 'SQL Server client is not connected' unless @connection
+
+          yield
+        end
       end
     end
   end

lib/active_record/connection_adapters/sqlserver_adapter.rb

@@ -374,7 +374,7 @@ def initialize_type_map(m = type_map)
       end
 
       def translate_exception(e, message:, sql:, binds:)
-        return ActiveRecord::ConnectionNotEstablished.new("SQLServer client is not connected") if e.is_a?(ActiveRecord::ConnectionNotEstablished)
+        return e if e.is_a?(ActiveRecord::ConnectionNotEstablished)
 
         case message
         when /(cannot insert duplicate key .* with unique index) | (violation of unique key constraint)/i

test/cases/disconnected_test_sqlserver.rb

@@ -18,7 +18,7 @@ def setup
   test "can't execute procuderes while disconnected" do
     @connection.execute_procedure :sp_tables, "sst_datatypes"
     @connection.disconnect!
-    assert_raises(ActiveRecord::ConnectionNotEstablished) do
+    assert_raises(ActiveRecord::ConnectionNotEstablished, 'SQL Server client is not connected') do
       @connection.execute_procedure :sp_tables, "sst_datatypes"
     end
   end
@@ -32,7 +32,7 @@ def setup
 
     @connection.exec_query sql, "TEST", binds
     @connection.disconnect!
-    assert_raises(ActiveRecord::ConnectionNotEstablished) do
+    assert_raises(ActiveRecord::ConnectionNotEstablished, 'SQL Server client is not connected') do
       @connection.exec_query sql, "TEST", binds
     end
   end