Add/remove support in the command line util

bin/soulmate

@@ -27,7 +27,9 @@ parser = OptionParser.new do |opts|
 
   opts.separator ""
   opts.separator "Commands:"
-  opts.separator "  load TYPE   Loads items of specified type read from stdin in the JSON lines format"
+  opts.separator "  load   TYPE   Replaces collection specified by TYPE with items read from stdin in the JSON lines format"
+  opts.separator "  add    TYPE   Adds items to collection specified by TYPE read from stdin in the JSON lines format"
+  opts.separator "  remove TYPE   Removes items from collection specified by TYPE read from stdin in the JSON lines format. Items only require an 'id', all other fields are ignored."
 end
 
 def load(type)
@@ -37,11 +39,35 @@ def load(type)
   puts "Loaded a total of #{total} items"
 end
 
+def add(type)
+  puts "Adding items of type #{type}..."
+  loader = Soulmate::Loader.new(type)
+  items = $stdin.read.split("\n").map { |l| JSON.parse(l) }
+  items.each do |item|
+    loader.add(item)
+  end
+  puts "Loaded a total of #{items.size} items"
+end
+
+def remove(type)
+  puts "Removing items of type #{type}..."
+  loader = Soulmate::Loader.new(type)
+  items = $stdin.read.split("\n").map { |l| JSON.parse(l) }
+  items.each do |item|
+    loader.remove(item)
+  end
+  puts "Removed a total of #{items.size} items"
+end
+
 parser.parse!
 
 case ARGV[0]
 when 'load'
   load ARGV[1]
+when 'add'
+  add ARGV[1]
+when 'remove'
+  remove ARGV[1]
 else
   puts parser.help
 end

lib/soulmate/loader.rb

@@ -19,14 +19,9 @@ def load(items)
       # delete the data stored for this type
       Soulmate.redis.del(database)
 
-      items_loaded = 0
       items.each_with_index do |item, i|
         add(item, :skip_duplicate_check => true)
-        items_loaded += 1
-        puts "added #{i} entries" if i % 100 == 0 and i != 0
       end
-
-      items_loaded
     end
 
     # "id", "term", "score", "aliases", "data"
@@ -35,7 +30,7 @@ def add(item, opts = {})
       raise ArgumentError unless item["id"] && item["term"]
 
       # kill any old items with this id
-      remove(item["id"]) unless opts[:skip_duplicate_check]
+      remove("id" => item["id"]) unless opts[:skip_duplicate_check]
 
       # store the raw data in a separate key to reduce memory usage
       Soulmate.redis.hset(database, item["id"], JSON.dump(item))
@@ -46,8 +41,9 @@ def add(item, opts = {})
       end
     end
 
-    def remove(id, opts = {})
-      prev_item = Soulmate.redis.hget(database, id)
+    # remove only cares about an item's id, but for consistency takes an object
+    def remove(item)
+      prev_item = Soulmate.redis.hget(database, item["id"])
       if prev_item
         prev_item = JSON.load(prev_item)
         # undo the operations done in add

test/test_soulmate.rb

@@ -10,7 +10,7 @@ def test_integration_can_load_values_and_query
 
     items_loaded = Soulmate::Loader.new('venues').load(items)
 
-    assert_equal 6, items_loaded
+    assert_equal 6, items_loaded.size
 
     matcher = Soulmate::Matcher.new('venues')
     results = matcher.matches_for_term('stad', :limit => 5)
@@ -28,7 +28,7 @@ def test_integration_can_load_values_and_query_via_aliases
 
     items_loaded = Soulmate::Loader.new('venues').load(items)
 
-    assert_equal 6, items_loaded
+    assert_equal 6, items_loaded.size
 
     matcher = Soulmate::Matcher.new('venues')
     results = matcher.matches_for_term('land shark stadium', :limit => 5)
@@ -57,7 +57,7 @@ def test_can_remove_items
     results = matcher.matches_for_term("te", :cache => false)
     assert_equal 1, results.size
 
-    loader.remove(1)
+    loader.remove("id" => 1)
     results = matcher.matches_for_term("te", :cache => false)
     assert_equal 0, results.size