From 27e5dfc1b5f606eaf73b98b8aa7dbc4534621960 Mon Sep 17 00:00:00 2001
From: Sebastian Kippe <sebastian@kip.pe>
Date: Mon, 22 Apr 2013 17:11:30 +0200
Subject: [PATCH] Add per-category object counter

---
 lib/remote_storage/riak.rb | 17 +++++++++++++++++
 spec/riak_spec.rb          | 34 ++++++++++++++++++++++++----------
 2 files changed, 41 insertions(+), 10 deletions(-)

diff --git a/lib/remote_storage/riak.rb b/lib/remote_storage/riak.rb
index f812134..715d5e0 100644
--- a/lib/remote_storage/riak.rb
+++ b/lib/remote_storage/riak.rb
@@ -87,6 +87,7 @@ module RemoteStorage
     def put_data(user, directory, key, data, content_type=nil)
       object = build_data_object(user, directory, key, data, content_type)
 
+      object_exists = !object.data.nil?
       existing_object_size = object_size(object)
 
       timestamp = (Time.now.to_f * 1000).to_i
@@ -102,6 +103,7 @@ module RemoteStorage
 
       object.store
 
+      log_object_count(user, directory, 1) unless object_exists
       log_object_size(user, directory, new_object_size, existing_object_size)
       update_all_directory_objects(user, directory, timestamp)
 
@@ -120,6 +122,7 @@ module RemoteStorage
 
       riak_response = data_bucket.delete("#{user}:#{directory}:#{key}")
 
+      log_object_count(user, directory, -1)
       log_object_size(user, directory, 0, existing_object_size)
 
       timestamp = (Time.now.to_f * 1000).to_i
@@ -166,6 +169,20 @@ module RemoteStorage
       info.store
     end
 
+    def log_object_count(user, directory, change)
+      category = extract_category(directory)
+
+      info = info_bucket.get_or_new("usage:count:#{user}:#{category}")
+      info.content_type = "text/plain"
+
+      count = change.to_i
+      count += info.data.to_i
+
+      info.data = count.to_s
+      info.indexes.merge!({:user_id_bin => [user]})
+      info.store
+    end
+
     def object_size(object)
       if binary_link = object.links.select {|l| l.tag == "binary"}.first
         response = head(LiquorCabinet.config['buckets']['binaries'], escape(binary_link.key))
diff --git a/spec/riak_spec.rb b/spec/riak_spec.rb
index f43dcf4..c985fa8 100644
--- a/spec/riak_spec.rb
+++ b/spec/riak_spec.rb
@@ -1,9 +1,9 @@
 require_relative "spec_helper"
 
-def set_usage_size_info(user, category, size)
-  object = info_bucket.get_or_new("usage:size:#{user}:#{category}")
+def set_usage_info(user, category, type, value)
+  object = info_bucket.get_or_new("usage:#{type}:#{user}:#{category}")
   object.content_type = "text/plain"
-  object.data = size.to_s
+  object.data = value.to_s
   object.store
 end
 
@@ -93,7 +93,7 @@ describe "App with Riak backend" do
     describe "PUT" do
       before do
         header "Authorization", "Bearer 123"
-        set_usage_size_info "jimmy", "documents", "23"
+        set_usage_info "jimmy", "documents", "size", "23"
       end
 
       describe "with implicit content type" do
@@ -206,7 +206,7 @@ describe "App with Riak backend" do
 
       describe "with existing content" do
         before do
-          set_usage_size_info "jimmy", "documents", "10"
+          set_usage_info "jimmy", "documents", "size", "10"
           put "/jimmy/documents/archive/foo", "lorem ipsum"
           put "/jimmy/documents/archive/foo", "some awesome content"
         end
@@ -223,18 +223,27 @@ describe "App with Riak backend" do
 
       describe "public data" do
         before do
-          set_usage_size_info "jimmy", "public/documents", "10"
+          set_usage_info "jimmy", "public/documents", "size", "10"
+          set_usage_info "jimmy", "public/documents", "count", "100"
           put "/jimmy/public/documents/notes/foo", "note to self"
         end
 
+        # after do
+        #   info_bucket.delete "usage:size:jimmy:public/documents"
+        # end
+
         it "saves the value" do
           last_response.status.must_equal 200
           data_bucket.get("jimmy:public/documents/notes:foo").data.must_equal "note to self"
         end
 
-        it "increases the overall category size" do
+        it "increases the category size counter" do
           info_bucket.get("usage:size:jimmy:public/documents").data.must_equal "22"
         end
+
+        it "increases the category object counter" do
+          info_bucket.get("usage:count:jimmy:public/documents").data.must_equal "101"
+        end
       end
 
       context "with binary data" do
@@ -348,7 +357,8 @@ describe "App with Riak backend" do
     describe "DELETE" do
       before do
         header "Authorization", "Bearer 123"
-        set_usage_size_info "jimmy", "documents", "123"
+        set_usage_info "jimmy", "documents", "size", "123"
+        set_usage_info "jimmy", "documents", "count", "1000"
         delete "/jimmy/documents/foo"
       end
 
@@ -359,17 +369,21 @@ describe "App with Riak backend" do
         }.must_raise Riak::HTTPFailedRequest
       end
 
-      it "decreases the overall category size" do
+      it "decreases the category size counter" do
         info_bucket.get("usage:size:jimmy:documents").data.must_equal "101"
       end
 
+      it "decreases the category object counter" do
+        info_bucket.get("usage:count:jimmy:documents").data.must_equal "999"
+      end
+
       context "binary data" do
         before do
           header "Content-Type", "image/jpeg; charset=binary"
           filename = File.join(File.expand_path(File.dirname(__FILE__)), "fixtures", "rockrule.jpeg")
           @image = File.open(filename, "r").read
           put "/jimmy/documents/jaypeg", @image
-          set_usage_size_info "jimmy", "documents", "100000"
+          set_usage_info "jimmy", "documents", "size", "100000"
 
           delete "/jimmy/documents/jaypeg"
         end