From 6c65b0176223c90e08903ebc1a9a1075659f2c43 Mon Sep 17 00:00:00 2001
From: Jesse Storimer <jstorimer@gmail.com>
Date: Thu, 20 Jan 2011 00:37:05 -0500
Subject: Allow database specific locking clauses to be used

---
 lib/arel/nodes/lock.rb          |  4 ++++
 lib/arel/select_manager.rb      |  2 +-
 lib/arel/visitors/mysql.rb      |  6 +++++-
 lib/arel/visitors/postgresql.rb |  6 +++++-
 test/visitors/test_mysql.rb     | 15 ++++++++++-----
 test/visitors/test_postgres.rb  | 17 +++++++++++++----
 6 files changed, 38 insertions(+), 12 deletions(-)

diff --git a/lib/arel/nodes/lock.rb b/lib/arel/nodes/lock.rb
index e5fb258e26..f4eaf125e0 100644
--- a/lib/arel/nodes/lock.rb
+++ b/lib/arel/nodes/lock.rb
@@ -1,6 +1,10 @@
 module Arel
   module Nodes
     class Lock < Arel::Nodes::Node
+      attr_reader :locking
+      def initialize locking = true
+        @locking = locking
+      end
     end
   end
 end
diff --git a/lib/arel/select_manager.rb b/lib/arel/select_manager.rb
index b4103144fe..fce6ca17d8 100644
--- a/lib/arel/select_manager.rb
+++ b/lib/arel/select_manager.rb
@@ -40,7 +40,7 @@ module Arel
     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?
-      @ast.lock = Nodes::Lock.new
+      @ast.lock = Nodes::Lock.new(locking)
       self
     end
 
diff --git a/lib/arel/visitors/mysql.rb b/lib/arel/visitors/mysql.rb
index e90161eee4..8f9c3cc809 100644
--- a/lib/arel/visitors/mysql.rb
+++ b/lib/arel/visitors/mysql.rb
@@ -3,7 +3,11 @@ module Arel
     class MySQL < Arel::Visitors::ToSql
       private
       def visit_Arel_Nodes_Lock o
-        "FOR UPDATE"
+        if o.locking.is_a?(String)
+          o.locking
+        else
+          "FOR UPDATE"
+        end
       end
 
       ###
diff --git a/lib/arel/visitors/postgresql.rb b/lib/arel/visitors/postgresql.rb
index 0e82a703ca..45e7349fd3 100644
--- a/lib/arel/visitors/postgresql.rb
+++ b/lib/arel/visitors/postgresql.rb
@@ -3,7 +3,11 @@ module Arel
     class PostgreSQL < Arel::Visitors::ToSql
       private
       def visit_Arel_Nodes_Lock o
-        "FOR UPDATE"
+        if o.locking.is_a?(String)
+          o.locking
+        else
+          "FOR UPDATE"
+        end
       end
 
       def visit_Arel_Nodes_SelectStatement o
diff --git a/test/visitors/test_mysql.rb b/test/visitors/test_mysql.rb
index c3b79ca667..c22cbaff19 100644
--- a/test/visitors/test_mysql.rb
+++ b/test/visitors/test_mysql.rb
@@ -29,11 +29,16 @@ module Arel
         sql.must_be_like "SELECT FROM DUAL"
       end
 
-      it 'uses FOR UPDATE when locking' do
-        stmt = Nodes::SelectStatement.new
-        stmt.lock = Nodes::Lock.new
-        sql = @visitor.accept(stmt)
-        sql.must_be_like "SELECT FROM DUAL FOR UPDATE"
+      describe 'locking' do
+        it 'defaults to FOR UPDATE when locking' do
+          node = Nodes::Lock.new
+          @visitor.accept(node).must_be_like "FOR UPDATE"
+        end
+
+        it 'allows a custom string to be used as a lock' do
+          node = Nodes::Lock.new('LOCK IN SHARE MODE')
+          @visitor.accept(node).must_be_like "LOCK IN SHARE MODE"
+        end
       end
     end
   end
diff --git a/test/visitors/test_postgres.rb b/test/visitors/test_postgres.rb
index 6e8f399b6b..ec55e555a6 100644
--- a/test/visitors/test_postgres.rb
+++ b/test/visitors/test_postgres.rb
@@ -7,10 +7,19 @@ module Arel
         @visitor = PostgreSQL.new Table.engine
       end
 
-      it 'should produce a lock value' do
-        @visitor.accept(Nodes::Lock.new).must_be_like %{
-          FOR UPDATE
-        }
+      describe 'locking' do
+        it 'defaults to FOR UPDATE' do
+          @visitor.accept(Nodes::Lock.new).must_be_like %{
+            FOR UPDATE
+          }
+        end
+
+        it 'allows a custom string to be used as a lock' do
+          node = Nodes::Lock.new('FOR SHARE')
+          @visitor.accept(node).must_be_like %{
+            FOR SHARE
+          }
+        end
       end
 
       it "should escape LIMIT" do
-- 
cgit v1.2.3