From e6ef07b2576e802a1db3b2c3764d2bb678ceddb2 Mon Sep 17 00:00:00 2001 From: "E.J. Finneran" Date: Wed, 19 Dec 2012 22:17:29 -0800 Subject: [PATCH] Allow dumping to a Base64 encoded string for Heroku compatibility Heroku can't handle non-UTF-8 encoded data being output from a 'heroku run' command. This change allows replicate to dump and load Base64 encoded objects so Heroku can handle it. Heroku also sends stderr output to stdout so the loading code attempts to be smart about removing non-base64 lines from the input before loading it. --- README.md | 13 +++++++++++++ bin/replicate | 36 ++++++++++++++++++++++++++++++++++-- lib/replicate.rb | 1 + lib/replicate/serializer.rb | 25 +++++++++++++++++++++++++ test/dumper_test.rb | 12 +++++++++++- 5 files changed, 84 insertions(+), 3 deletions(-) create mode 100644 lib/replicate/serializer.rb diff --git a/README.md b/README.md index 06c0b3d..77b1540 100644 --- a/README.md +++ b/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 ------------ diff --git a/bin/replicate b/bin/replicate index 55958b3..29cdb19 100755 --- a/bin/replicate +++ b/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, '', 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 diff --git a/lib/replicate.rb b/lib/replicate.rb index bc13646..f2f8137 100644 --- a/lib/replicate.rb +++ b/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 diff --git a/lib/replicate/serializer.rb b/lib/replicate/serializer.rb new file mode 100644 index 0000000..4eed2bc --- /dev/null +++ b/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 diff --git a/test/dumper_test.rb b/test/dumper_test.rb index e5b0fd3..0385923 100644 --- a/test/dumper_test.rb +++ b/test/dumper_test.rb @@ -45,7 +45,7 @@ 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 @@ -53,6 +53,16 @@ def test_writing_to_io 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)