diff --git a/lib/remote_storage/riak.rb b/lib/remote_storage/riak.rb index 7ead0cd..ac8f927 100644 --- a/lib/remote_storage/riak.rb +++ b/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.content_type = content_type || "text/plain; charset=utf-8" + object = build_data_object(user, directory, key, data, content_type) - directory_index = directory == "" ? "/" : directory - object.indexes.merge!({:user_id_bin => [user], - :directory_bin => [CGI.escape(directory_index)]}) + existing_object_size = object_size(object) 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 = "/" diff --git a/spec/riak_spec.rb b/spec/riak_spec.rb index 2bc5f80..6f42e63 100644 --- a/spec/riak_spec.rb +++ b/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 diff --git a/spec/spec_helper.rb b/spec/spec_helper.rb index 7172168..44c2063 100644 --- a/spec/spec_helper.rb +++ b/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