add Curator module around riak data store.

lib/curator/riak/data_store.rb

@@ -1,71 +1,73 @@
 require 'riak'
 require 'yaml'
 
-module Riak
-  class DataStore
-    def self.client
-      return @client if @client
-      yml_config = YAML.load(File.read(Curator.riak_config_file))[Curator.environment]
-      @client = Riak::Client.new(yml_config)
-    end
+module Curator
+  module Riak
+    class DataStore
+      def self.client
+        return @client if @client
+        yml_config = YAML.load(File.read(Curator.riak_config_file))[Curator.environment]
+        @client = ::Riak::Client.new(yml_config)
+      end
 
-    def self.delete(bucket_name, key)
-      bucket = _bucket(bucket_name)
-      object = bucket.get(key)
-      object.delete
-    end
+      def self.delete(bucket_name, key)
+        bucket = _bucket(bucket_name)
+        object = bucket.get(key)
+        object.delete
+      end
 
-    def self.ping
-      client.ping
-    end
+      def self.ping
+        client.ping
+      end
 
-    def self.save(options)
-      bucket = _bucket(options[:collection_name])
-      object = Riak::RObject.new(bucket, options[:key])
-      object.content_type = "application/json"
-      object.data = options[:value]
-      options.fetch(:index, {}).each do |index_name, index_value|
-        object.indexes["#{index_name}_bin"] << index_value
+      def self.save(options)
+        bucket = _bucket(options[:collection_name])
+        object = ::Riak::RObject.new(bucket, options[:key])
+        object.content_type = "application/json"
+        object.data = options[:value]
+        options.fetch(:index, {}).each do |index_name, index_value|
+          object.indexes["#{index_name}_bin"] << index_value
+        end
+        object.store
       end
-      object.store
-    end
 
-    def self.find_by_key(bucket_name, key)
-      bucket = _bucket(bucket_name)
-      begin
-        object = bucket.get(key)
-        { :key => object.key, :data => object.data } unless object.data.empty?
-      rescue Riak::HTTPFailedRequest => failed_request
-        raise failed_request unless failed_request.not_found?
+      def self.find_by_key(bucket_name, key)
+        bucket = _bucket(bucket_name)
+        begin
+          object = bucket.get(key)
+          { :key => object.key, :data => object.data } unless object.data.empty?
+        rescue ::Riak::HTTPFailedRequest => failed_request
+          raise failed_request unless failed_request.not_found?
+        end
       end
-    end
 
-    def self.find_by_index(bucket_name, index_name, query)
-      return [] if query.nil?
+      def self.find_by_index(bucket_name, index_name, query)
+        return [] if query.nil?
 
-      bucket = _bucket(bucket_name)
-      begin
-        keys = _find_key_by_index(bucket, index_name.to_s, query)
-        keys.map { |key| find_by_key(bucket_name, key) }
-      rescue Riak::HTTPFailedRequest => failed_request
-        raise failed_request unless failed_request.not_found?
+        bucket = _bucket(bucket_name)
+        begin
+          keys = _find_key_by_index(bucket, index_name.to_s, query)
+          keys.map { |key| find_by_key(bucket_name, key) }
+        rescue ::Riak::HTTPFailedRequest => failed_request
+          raise failed_request unless failed_request.not_found?
+        end
       end
-    end
 
-    def self._bucket(name)
-      client.bucket(_bucket_name(name))
-    end
+      def self._bucket(name)
+        client.bucket(_bucket_name(name))
+      end
 
-    def self._bucket_name(name)
-      bucket_prefix + ":" + name
-    end
+      def self._bucket_name(name)
+        bucket_prefix + ":" + name
+      end
 
-    def self.bucket_prefix
-      "#{Curator.bucket_prefix}:#{Curator.environment}"
-    end
+      def self.bucket_prefix
+        "#{Curator.bucket_prefix}:#{Curator.environment}"
+      end
 
-    def self._find_key_by_index(bucket, index_name, query)
-      bucket.get_index("#{index_name}_bin", query)
+      def self._find_key_by_index(bucket, index_name, query)
+        bucket.get_index("#{index_name}_bin", query)
+      end
     end
   end
 end

lib/curator/riak/test/cleaner.rb

@@ -5,10 +5,10 @@ module Riak
     module Test
       class Cleaner
         def self.remove_all_keys
-          riak = ::Riak::DataStore.client
+          riak = DataStore.client
 
           silence_warnings do
