Eliminate missed lease_connection calls

lib/active_record/tasks/sqlserver_database_tasks.rb

@@ -10,7 +10,7 @@ module Tasks
     class SQLServerDatabaseTasks
       DEFAULT_COLLATION = "SQL_Latin1_General_CP1_CI_AS"
 
-      delegate :lease_connection, :establish_connection, to: ActiveRecord::Base
+      delegate :with_connection, :establish_connection, to: ActiveRecord::Base
 
       def self.using_database_configurations?
         true
@@ -23,8 +23,10 @@ def initialize(configuration)
 
       def create(master_established = false)
         establish_master_connection unless master_established
-        lease_connection.create_database configuration.database, configuration_hash.merge(collation: default_collation)
-        establish_connection configuration
+        with_connection do |connection|
+          connection.create_database(configuration.database, configuration_hash.merge(collation: default_collation))
+        end
+        establish_connection(configuration)
       rescue ActiveRecord::StatementInvalid => e
         if /database .* already exists/i === e.message
           raise DatabaseAlreadyExists
@@ -35,15 +37,15 @@ def create(master_established = false)
 
       def drop
         establish_master_connection
-        lease_connection.drop_database configuration.database
+        with_connection { |connection| connection.drop_database(configuration.database) }
       end
 
       def charset
-        lease_connection.charset
+        with_connection { |connection| connection.charset }
       end
 
       def collation
-        lease_connection.collation
+        with_connection { |connection| connection.collation }
       end
 
       def purge
@@ -56,34 +58,38 @@ def clear_active_connections!
         ActiveRecord::Base.connection_handler.clear_active_connections!(:all)
       end
 
-      def structure_dump(filename, extra_flags)
-        server_arg = "-S #{Shellwords.escape(configuration_hash[:host])}"
-        server_arg += ":#{Shellwords.escape(configuration_hash[:port])}" if configuration_hash[:port]
-        command = [
-          "defncopy-ttds",
-          server_arg,
-          "-D #{Shellwords.escape(configuration_hash[:database])}",
-          "-U #{Shellwords.escape(configuration_hash[:username])}",
-          "-P #{Shellwords.escape(configuration_hash[:password])}",
-          "-o #{Shellwords.escape(filename)}",
-        ]
-        table_args = lease_connection.tables.map { |t| Shellwords.escape(t) }
-        command.concat(table_args)
-        view_args = lease_connection.views.map { |v| Shellwords.escape(v) }
-        command.concat(view_args)
-        raise "Error dumping database" unless Kernel.system(command.join(" "))
-
-        dump = File.read(filename)
-        dump.gsub!(/^USE .*$\nGO\n/, "")                      # Strip db USE statements
-        dump.gsub!(/^GO\n/, "")                               # Strip db GO statements
-        dump.gsub!(/nvarchar\(8000\)/, "nvarchar(4000)")      # Fix nvarchar(8000) column defs
-        dump.gsub!(/nvarchar\(-1\)/, "nvarchar(max)")         # Fix nvarchar(-1) column defs
-        dump.gsub!(/text\(\d+\)/, "text")                     # Fix text(16) column defs
-        File.open(filename, "w") { |file| file.puts dump }
+      def structure_dump(filename, _extra_flags)
+        with_connection do |connection|
+          server_arg = "-S #{Shellwords.escape(configuration_hash[:host])}"
+          server_arg += ":#{Shellwords.escape(configuration_hash[:port])}" if configuration_hash[:port]
+          command = [
+            "defncopy-ttds",
+            server_arg,
+            "-D #{Shellwords.escape(configuration_hash[:database])}",
+            "-U #{Shellwords.escape(configuration_hash[:username])}",
+            "-P #{Shellwords.escape(configuration_hash[:password])}",
+            "-o #{Shellwords.escape(filename)}",
+          ]
+          table_args = connection.tables.map { |t| Shellwords.escape(t) }
+          command.concat(table_args)
+          view_args = connection.views.map { |v| Shellwords.escape(v) }
+          command.concat(view_args)
+          raise "Error dumping database" unless Kernel.system(command.join(" "))
+
+          dump = File.read(filename)
+          dump.gsub!(/^USE .*$\nGO\n/, "")                      # Strip db USE statements
+          dump.gsub!(/^GO\n/, "")                               # Strip db GO statements
+          dump.gsub!(/nvarchar\(8000\)/, "nvarchar(4000)")      # Fix nvarchar(8000) column defs
+          dump.gsub!(/nvarchar\(-1\)/, "nvarchar(max)")         # Fix nvarchar(-1) column defs
+          dump.gsub!(/text\(\d+\)/, "text")                     # Fix text(16) column defs
+          File.open(filename, "w") { |file| file.puts dump }
+        end
       end
 
-      def structure_load(filename, extra_flags)
-        lease_connection.execute File.read(filename)
+      def structure_load(filename, _extra_flags)
+        with_connection do |connection|
+          connection.execute File.read(filename)
+        end
       end
 
       private

lib/arel/visitors/sqlserver.rb

@@ -129,10 +129,12 @@ def visit_Arel_Table(o, collector)
         # github.com/rails-sqlserver/activerecord-sqlserver-adapter/issues/450
         table_name =
           begin
-            if o.class.engine.lease_connection.respond_to?(:sqlserver?) && o.class.engine.lease_connection.database_prefix_remote_server?
-              remote_server_table_name(o)
-            else
-              quote_table_name(o.name)
+            o.class.engine.with_connection do |connection|
+              if connection.respond_to?(:sqlserver?) && connection.database_prefix_remote_server?
+                remote_server_table_name(o)
+              else
+                quote_table_name(o.name)
+              end
             end
           rescue Exception
             quote_table_name(o.name)
@@ -315,14 +317,16 @@ def primary_Key_From_Table(t)
       end
 
       def remote_server_table_name(o)
-        ActiveRecord::ConnectionAdapters::SQLServer::Utils.extract_identifiers(
-          "#{o.class.engine.lease_connection.database_prefix}#{o.name}"
-        ).quoted
+        o.class.engine.with_connection do |connection|
+          ActiveRecord::ConnectionAdapters::SQLServer::Utils.extract_identifiers(
+            "#{connection.database_prefix}#{o.name}"
+          ).quoted
+        end
       end
 
-      # Need to remove ordering from subqueries unless TOP/OFFSET also used. Otherwise, SQLServer
+      # Need to remove ordering from sub-queries unless TOP/OFFSET also used. Otherwise, SQLServer
       # returns error "The ORDER BY clause is invalid in views, inline functions, derived tables,
-      # subqueries, and common table expressions, unless TOP, OFFSET or FOR XML is also specified."
+      # sub-queries, and common table expressions, unless TOP, OFFSET or FOR XML is also specified."
       def remove_invalid_ordering_from_select_statement(node)
         return unless Arel::Nodes::SelectStatement === node