Get MSSQL working with ARel 2, refactor limit helper code into shared…

… module

lib/arel/visitors/mssql.rb

@@ -1,13 +1,44 @@
 module Arel
   module Visitors
     class SQLServer < Arel::Visitors::ToSql
+      include ArJdbc::MsSQL::LimitHelpers::SqlServerReplaceLimitOffset
+
+      def select_count? o
+        sel = o.cores.length == 1 && o.cores.first
+        projections = sel.projections.length == 1 && sel.projections
+        Arel::Nodes::Count === projections.first
+      end
+
+      # Need to mimic the subquery logic in ARel 1.x for select count with limit
+      # See arel/engines/sql/compilers/mssql_compiler.rb for details
       def visit_Arel_Nodes_SelectStatement o
         order = "ORDER BY #{o.orders.map { |x| visit x }.join(', ')}" unless o.orders.empty?
-        add_limit_offset([o.cores.map { |x| visit_Arel_Nodes_SelectCore x }.join, order].compact.join(' '), o, order)
+        if o.limit
+          if select_count?(o)
+            subquery = true
+            sql = o.cores.map do |x|
+              x = x.dup
+              x.projections = [Arel::Nodes::SqlLiteral.new("*")]
+              visit_Arel_Nodes_SelectCore x
+            end.join
+          else
+            sql = o.cores.map { |x| visit_Arel_Nodes_SelectCore x }.join
+          end
+
+          order ||= "ORDER BY #{@connection.determine_order_clause(sql)}"
+          replace_limit_offset!(sql, o.limit.to_i, o.offset && o.offset.value.to_i, order)
+          sql = "SELECT COUNT(*) AS count_id FROM (#{sql}) AS subquery" if subquery
+        elsif order
+          sql << " #{order}"
+        else
+          sql = super
+        end
+        sql
       end
     end
 
     class SQLServer2000 < SQLServer
+      include ArJdbc::MsSQL::LimitHelpers::SqlServer2000ReplaceLimitOffset
     end
   end
 end

lib/arjdbc/mssql/adapter.rb

@@ -46,9 +46,9 @@ def sqlserver_version
 
     def add_version_specific_add_limit_offset
       if sqlserver_version == "2000"
-        extend LimitHelpers::SqlServer2000LimitOffset
+        extend LimitHelpers::SqlServer2000AddLimitOffset
       else
-        extend LimitHelpers::SqlServerLimitOffset
+        extend LimitHelpers::SqlServerAddLimitOffset
       end
     end
 
@@ -392,7 +392,6 @@ def add_lock!(sql, options)
       sql
     end
 
-    private
     # Turns IDENTITY_INSERT ON for table during execution of the block
     # N.B. This sets the state of IDENTITY_INSERT to OFF after the
     # block has been executed without regard to its previous state

lib/arjdbc/mssql/limit_helpers.rb

@@ -1,6 +1,7 @@
 module ::ArJdbc
   module MsSQL
     module LimitHelpers
