refactoring custom lock logic to use sql literals

3 files changed, 11 insertions, 13 deletions

diff --git a/lib/arel/select_manager.rb b/lib/arel/select_manager.rb
index 249ad680a7..adffcca764 100644
--- a/lib/arel/select_manager.rb
+++ b/lib/arel/select_manager.rb
@@ -37,9 +37,15 @@ module Arel
       @ctx.wheres.map { |c| to_sql.accept c }
     end
 
-    def lock locking = true
-      # FIXME: do we even need to store this?  If locking is +false+ shouldn't
-      # we just remove the node from the AST?
+    def lock locking = Arel.sql('FOR UPDATE')
+      case locking
+      when true
+        locking = Arel.sql('FOR UPDATE')
+      when Arel::Nodes::SqlLiteral
+      when String
+        locking = Arel.sql locking
+      end
+
       @ast.lock = Nodes::Lock.new(locking)
       self
     end
diff --git a/lib/arel/visitors/mysql.rb b/lib/arel/visitors/mysql.rb
index dad4d5e3b2..4f029245ac 100644
--- a/lib/arel/visitors/mysql.rb
+++ b/lib/arel/visitors/mysql.rb
@@ -3,11 +3,7 @@ module Arel
     class MySQL < Arel::Visitors::ToSql
       private
       def visit_Arel_Nodes_Lock o
-        if o.expr.is_a?(String)
-          o.expr
-        else
-          "FOR UPDATE"
-        end
+        visit o.expr
       end
 
       ###
diff --git a/lib/arel/visitors/postgresql.rb b/lib/arel/visitors/postgresql.rb
index 68b483910e..c423dc6fc6 100644
--- a/lib/arel/visitors/postgresql.rb
+++ b/lib/arel/visitors/postgresql.rb
@@ -3,11 +3,7 @@ module Arel
     class PostgreSQL < Arel::Visitors::ToSql
       private
       def visit_Arel_Nodes_Lock o
-        if o.expr.is_a?(String)
-          o.expr
-        else
-          "FOR UPDATE"
-        end
+        visit o.expr
       end
 
       def visit_Arel_Nodes_SelectStatement o