Log object sizes to a per-user-and-category counter (closes #22)

lib/remote_storage/riak.rb

@@ -27,6 +27,10 @@ module RemoteStorage
       @binary_bucket ||= client.bucket(LiquorCabinet.config['buckets']['binaries'])
     end
 
+    def info_bucket
+      @info_bucket ||= client.bucket(LiquorCabinet.config['buckets']['info'])
+    end
+
     def authorize_request(user, directory, token, listing=false)
       request_method = env["REQUEST_METHOD"]
 
@@ -81,24 +85,24 @@ module RemoteStorage
     end
 
     def put_data(user, directory, key, data, content_type=nil)
-      object = data_bucket.new("#{user}:#{directory}:#{key}")
+      object = build_data_object(user, directory, key, data, content_type)
-      object.content_type = content_type || "text/plain; charset=utf-8"
 
-      directory_index = directory == "" ? "/" : directory
+      existing_object_size = object_size(object)
-      object.indexes.merge!({:user_id_bin => [user],
-                             :directory_bin => [CGI.escape(directory_index)]})
 
       timestamp = (Time.now.to_f * 1000).to_i
       object.meta["timestamp"] = timestamp
 
       if binary_data?(object.content_type, data)
         save_binary_data(object, data) or halt 422
+        new_object_size = data.size
       else
         set_object_data(object, data) or halt 422
+        new_object_size = object.raw_data.size
       end
 
       object.store
 
+      log_object_size(user, directory, new_object_size, existing_object_size)
       update_all_directory_objects(user, directory, timestamp)
 
       halt 200
@@ -108,6 +112,7 @@ module RemoteStorage
 
     def delete_data(user, directory, key)
       object = data_bucket.get("#{user}:#{directory}:#{key}")
+      existing_object_size = object_size(object)
 
       if binary_link = object.links.select {|l| l.tag == "binary"}.first
         client[binary_link.bucket].delete(binary_link.key)
@@ -115,6 +120,8 @@ module RemoteStorage
 
       riak_response = data_bucket.delete("#{user}:#{directory}:#{key}")
 
+      log_object_size(user, directory, 0, existing_object_size)
+
       timestamp = (Time.now.to_f * 1000).to_i
       delete_or_update_directory_objects(user, directory, timestamp)
 
@@ -123,8 +130,68 @@ module RemoteStorage
       halt 404
     end
 
+
     private
 
+    def extract_category(directory)
+      if directory.match(/^public\//)
+        "public/#{directory.split('/')[1]}"
+      else
+        directory.split('/').first
+      end
+    end
+
+    def build_data_object(user, directory, key, data, content_type=nil)
+      object = data_bucket.get_or_new("#{user}:#{directory}:#{key}")
+
+      object.content_type = content_type || "text/plain; charset=utf-8"
+
+      directory_index = directory == "" ? "/" : directory
+      object.indexes.merge!({:user_id_bin => [user],
+                             :directory_bin => [CGI.escape(directory_index)]})
+
+      object
+    end
+
+    def log_object_size(user, directory, new_size=0, old_size=0)
+      category = extract_category(directory)
+      info = info_bucket.get_or_new("usage:size:#{user}:#{category}")
+
+      info.content_type = "text/plain"
+      size = -old_size + new_size
+      size += info.data.to_i
+
+      info.data = size.to_s
+      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))
+        response[:headers]["content-length"].first.to_i
+      else
+        object.raw_data.nil? ? 0 : object.raw_data.size
+      end
+    end
+
+    def escape(string)
+      ::Riak.escaper.escape(string).gsub("+", "%20").gsub('/', "%2F")
+    end
+
+    # Perform a HEAD request via the backend method
+    def head(bucket, key)
+      client.http do |h|
+        url = riak_uri(bucket, key)
+        h.head [200], url
+      end
+    end
+
+    # A URI object that can be used with HTTP backend methods
+    def riak_uri(bucket, key)
+      rc = LiquorCabinet.config['riak'].symbolize_keys
+      URI.parse "http://#{rc[:host]}:#{rc[:http_port]}/riak/#{bucket}/#{key}"
+    end
+
     def serializer_for(content_type)
       ::Riak::Serializers[content_type[/^[^;\s]+/]]
     end
@@ -146,7 +213,7 @@ module RemoteStorage
       permission = authorizations[""]
 
       authorizations.each do |key, value|
