Using Replicate with Heroku

README.md

@@ -97,6 +97,19 @@ Run the dump script:
     $ remote_command="replicate -r /app/config/environment -d 'User.find(1234)'"
     $ ssh example.org "$remote_command" |replicate -r ./config/environment -l
 
+### Dumping and loading from Heroku
+
+Use the -h options when dumping data from Heroku and use the -h or -b options when loading the data.
+
+	$ remote_command="replicate -r /app/config/environment -dh 'User.find(1234)'"
+	$ heroku run "$remote_command" | replicate -r ./config/environment -lb
+
+	Ignoring line: Running `replicate -r ./config/environment -dbq "Post.all"` attached to terminal... up, run.6283
+    ==> loaded 6 total objects:
+    Post      4
+    User      2
+
+
 ActiveRecord
 ------------
 

bin/replicate

@@ -22,6 +22,8 @@
 #/ Options:
 #/   -r, --require        Require the library. Often used with 'config/environment'.
 #/   -i, --keep-id        Use replicated ids when loading dump file.
+#/   -b, --base64         Use Base64 encoding when dumping and loading
+#/   -h, --heroku         Heroku friendly settings. Base64 encoding and no logging to stderr (Equivilant to -b -q options)
 #/   -f, --force          Allow loading in production environments.
 #/   -v, --verbose        Write more status output.
 #/   -q, --quiet          Write less status output.
@@ -33,6 +35,7 @@ require 'optparse'
 # default options
 mode     = nil
 verbose  = false
+base64   = false
 quiet    = false
 keep_id  = false
 out      = STDOUT
@@ -47,6 +50,8 @@ ARGV.options do |opts|
   opts.on("-l", "--load")      { mode = :load }
   opts.on("-r", "--require=f") { |file| require file }
   opts.on("-v", "--verbose")   { verbose = true }
+  opts.on("-b", "--base64")    { base64 = true }
+  opts.on("-h", "--heroku")    { base64 = true; quiet = true }
   opts.on("-q", "--quiet")     { quiet = true }
   opts.on("-i", "--keep-id")   { keep_id = true }
   opts.on("--force")           { force = true }
@@ -67,8 +72,12 @@ end
 if mode == :dump
   script = ARGV.shift
   usage.call if script.empty?
+  serializer = Replicate::Serializer.new
+  if base64
+    serializer.mode = :base64
+  end
   Replicate::Dumper.new do |dumper|
-    dumper.marshal_to out
+    dumper.marshal_to serializer
     dumper.log_to $stderr, verbose, quiet
     if script == '-'
       code = $stdin.read
@@ -79,18 +88,41 @@ if mode == :dump
       objects = dumper.instance_eval(script, '<argv>', 0)
       dumper.dump objects
     end
+    serializer.flush
   end
 
 # load mode means we're reading objects from stdin and creating them under
 # the current environment.
 elsif mode == :load
+  require 'base64'
+  if base64
+    # Remove newlines and other characters that Base64 doesn't like
+    input = $stdin.readlines.map(&:strip)
+
+    begin
+      data = Base64.strict_decode64(input.join)
+    rescue ArgumentError
+      # Heroku injects lines into STDOUT so we need to account for that here:
+      if input.size > 1
+        ignored_line = input.shift
+        $stderr.puts "Ignoring line: #{ignored_line}"
+        retry
+      else
+        abort "Check that the input is actually Base64 encoded"
+      end
+    end
+
+    io = StringIO.new(data)
+  else
+    io = $stdin
+  end
   if Replicate.production_environment? && !force
     abort "error: refusing to load in production environment\n" +
           "       manual override: #{File.basename($0)} --force #{original_argv.join(' ')}"
   else
     Replicate::Loader.new do |loader|
       loader.log_to $stderr, verbose, quiet
-      loader.read $stdin
+      loader.read io
     end
   end
 

lib/replicate.rb

@@ -4,6 +4,7 @@ module Replicate
   autoload :Loader,  'replicate/loader'
   autoload :Object,  'replicate/object'
   autoload :Status,  'replicate/status'
+  autoload :Serializer, 'replicate/serializer'
   autoload :AR,      'replicate/active_record'
 
   # Determine if this is a production looking environment. Used in bin/replicate

lib/replicate/serializer.rb

@@ -0,0 +1,25 @@
+require 'base64'
+require 'stringio'
+
+module Replicate
+  class Serializer < StringIO
+    attr_accessor :write_to, :mode
+
+    def write_to
+      @write_to ||= STDOUT
+    end
+
+    def flush
+      write_to.puts(self.string)
+    end
+
+    def string
+      if mode == :base64
+        Base64.strict_encode64(super)
+      else
+        super
+      end
+    end
+
+  end
+end

test/dumper_test.rb

@@ -45,14 +45,24 @@ def test_never_dumps_objects_more_than_once
   end
 
   def test_writing_to_io
-    io = StringIO.new
+    io = Replicate::Serializer.new
     io.set_encoding 'BINARY' if io.respond_to?(:set_encoding)
     @dumper.marshal_to io
     @dumper.dump object = thing
     data = Marshal.dump(['Replicate::Object', object.id, object.attributes])
     assert_equal data, io.string
   end
 
+  def test_writing_to_io_with_base64_encoding
+    io = Replicate::Serializer.new
+    io.mode = :base64
+    io.set_encoding 'BINARY' if io.respond_to?(:set_encoding)
+    @dumper.marshal_to io
+    @dumper.dump object = thing
+    data = Base64.strict_encode64(Marshal.dump(['Replicate::Object', object.id, object.attributes]))
+    assert_equal data, io.string
+  end
+
   def test_stats
     10.times { @dumper.dump thing }
     assert_equal({'Replicate::Object' => 10}, @dumper.stats)