From 9f18dd9925de6421a269530992a509352f921c6f Mon Sep 17 00:00:00 2001
From: Aidan Haran <aharan@fusioneer.com>
Date: Wed, 21 Apr 2021 15:47:39 +0100
Subject: [PATCH] Refactor

---
 .../sqlserver/database_statements.rb           | 18 +++++++++---------
 .../connection_adapters/sqlserver_adapter.rb   |  2 +-
 test/cases/disconnected_test_sqlserver.rb      |  4 ++--
 3 files changed, 12 insertions(+), 12 deletions(-)

diff --git a/lib/active_record/connection_adapters/sqlserver/database_statements.rb b/lib/active_record/connection_adapters/sqlserver/database_statements.rb
index 6ad1b81d4..1ef2dbca0 100644
--- a/lib/active_record/connection_adapters/sqlserver/database_statements.rb
+++ b/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
diff --git a/lib/active_record/connection_adapters/sqlserver_adapter.rb b/lib/active_record/connection_adapters/sqlserver_adapter.rb
index 6cedc2285..75988dced 100644
--- a/lib/active_record/connection_adapters/sqlserver_adapter.rb
+++ b/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
diff --git a/test/cases/disconnected_test_sqlserver.rb b/test/cases/disconnected_test_sqlserver.rb
index 7529ee864..207605d4e 100644
--- a/test/cases/disconnected_test_sqlserver.rb
+++ b/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