Handle offset on update.

Related issue on Rails: rails/rails/issues/10849

lib/arel/nodes/update_statement.rb

@@ -2,14 +2,15 @@ module Arel
   module Nodes
     class UpdateStatement < Arel::Nodes::Node
       attr_accessor :relation, :wheres, :values, :orders, :limit
-      attr_accessor :key
+      attr_accessor :key, :offset
 
       def initialize
         @relation = nil
         @wheres   = []
         @values   = []
         @orders   = []
         @limit    = nil
+        @offset   = nil
         @key      = nil
       end
 
@@ -20,7 +21,7 @@ def initialize_copy other
       end
 
       def hash
-        [@relation, @wheres, @values, @orders, @limit, @key].hash
+        [@relation, @wheres, @values, @orders, @limit, @key, @offset].hash
       end
 
       def eql? other
@@ -30,6 +31,7 @@ def eql? other
           self.values == other.values &&
           self.orders == other.orders &&
           self.limit == other.limit &&
+          self.offset == other.offset &&
           self.key == other.key
       end
       alias :== :eql?

lib/arel/update_manager.rb

@@ -24,6 +24,20 @@ def order *expr
       self
     end
 
+    def offset
+      @ast.offset && @ast.offset.expr
+    end
+
+    def skip amount
+      if amount
+        @ast.offset = Nodes::Offset.new(amount)
+      else
+        @ast.offset = nil
+      end
+      self
+    end
+    alias :offset= :skip
+
     ###
     # UPDATE +table+
     def table table

lib/arel/visitors/to_sql.rb

@@ -91,11 +91,12 @@ def build_subselect key, o
         core.projections = [key]
         stmt.limit       = o.limit
         stmt.orders      = o.orders
+        stmt.offset      = o.offset
         stmt
       end
 
       def visit_Arel_Nodes_UpdateStatement o, collector
-        if o.orders.empty? && o.limit.nil?
+        if o.orders.empty? && o.limit.nil? && o.offset.nil?
           wheres = o.wheres
         else
           wheres = [Nodes::In.new(o.key, [build_subselect(o.key, o)])]

test/nodes/test_update_statement.rb

@@ -24,13 +24,15 @@
       statement1.values   = false
       statement1.orders   = %w[x y z]
       statement1.limit    = 42
+      statement1.offset   = 10
       statement1.key      = 'zomg'
       statement2 = Arel::Nodes::UpdateStatement.new
       statement2.relation = 'zomg'
       statement2.wheres   = 2
       statement2.values   = false
       statement2.orders   = %w[x y z]
       statement2.limit    = 42
+      statement2.offset   = 10
       statement2.key      = 'zomg'
       array = [statement1, statement2]
       assert_equal 1, array.uniq.size
@@ -43,13 +45,15 @@
       statement1.values   = false
       statement1.orders   = %w[x y z]
       statement1.limit    = 42
+      statement1.offset   = 11
       statement1.key      = 'zomg'
       statement2 = Arel::Nodes::UpdateStatement.new
       statement2.relation = 'zomg'
       statement2.wheres   = 2
       statement2.values   = false
       statement2.orders   = %w[x y z]
       statement2.limit    = 42
+      statement2.offset   = 10
       statement2.key      = 'wth'
       array = [statement1, statement2]
       assert_equal 2, array.uniq.size

test/test_update_manager.rb

@@ -26,6 +26,16 @@ module Arel
       assert_match(/LIMIT 10/, um.to_sql)
     end
 
+    it 'handles offset properly' do
+      table = Table.new(:users)
+      um = Arel::UpdateManager.new Table.engine
+      um.key = 'id'
+      um.skip 10
+      um.table table
+      um.set [[table[:name], nil]]
+      assert_match(/OFFSET 10/, um.to_sql)
+    end
+
     describe 'set' do
       it "updates with null" do
         table = Table.new(:users)