get_file_data: avoid exposing users to copy_stream invocation

We'll now accept an optional argument which can be passed to
IO.copy_stream directly.  This should make life easier on users
so they won't be exposed to our internals to make efficient
copies of large files.

5 files changed, 23 insertions, 13 deletions

diff --git a/bin/mog b/bin/mog
index a308db5..93434bc 100755
--- a/bin/mog
+++ b/bin/mog
@@ -128,7 +128,7 @@ begin
       if (!cat[:raw] && key =~ /^_big_info:/)
         mg.bigfile_write(key, STDOUT, {:verify => true})
       else
-        mg.get_file_data(key) { |fp| MogileFS::X.copy_stream(fp, STDOUT) }
+        mg.get_file_data(key, STDOUT)
       end
     end
   when 'ls'
diff --git a/lib/mogilefs/copy_stream.rb b/lib/mogilefs/copy_stream.rb
index b928031..7531f54 100644
--- a/lib/mogilefs/copy_stream.rb
+++ b/lib/mogilefs/copy_stream.rb
@@ -3,24 +3,26 @@
 # internal compatibility class for older Rubies
 module MogileFS::CopyStream # :nodoc:
   def self.copy_stream(src, dst)
-    src_io = String === src ? File.open(src) : src
+    src_io = src.respond_to?(:to_str) ? File.open(src) : src
+    dst_io = dst.respond_to?(:to_str) ? File.open(dst, "w") : dst
     buf = ""
     written = 0
     if src_io.respond_to?(:readpartial)
       begin
         src_io.readpartial(0x4000, buf)
-        written += dst.write(buf)
+        written += dst_io.write(buf)
       rescue EOFError
         break
       end while true
     else
       while src_io.read(0x4000, buf)
-        written += dst.write(buf)
+        written += dst_io.write(buf)
       end
     end
-    dst.flush if dst.respond_to?(:flush)
+    dst_io.flush if dst_io.respond_to?(:flush)
     written
     ensure
-      src_io.close if String === src
+      src_io.close if src_io != src
+      dst_io.close if dst_io != dst
   end
 end
diff --git a/lib/mogilefs/mogilefs.rb b/lib/mogilefs/mogilefs.rb
index 2b02fb7..625e26f 100644
--- a/lib/mogilefs/mogilefs.rb
+++ b/lib/mogilefs/mogilefs.rb
@@ -52,12 +52,22 @@ class MogileFS::MogileFS < MogileFS::Client
   end
 
   ##
-  # Retrieves the contents of +key+.
+  # Retrieves the contents of +key+.  If +dest+ is specified, +dest+
+  # should be an IO-like object capable of receiving the +write+ method
+  # or a path name.
 
-  def get_file_data(key)
+  def get_file_data(key, dest = nil)
     paths = get_paths(key)
     sock = MogileFS::HTTPReader.first(paths, "GET", @get_file_data_timeout)
-    block_given? ? yield(sock) : sock.to_s
+    if dest
+      MogileFS::X.copy_stream(sock, dest)
+    elsif block_given?
+      yield(sock)
+    else
+      sock.to_s
+    end
+    ensure
+      sock.close unless sock.closed?
   end
 
   ##
diff --git a/test/test_mogilefs_integration.rb b/test/test_mogilefs_integration.rb
index 47dcb37..0556c5b 100644
--- a/test/test_mogilefs_integration.rb
+++ b/test/test_mogilefs_integration.rb
@@ -12,9 +12,7 @@ class TestMogileFSIntegration < TestMogIntegration
     assert_equal 4, @client.size("CRUD")
     assert_equal "DATA", @client.get_file_data("CRUD")
     sio = StringIO.new("")
-    rv = @client.get_file_data("CRUD") do |rd|
-      MogileFS::X.copy_stream(rd, sio)
-    end
+    rv = @client.get_file_data("CRUD", sio)
     assert_equal 4, rv
     assert_equal "DATA", sio.string
     assert_equal 8, @client.store_content("CRUD", "default", "MOARDATA")
diff --git a/test/test_mogtool_bigfile.rb b/test/test_mogtool_bigfile.rb
index a6de891..05769cd 100644
--- a/test/test_mogtool_bigfile.rb
+++ b/test/test_mogtool_bigfile.rb
@@ -43,7 +43,7 @@ class TestMogtoolBigfile < TestMogIntegration
     part1 = "#@big_uuid,1"
     tmp = tmpfile("part1")
     before_uris = @client.get_uris(part1)
-    @client.get_file_data(part1) { |sock| MogileFS::X.copy_stream(sock, tmp) }
+    @client.get_file_data(part1, tmp)
     @client.delete(part1)
     @client.store_file(part1, nil, tmp.path)
     wait_for_DELETE(before_uris)