-            buckets = riak.buckets.select { |bucket| bucket.name.start_with?(::Riak::DataStore.bucket_prefix) }
+            buckets = riak.buckets.select { |bucket| bucket.name.start_with?(DataStore.bucket_prefix) }
             buckets.each do |bucket|
               bucket.keys do |keys|
                 keys.each { |key| bucket.delete(key) }

lib/curator/riak/test/resetable_data_store.rb

@@ -1,7 +1,7 @@
 module Curator
   module Riak
     module Test
-      class ResetableDataStore < ::Riak::DataStore
+      class ResetableDataStore < DataStore
         def self.bucket_prefix
           job = "#{ENV['JOB_NAME'].gsub(/[^[:alnum:]]/, '_')}" if ENV['JOB_NAME'].present?
           [Curator.bucket_prefix, job, Curator.environment].compact.join(':')

spec/curator/riak/data_store_spec.rb

@@ -0,0 +1,91 @@
+require 'spec_helper'
+
+module Curator
+  module Riak
+    describe Curator::Riak::DataStore do
+      describe "self.client" do
+        it "returns a riak client with a config read from the yml file provided" do
+          begin
+            File.should_receive(:read).and_return(<<-YML)
+            test:
+              :host: somehost
+              :http_port: 1234
+              :pb_port: 5678
+            YML
+            DataStore.instance_variable_set('@client', nil)
+            client = DataStore.client
+            client.node.host.should == "somehost"
+            client.node.http_port.should == 1234
+            client.node.pb_port.should == 5678
+          ensure
+            DataStore.instance_variable_set("@client", nil)
+          end
+        end
+      end
+
+      describe "self.save" do
+        it "namespaces buckets with app and environment" do
+          DataStore.save(:collection_name => "fake_things", :key => "blah", :value => {"foo" => "bar"})
+          DataStore.client.bucket(DataStore._bucket_name("fake_things")).get("blah").data.should == {"foo" => "bar"}
+        end
+
+        it "can index by multiple things" do
+          begin
+            DataStore.save(
+              :collection_name => "fake_things",
+              :key => "blah",
+              :value => {:foo => "foo-data", :bar => "bar-data"},
+              :index => {:foo => "foo-data", :bar => "bar-data"}
+            )
+
+            foo_result = DataStore.find_by_index("fake_things", "foo", "foo-data").first
+            foo_result[:key].should == "blah"
+            bar_result = DataStore.find_by_index("fake_things", "bar", "bar-data").first
+            bar_result[:key].should == "blah"
+          ensure
+            DataStore._bucket("fake_things").delete("blah")
+          end
+        end
+      end
+
+      describe "self.delete" do
+        it "deletes an object in a buket for a key" do
+          DataStore.save(:collection_name => "heap", :key => "some_key", :value => {"k" => "v"})
+          DataStore.delete("heap", "some_key")
+          DataStore.find_by_key("heap", "some_key").should be_nil
+        end
+      end
+
+      describe "find_by_index" do
+        it "returns an empty array if key is not found" do
+          DataStore.find_by_index("abyss","invalid_index","invalid_key").should be_empty
+        end
+
+        it "returns an empty array if key is nil" do
+          DataStore.find_by_index("abyss","invalid_index", nil).should be_empty
+        end
+
+        it "returns multiple objects" do
+          DataStore.save(:collection_name => "test_collection", :key => "key1", :value => {:indexed_key => "indexed_value"}, :index => {:indexed_key => "indexed_value"})
+          DataStore.save(:collection_name => "test_collection", :key => "key2", :value => {:indexed_key => "indexed_value"}, :index => {:indexed_key => "indexed_value"})
+
+          keys = DataStore.find_by_index("test_collection", :indexed_key, "indexed_value").map { |data| data[:key] }
+          keys.sort.should == ["key1", "key2"]
+        end
+      end
+
+      describe "find_by_key" do
+        it "returns an object by key" do
+          DataStore.save(:collection_name => "heap", :key => "some_key", :value => {"k" => "v"})
+          DataStore.find_by_key("heap", "some_key").should == {:key => "some_key", :data => {"k" => "v"}}
+        end
+      end
+
+      context "bucket name dependent on environment" do
+        it "defaults bucket name" do
+          DataStore::_bucket_name("my_bucket").should == "#{Curator.bucket_prefix}:test:my_bucket"
+        end
+      end
+    end
+  end
+end