-        if directory.match /^(public\/)?#{key}(\/|$)/
+        if directory.match(/^(public\/)?#{key}(\/|$)/)
           if permission.nil? || permission == "r"
             permission = value
           end
@@ -253,7 +320,7 @@ module RemoteStorage
     end
 
     def update_directory_object(user, directory, timestamp)
-      if directory.match /\//
+      if directory.match(/\//)
         parent_directory = directory[0..directory.rindex("/")-1]
       elsif directory != ""
         parent_directory = "/"

spec/riak_spec.rb

@@ -1,5 +1,12 @@
 require_relative "spec_helper"
 
+def set_usage_size_info(user, category, size)
+  object = info_bucket.get_or_new("usage:size:#{user}:#{category}")
+  object.content_type = "text/plain"
+  object.data = size.to_s
+  object.store
+end
+
 describe "App with Riak backend" do
   include Rack::Test::Methods
   include RemoteStorage::Riak
@@ -23,6 +30,7 @@ describe "App with Riak backend" do
       last_response.body.must_equal "some text data"
     end
 
+    # If this one fails, try restarting Riak
     it "has a Last-Modified header set" do
       last_response.status.must_equal 200
       last_response.headers["Last-Modified"].wont_be_nil
@@ -85,6 +93,7 @@ describe "App with Riak backend" do
     describe "PUT" do
       before do
         header "Authorization", "Bearer 123"
+        set_usage_size_info "jimmy", "documents", "23"
       end
 
       describe "with implicit content type" do
@@ -102,6 +111,10 @@ describe "App with Riak backend" do
           data_bucket.get("jimmy:documents:bar").content_type.must_equal "text/plain; charset=utf-8"
         end
 
+        it "increases the overall category size" do
+          info_bucket.get("usage:size:jimmy:documents").data.must_equal "35"
+        end
+
         it "indexes the data set" do
           indexes = data_bucket.get("jimmy:documents:bar").indexes
           indexes["user_id_bin"].must_be_kind_of Set
@@ -127,6 +140,10 @@ describe "App with Riak backend" do
           data_bucket.get("jimmy:documents:jason").content_type.must_equal "application/json"
         end
 
+        it "increases the overall category size" do
+          info_bucket.get("usage:size:jimmy:documents").data.must_equal "49"
+        end
+
         it "delivers the data correctly" do
           header "Authorization", "Bearer 123"
           get "/jimmy/documents/jason"
@@ -184,6 +201,40 @@ describe "App with Riak backend" do
         end
       end
 
+      describe "with existing content" do
+        before do
+          set_usage_size_info "jimmy", "documents", "10"
+          put "/jimmy/documents/archive/foo", "lorem ipsum"
+          put "/jimmy/documents/archive/foo", "some awesome content"
+        end
+
+        it "saves the value" do
+          last_response.status.must_equal 200
+          data_bucket.get("jimmy:documents/archive:foo").data.must_equal "some awesome content"
+        end
+
+        it "increases the overall category size" do
+          puts info_bucket.keys.inspect
+          info_bucket.get("usage:size:jimmy:documents").data.must_equal "30"
+        end
+      end
+
+      describe "public data" do
+        before do
+          set_usage_size_info "jimmy", "public/documents", "10"
+          put "/jimmy/public/documents/notes/foo", "note to self"
+        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
+          info_bucket.get("usage:size:jimmy:public/documents").data.must_equal "22"
+        end
+      end
+
       context "with binary data" do
         context "binary charset in content-type header" do
           before do
@@ -207,6 +258,10 @@ describe "App with Riak backend" do
             last_response.body.must_equal @image
           end
 
+          it "increases the overall category size" do
+            info_bucket.get("usage:size:jimmy:documents").data.must_equal "16067"
+          end
+
           it "indexes the binary set" do
             indexes = binary_bucket.get("jimmy:documents:jaypeg").indexes
             indexes["user_id_bin"].must_be_kind_of Set
@@ -291,28 +346,33 @@ describe "App with Riak backend" do
     describe "DELETE" do
       before do
         header "Authorization", "Bearer 123"
+        set_usage_size_info "jimmy", "documents", "123"
+        delete "/jimmy/documents/foo"
       end
 
       it "removes the key" do
-        delete "/jimmy/documents/foo"
-
         last_response.status.must_equal 204
         lambda {
           data_bucket.get("jimmy:documents:foo")
         }.must_raise Riak::HTTPFailedRequest
       end
 
+      it "decreases the overall category size" do
+        info_bucket.get("usage:size:jimmy:documents").data.must_equal "101"
+      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"
+
+          delete "/jimmy/documents/jaypeg"
         end
 
         it "removes the main object" do
-          delete "/jimmy/documents/jaypeg"
-
           last_response.status.must_equal 204
           lambda {
             data_bucket.get("jimmy:documents:jaypeg")
@@ -320,13 +380,15 @@ describe "App with Riak backend" do
         end
 
         it "removes the binary object" do
-          delete "/jimmy/documents/jaypeg"
-
           last_response.status.must_equal 204
           lambda {
             binary_bucket.get("jimmy:documents:jaypeg")
           }.must_raise Riak::HTTPFailedRequest
         end
+
+        it "decreases the overall category size" do
+          info_bucket.get("usage:size:jimmy:documents").data.must_equal "83956"
+        end
       end
     end
   end

spec/spec_helper.rb

@@ -42,8 +42,12 @@ def binary_bucket
   @binary_bucket ||= storage_client.bucket(settings.bucket_config['binaries'])
 end
 
+def info_bucket
+  @info_bucket ||= storage_client.bucket(settings.bucket_config['info'])
+end
+
 def purge_all_buckets
-  [data_bucket, directory_bucket, auth_bucket, binary_bucket].each do |bucket|
+  [data_bucket, directory_bucket, auth_bucket, binary_bucket, info_bucket].each do |bucket|
     bucket.keys.each {|key| bucket.delete key}
   end
 end