+      module_function
       def get_table_name(sql)
         if sql =~ /^\s*insert\s+into\s+([^\(\s,]+)\s*|^\s*update\s+([^\(\s,]+)\s*/i
           $1
@@ -11,15 +12,13 @@ def get_table_name(sql)
         end
       end
 
-      module SqlServer2000LimitOffset
-        def add_limit_offset!(sql, options)
-          limit = options[:limit]
+      module SqlServer2000ReplaceLimitOffset
+        module_function
+        def replace_limit_offset!(sql, limit, offset, order)
           if limit
-            offset = (options[:offset] || 0).to_i
+            offset ||= 0
             start_row = offset + 1
             end_row = offset + limit.to_i
-            order = (options[:order] || determine_order_clause(sql))
-            sql.sub!(/ ORDER BY.*$/i, '')
             find_select = /\b(SELECT(?:\s+DISTINCT)?)\b(.*)/im
             whole, select, rest_of_query = find_select.match(sql).to_a
             if (start_row == 1) && (end_row ==1)
@@ -31,41 +30,61 @@ def add_limit_offset!(sql, options)
               #removing out stuff before the FROM...
               rest = rest_of_query[/FROM/i=~ rest_of_query.. -1]
               #need the table name for avoiding amiguity
-              table_name = get_table_name(sql)
+              table_name = LimitHelpers.get_table_name(sql)
               #I am not sure this will cover all bases.  but all the tests pass
-              new_order = "#{order}, #{table_name}.id" if order.index("#{table_name}.id").nil?
+              new_order = "ORDER BY #{order}, #{table_name}.id" if order.index("#{table_name}.id").nil?
               new_order ||= order
 
               if (rest_of_query.match(/WHERE/).nil?)
-                new_sql = "#{select} TOP #{limit} #{rest_of_query} WHERE #{table_name}.id NOT IN (#{select} TOP #{offset} #{table_name}.id #{rest} ORDER BY #{new_order}) ORDER BY #{order} "
+                new_sql = "#{select} TOP #{limit} #{rest_of_query} WHERE #{table_name}.id NOT IN (#{select} TOP #{offset} #{table_name}.id #{rest} #{new_order}) #{order} "
               else
-                new_sql = "#{select} TOP #{limit} #{rest_of_query} AND #{table_name}.id NOT IN (#{select} TOP #{offset} #{table_name}.id #{rest} ORDER BY #{new_order}) ORDER BY #{order} "
+                new_sql = "#{select} TOP #{limit} #{rest_of_query} AND #{table_name}.id NOT IN (#{select} TOP #{offset} #{table_name}.id #{rest} #{new_order}) #{order} "
               end
 
               sql.replace(new_sql)
             end
           end
+          sql
         end
       end
 
-      module SqlServerLimitOffset
+      module SqlServer2000AddLimitOffset
         def add_limit_offset!(sql, options)
-          limit = options[:limit]
+          if options[:limit]
+            order = "ORDER BY #{options[:order] || determine_order_clause(sql)}"
+            sql.sub!(/ ORDER BY.*$/i, '')
+            SqlServerReplaceLimitOffset.replace_limit_offset!(sql, options[:limit], options[:offset], order)
+          end
+        end
+      end
+
+      module SqlServerReplaceLimitOffset
+        module_function
+        def replace_limit_offset!(sql, limit, offset, order)
           if limit
-            offset = (options[:offset] || 0).to_i
+            offset ||= 0
             start_row = offset + 1
             end_row = offset + limit.to_i
-            order = (options[:order] || determine_order_clause(sql))
-            sql.sub!(/ ORDER BY.*$/i, '')
             find_select = /\b(SELECT(?:\s+DISTINCT)?)\b(.*)/im
             whole, select, rest_of_query = find_select.match(sql).to_a
             if rest_of_query.strip!.first == '*'
               from_table = /.*FROM\s*\b(\w*)\b/i.match(rest_of_query).to_a[1]
             end
-            new_sql = "#{select} t.* FROM (SELECT ROW_NUMBER() OVER(ORDER BY #{order}) AS _row_num, #{from_table + '.' if from_table}#{rest_of_query}"
+            new_sql = "#{select} t.* FROM (SELECT ROW_NUMBER() OVER(#{order}) AS _row_num, #{from_table + '.' if from_table}#{rest_of_query}"
             new_sql << ") AS t WHERE t._row_num BETWEEN #{start_row.to_s} AND #{end_row.to_s}"
             sql.replace(new_sql)
           end
+          sql
+        end
+      end
+
+      module SqlServerAddLimitOffset
+        def add_limit_offset!(sql, options)
+          if options[:limit]
+            order = "ORDER BY #{options[:order] || determine_order_clause(sql)}"
+            sql.sub!(/ ORDER BY.*$/i, '')
+            SqlServerReplaceLimitOffset.replace_limit_offset!(sql, options[:limit], options[:offset], order)
+          end
         end
       end
     end