From cd5e2a53f4b480ae7078630e76e4bd506316ecc6 Mon Sep 17 00:00:00 2001
From: Mike Gunderloy <MikeG1@larkfarm.com>
Date: Sat, 21 Mar 2009 06:28:27 -0500
Subject: [PATCH] Fix to accomodate non-standard PK columns.

---
 README                   |  7 ++++++-
 db_populate.gemspec      |  6 +++---
 lib/create_or_update.rb  |  7 ++++---
 test/db_populate_test.rb | 21 +++++++++++++++++++++
 test/schema.rb           |  4 ++++
 5 files changed, 38 insertions(+), 7 deletions(-)

diff --git a/README b/README
index b92b3e6..3457f6f 100644
--- a/README
+++ b/README
@@ -54,4 +54,9 @@ db_populate rake tasks
 db_populate includes two rake tasks:
 
 rake db:populate loads all of the data for the current environment
-rake db:migrate_and_populate is the same as calling rake db:migrate followed by rake db:populate
\ No newline at end of file
+rake db:migrate_and_populate is the same as calling rake db:migrate followed by rake db:populate
+
+History
+=======
+2009-03-21 Patch from Ahmed El-Daly to allow PKs with names other than id
+2008-10-11 Initial release
\ No newline at end of file
diff --git a/db_populate.gemspec b/db_populate.gemspec
index 0b842e7..9272cec 100644
--- a/db_populate.gemspec
+++ b/db_populate.gemspec
@@ -1,10 +1,10 @@
 Gem::Specification.new do |s|
   s.name     = "db_populate"
-  s.version  = "0.2.0"
-  s.date     = "2008-10-11"
+  s.version  = "0.2.1"
+  s.date     = "2009-03-21"
   s.summary  = "Seed data populator for Rails"
   s.email    = "MikeG1@larkfarm.com"
-  s.homepage = "http://github.com/ffmike/db_populate"
+  s.homepage = "http://github.com/ffmike/db-populate/tree/master"
   s.description = "db_populate provides rake and code support for adding seed data to Rails projects. Forked from a rake task by
   Josh Knowles, plus code by Luke Francl."
   s.has_rdoc = false
diff --git a/lib/create_or_update.rb b/lib/create_or_update.rb
index b53b6aa..374900a 100644
--- a/lib/create_or_update.rb
+++ b/lib/create_or_update.rb
@@ -1,5 +1,6 @@
 class ActiveRecord::Base
-  # given a hash of attributes including the ID, look up the record by ID. 
+  # given a hash of attributes including the ID, look up the record by ID.
+  # uses whatever the PK of the model is to do the lookup 
   # If it does not exist, it is created with the rest of the options. 
   # If it exists, it is updated with the given options. 
   #
@@ -7,8 +8,8 @@ class ActiveRecord::Base
   # 
   # Returns the record.
   def self.create_or_update(options = {})
-    id = options.delete(:id)
-    record = find_by_id(id) || new
+    id = options.delete(primary_key.to_sym)
+    record = send("find_by_#{primary_key}", id) || new
     record.id = id
     record.attributes = options
     record.save!
diff --git a/test/db_populate_test.rb b/test/db_populate_test.rb
index f543790..36abf7f 100644
--- a/test/db_populate_test.rb
+++ b/test/db_populate_test.rb
@@ -6,6 +6,10 @@ require 'mocha'
 class User < ActiveRecord::Base
 end
 
+class Customer < ActiveRecord::Base
+  set_primary_key "cust_id"
+end
+
 class DbPopulateTest < Test::Unit::TestCase
   
   def test_creates_new_record
@@ -25,5 +29,22 @@ class DbPopulateTest < Test::Unit::TestCase
     assert_equal u.name, "George"
   end
   
+  def test_creates_new_record_with_nonstandard_pk
+    Customer.delete_all
+    Customer.create_or_update(:cust_id => 1, :name => "Fred")
+    assert_equal Customer.count, 1
+    c = Customer.find(:first)
+    assert_equal c.name, "Fred"
+  end
+  
+  def test_updates_existing_record
+    Customer.delete_all
+    Customer.create_or_update(:cust_id => 1, :name => "Fred")
+    Customer.create_or_update(:cust_id => 1, :name => "George")
+    assert_equal Customer.count, 1
+    c = Customer.find(:first)
+    assert_equal c.name, "George"
+  end
+  
 end
 
diff --git a/test/schema.rb b/test/schema.rb
index 9828fbc..b24a850 100644
--- a/test/schema.rb
+++ b/test/schema.rb
@@ -4,4 +4,8 @@ ActiveRecord::Schema.define(:version => 0) do
     t.string   "name"
   end
     
+  create_table "customers", :primary_key => 'cust_id', :force => true do |t|
+    t.string "name"
+  end
+  
 end