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)