4 files changed, 30 insertions, 11 deletions

diff --git a/app/models/node.rb b/app/models/node.rb
index da26164..0f34b73 100644
--- a/app/models/node.rb
+++ b/app/models/node.rb
@@ -8,6 +8,7 @@ class Node < ActiveRecord::Base
   belongs_to  :draft, :class_name => "Page",  :foreign_key => :draft_id
   has_many    :permissions
   has_one     :event
+  belongs_to  :user,                          :foreign_key => :locking_user_id
   
   # Callbacks
   after_create :initialize_empty_page
@@ -43,17 +44,17 @@ class Node < ActiveRecord::Base
   
   # Instance Methods
   
-  def find_or_create_draft user
-    if draft && draft.user == user
+  def find_or_create_draft current_user
+    if draft && self.user == current_user
       draft
-    elsif draft && draft.user.nil?
-      draft.user = user
-      draft.save
+    elsif draft && self.user.nil?
+      lock_for! current_user
       draft
-    elsif draft && draft.user != user
+    elsif draft && self.user != current_user
       raise "Page is locked"
     else
-      create_new_draft user
+      lock_for! current_user
+      create_new_draft current_user
     end
   end
   
@@ -73,6 +74,7 @@ class Node < ActiveRecord::Base
       self.head.save!
       self.draft = nil
       self.save!
+      self.unlock!
     else
       nil
     end
@@ -95,6 +97,15 @@ class Node < ActiveRecord::Base
   end
   
   protected
+    def lock_for! current_user
+      self.user = current_user
+      self.save
+    end
+  
+    def unlock!
+      self.user = nil
+      self.save
+    end
   
     # Creates an empty page and associates it to the given node. This means
     # freshly created node has an empty draft. A user can create nodes as he
diff --git a/app/views/nodes/index.html.erb b/app/views/nodes/index.html.erb
index 103617c..53ddb5b 100644
--- a/app/views/nodes/index.html.erb
+++ b/app/views/nodes/index.html.erb
@@ -28,7 +28,7 @@
           <%= link_to 'Unlock', unlock_node_path(node), :method => :put, :confirm => "Are you sure you want to unlock?" %>
       </td>
       <td>
-        <%= "#{node.draft.user.login}" if node.draft && node.draft.user %>
+        <%= node.user.login if node.user %>
       </td>
       <td>
         <%= node.draft ? node.draft.revision : node.head.revision %>
diff --git a/db/migrate/20090321150359_add_locking_user_id_to_nodes.rb b/db/migrate/20090321150359_add_locking_user_id_to_nodes.rb
new file mode 100644
index 0000000..b40dd71
--- /dev/null
+++ b/db/migrate/20090321150359_add_locking_user_id_to_nodes.rb
@@ -0,0 +1,9 @@
+class AddLockingUserIdToNodes < ActiveRecord::Migration
+  def self.up
+    add_column :nodes, :locking_user_id, :integer
+  end
+
+  def self.down
+    remove_column :nodes, :locking_user_id
+  end
+end
diff --git a/test/unit/node_test.rb b/test/unit/node_test.rb
index 2558ca3..a7405f0 100644
--- a/test/unit/node_test.rb
+++ b/test/unit/node_test.rb
@@ -30,7 +30,7 @@ class NodeTest < ActiveSupport::TestCase
     assert_not_nil @first_child.draft
     assert_nil @first_child.draft.user
     @first_child.find_or_create_draft @user1
-    assert_equal @user1, @first_child.draft.user
+    assert_equal @user1, @first_child.user
   end
   
   def test_unique_path_returns_an_array
@@ -121,9 +121,8 @@ class NodeTest < ActiveSupport::TestCase
     node = Node.create :slug => "third_child"
     node.move_to_child_of @root
     node.publish_draft!
-
     node.find_or_create_draft @user1
-
+    assert_equal @user1, node.user
     assert_raise(RuntimeError) do
       node.find_or_create_draft @user2
     end