lib/flipper/adapters/mongo.rb in flipper-mongo-0.1.1 vs lib/flipper/adapters/mongo.rb in flipper-mongo-0.2.0

@@ -1,65 +1,56 @@
 require 'set'
+require 'forwardable'
 require 'mongo'
 
 module Flipper
   module Adapters
     class Mongo
-      DefaultId = 'flipper'
+      extend Forwardable
 
-      def initialize(collection, options = {})
+      def initialize(collection)
         @collection = collection
-        @options = options
-        @mongo_criteria = {:_id => id}
-        @mongo_options = {:upsert => true, :safe => true}
+        @update_options = {:safe => true, :upsert => true}
       end
 
       def read(key)
-        read_key(key)
+        find_one key
       end
 
       def write(key, value)
-        update '$set' => {key => value}
+        update key, {'$set' => {'v' => value}}
       end
 
       def delete(key)
-        update '$unset' => {key => 1}
+        @collection.remove criteria(key)
       end
 
+      def set_members(key)
+        (find_one(key) || Set.new).to_set
+      end
+
       def set_add(key, value)
-        update '$addToSet' => {key => value}
+        update key, {'$addToSet' => {'v' => value}}
       end
 
       def set_delete(key, value)
-        update '$pull' => {key => value}
+        update key, {'$pull' => {'v' => value}}
       end
 
-      def set_members(key)
-        read_key(key) { Set.new }.to_set
-      end
-
       private
 
-      def id
-        @id ||= @options.fetch(:id) { DefaultId }
+      def find_one(key)
+        if (doc = @collection.find_one(criteria(key)))
+          doc['v']
+        end
       end
 
-      def update(updates)
-        @collection.update(@mongo_criteria, updates, @mongo_options)
+      def update(key, updates)
+        @collection.update criteria(key), updates, @update_options
       end
 
-      def read_key(key, &block)
-        load
-
-        if block_given?
-          @document.fetch(key, &block)
-        else
-          @document[key]
-        end
-      end
-
-      def load
-        @document = @collection.find_one(@mongo_criteria) || {}
+      def criteria(key)
+        {:_id => key}
       end
     end
   end
 end