get_file_data supports offset and count

Just like IO.copy_stream

3 files changed, 22 insertions, 9 deletions

diff --git a/lib/mogilefs/http_reader.rb b/lib/mogilefs/http_reader.rb
index ca0e46b..d7a36e6 100644
--- a/lib/mogilefs/http_reader.rb
+++ b/lib/mogilefs/http_reader.rb
@@ -19,11 +19,17 @@ class MogileFS::HTTPReader < MogileFS::Socket
           "read=#{buf.size} bytes, expected=#@content_length from #@uri", []
   end
 
-  def self.first(paths, timeout)
+  def self.first(paths, timeout, count = nil, offset = nil)
     errors = nil
+    if offset || count
+      offset ||= 0
+      range_end = count ? offset + count - 1 : ""
+      range = "Range: bytes=#{offset}-#{range_end}\r\n"
+    end
+
     paths.each do |path|
       begin
-        sock = try(path, timeout) and return sock
+        sock = try(path, timeout, range) and return sock
       rescue => e
         errors ||= []
         errors << "#{path} - #{e.message} (#{e.class})"
@@ -35,10 +41,10 @@ class MogileFS::HTTPReader < MogileFS::Socket
 
   # given a path, this returns a readable socket with ready data from the
   # body of the response.
-  def self.try(path, timeout)
+  def self.try(path, timeout, range) # :nodoc:
     uri = URI.parse(path)
     sock = tcp(uri.host, uri.port, timeout)
-    buf = "GET #{uri.request_uri} HTTP/1.0\r\n\r\n" # no chunking
+    buf = "GET #{uri.request_uri} HTTP/1.0\r\n#{range}\r\n" # no chunking
     sock.timed_write(buf, timeout)
 
     sock.timed_peek(2048, buf, timeout) or
@@ -48,16 +54,19 @@ class MogileFS::HTTPReader < MogileFS::Socket
 
     # we're dealing with a seriously slow/stupid HTTP server if we can't
     # get the header in a single recv(2) syscall.
-    if head =~ %r{\AHTTP/\d+\.\d+\s+200\s*} &&
+    if ((range && head =~ %r{\AHTTP/\d+\.\d+\s+206\s*}) ||
+        (!range && head =~ %r{\AHTTP/\d+\.\d+\s+200\s*})) &&
        head =~ %r{^Content-Length:\s*(\d+)}i
       sock.content_length = $1.to_i
       sock.uri = uri
       sock.timed_read(head.bytesize + 4, buf, 0)
       return sock
     end
-    raise MogileFS::InvalidResponseError, "header=#{head.inspect}", []
+    msg = range ? "Expected 206 w/#{range.strip}: " : "header="
+    msg << head.inspect
+    raise MogileFS::InvalidResponseError, msg, []
   rescue
-    sock.close if sock && ! sock.closed?
+    sock.close if sock
     raise
   end
 end
diff --git a/lib/mogilefs/mogilefs.rb b/lib/mogilefs/mogilefs.rb
index f132dd7..86eba42 100644
--- a/lib/mogilefs/mogilefs.rb
+++ b/lib/mogilefs/mogilefs.rb
@@ -57,9 +57,10 @@ class MogileFS::MogileFS < MogileFS::Client
   # should be an IO-like object capable of receiving the +write+ method
   # or a path name.
 
-  def get_file_data(key, dest = nil)
+  def get_file_data(key, dest = nil, count = nil, offset = nil)
     paths = get_paths(key)
-    sock = MogileFS::HTTPReader.first(paths, @get_file_data_timeout)
+    sock = MogileFS::HTTPReader.first(paths, @get_file_data_timeout,
+                                      count, offset)
     if dest
       MogileFS::X.copy_stream(sock, dest)
     elsif block_given?
diff --git a/test/test_mogilefs_integration.rb b/test/test_mogilefs_integration.rb
index 0556c5b..e42185a 100644
--- a/test/test_mogilefs_integration.rb
+++ b/test/test_mogilefs_integration.rb
@@ -11,6 +11,9 @@ class TestMogileFSIntegration < TestMogIntegration
     assert_equal 4, @client.store_content("CRUD", "default", "DATA")
     assert_equal 4, @client.size("CRUD")
     assert_equal "DATA", @client.get_file_data("CRUD")
+    assert_equal "DAT", @client.get_file_data("CRUD", nil, 3)
+    assert_equal "AT", @client.get_file_data("CRUD", nil, 2, 1)
+
     sio = StringIO.new("")
     rv = @client.get_file_data("CRUD", sio)
     assert_equal 4, rv