taps2 0.5.1

data/lib/taps/db_session.rb

@@ -0,0 +1,20 @@
+Sequel::Model.db = Sequel.connect(Taps::Config.taps_database_url)
+
+class DbSession < Sequel::Model
+  plugin :schema
+  set_schema do
+    primary_key :id
+    text :key
+    text :database_url
+    timestamp :started_at
+    timestamp :last_access
+  end
+
+  def conn
+    Sequel.connect(database_url) do |db|
+      yield db if block_given?
+    end
+  end
+end
+
+DbSession.create_table! unless DbSession.table_exists?

data/lib/taps/errors.rb

@@ -0,0 +1,15 @@
+module Taps
+  class BaseError < StandardError
+    attr_reader :original_backtrace
+
+    def initialize(message, opts={})
+      @original_backtrace = opts.delete(:backtrace)
+      super(message)
+    end
+  end
+
+  class NotImplemented < BaseError; end
+  class DuplicatePrimaryKeyError < BaseError; end
+  class CorruptedData < BaseError; end
+  class InvalidData < BaseError; end
+end

data/lib/taps/log.rb

@@ -0,0 +1,15 @@
+module Taps
+  def self.log=(log)
+    @@log = log
+  end
+
+  def self.log
+    @@log ||= begin
+      require 'logger'
+      log = Logger.new($stderr)
+      log.level = Logger::ERROR
+      log.datetime_format = "%Y-%m-%d %H:%M:%S"
+      log
+    end
+  end
+end

data/lib/taps/monkey.rb

@@ -0,0 +1,21 @@
+class Hash
+  def symbolize_keys
+    inject({}) do |options, (key, value)|
+      options[(key.to_sym rescue key) || key] = value
+      options
+    end
+  end
+
+  def symbolize_keys!
+    self.replace(symbolize_keys)
+  end
+
+  def symbolize_recursively!
+    self.replace(symbolize_keys)
+    self.each do |k, v|
+      if v.kind_of?(Hash)
+        v.symbolize_keys!
+      end
+    end
+  end
+end

data/lib/taps/multipart.rb

@@ -0,0 +1,73 @@
+require 'restclient'
+require 'rack/utils'
+require 'stringio'
+require 'vendor/okjson'
+
+module Taps
+  class Multipart
+    class Container
+      attr_accessor :attachments
+
+      def initialize
+        @attachments = []
+      end
+
+      def attach(opts)
+        mp = Taps::Multipart.new(opts)
+        attachments << mp
+      end
+
+      def generate
+        hash = {}
+        attachments.each do |mp|
+          hash[mp.name] = mp
+        end
+        m = RestClient::Payload::Multipart.new(hash)
+        [m.to_s, m.headers['Content-Type']]
+      end
+    end
+
+    attr_reader :opts
+
+    def initialize(opts={})
+      @opts = opts
+    end
+
+    def name
+      opts[:name]
+    end
+
+    def to_s
+      opts[:payload]
+    end
+
+    def content_type
+      opts[:content_type] || 'text/plain'
+    end
+
+    def original_filename
+      opts[:original_filename]
+    end
+
+    def self.create
+      c = Taps::Multipart::Container.new
+      yield c
+      c.generate
+    end
+
+    # response is a rest-client response
+    def self.parse(response)
+      content = response.to_s
+      env = {
+        'CONTENT_TYPE' => response.headers[:content_type],
+        'CONTENT_LENGTH' => content.size,
+        'rack.input' => StringIO.new(content)
+      }
+
+      params = Rack::Utils::Multipart.parse_multipart(env)
+      params.symbolize_keys!
+      params
+    end
+
+  end
+end

data/lib/taps/operation.rb

@@ -0,0 +1,578 @@
+require 'rest_client'
+require 'sequel'
+require 'zlib'
+
+require 'taps/progress_bar'
+require 'taps/config'
+require 'taps/utils'
+require 'taps/data_stream'
+require 'taps/errors'
+
+# disable warnings, rest client makes a lot of noise right now
+$VERBOSE = nil
+
+module Taps
+  class Operation
+    attr_reader :database_url, :remote_url, :opts
+    attr_reader :session_uri
+
+    def initialize(database_url, remote_url, opts={})
+      @database_url = database_url
+      @remote_url = remote_url
+      @opts = opts
+      @exiting = false
+      @session_uri = opts[:session_uri]
+    end
+
+    def file_prefix
+      "op"
+    end
+
+    def skip_schema?
+      !!opts[:skip_schema]
+    end
+
+    def indexes_first?
+      !!opts[:indexes_first]
+    end
+
+    def table_filter
+      opts[:table_filter]
+    end
+
+    def exclude_tables
+      opts[:exclude_tables] || []
+    end
+
+    def apply_table_filter(tables)
+      return tables unless table_filter || exclude_tables
+
+      re = table_filter ? Regexp.new(table_filter) : nil
+      if tables.kind_of?(Hash)
+        ntables = {}
+        tables.each do |t, d|
+          if !exclude_tables.include?(t.to_s) && (!re || !re.match(t.to_s).nil?)
+            ntables[t] = d
+          end
+        end
+        ntables
+      else
+        tables.reject { |t| exclude_tables.include?(t.to_s) || (re && re.match(t.to_s).nil?) }
+      end
+    end
+
+    def log
+      Taps.log
+    end
+
+    def store_session
+      file = "#{file_prefix}_#{Time.now.strftime("%Y%m%d%H%M")}.dat"
+      puts "\nSaving session to #{file}.."
+      File.open(file, 'w') do |f|
+        f.write(::OkJson.encode(to_hash))
+      end
+    end
+
+    def to_hash
+      {
+        :klass => self.class.to_s,
+        :database_url => database_url,
+        :remote_url => remote_url,
+        :session_uri => session_uri,
+        :stream_state => stream_state,
+        :completed_tables => completed_tables,
+        :table_filter => table_filter,
+      }
+    end
+
+    def exiting?
+      !!@exiting
+    end
+
+    def setup_signal_trap
+      trap("INT") {
+        puts "\nCompleting current action..."
+        @exiting = true
+      }
+
+      trap("TERM") {
+        puts "\nCompleting current action..."
+        @exiting = true
+      }
+    end
+
+    def resuming?
+      opts[:resume] == true
+    end
+
+    def default_chunksize
+      opts[:default_chunksize]
+    end
+
+    def completed_tables
+      opts[:completed_tables] ||= []
+    end
+
+    def stream_state
+      opts[:stream_state] ||= {}
+    end
+
+    def stream_state=(val)
+      opts[:stream_state] = val
+    end
+
+    def compression_disabled?
+      !!opts[:disable_compression]
+    end
+
+    def db
+      @db ||= Sequel.connect(database_url)
+    end
+
+    def server
+      @server ||= RestClient::Resource.new(remote_url)
+    end
+
+    def session_resource
+      @session_resource ||= begin
+        @session_uri ||= server['sessions'].post('', http_headers).to_s
+        server[@session_uri]
+      end
+    end
+
+    def set_session(uri)
+      session_uri = uri
+      @session_resource = server[session_uri]
+    end
+
+    def close_session
+      @session_resource.delete(http_headers) if @session_resource
+    end
+
+    def safe_url(url)
+      url.sub(/\/\/(.+?)?:(.*?)@/, '//\1:[hidden]@')
+    end
+
+    def safe_remote_url
+      safe_url(remote_url)
+    end
+
+    def safe_database_url
+      safe_url(database_url)
+    end
+
+    def http_headers(extra = {})
+      base = { :taps_version => Taps.version }
+      if compression_disabled?
+        base[:accept_encoding] = ""
+      else
+        base[:accept_encoding] = "gzip, deflate"
+      end
+      base.merge(extra)
+    end
+
+    def format_number(num)
+      num.to_s.gsub(/(\d)(?=(\d\d\d)+(?!\d))/, "\\1,")
+    end
+
+    def verify_server
+      begin
+        server['/'].get(http_headers)
+      rescue RestClient::RequestFailed => e
+        if e.http_code == 417
+          puts "#{safe_remote_url} is running a different minor version of taps."
+          puts "#{e.response.to_s}"
+          exit(1)
+        else
+          raise
+        end
+      rescue RestClient::Unauthorized
+        puts "Bad credentials given for #{safe_remote_url}"
+        exit(1)
+      rescue Errno::ECONNREFUSED
+        puts "Can't connect to #{safe_remote_url}. Please check that it's running"
+        exit(1)
+      end
+    end
+
+    def catch_errors(&blk)
+      verify_server
+
+      begin
+        blk.call
+        close_session
+      rescue RestClient::Exception, Taps::BaseError => e
+        store_session
+        if e.kind_of?(Taps::BaseError)
+          puts "!!! Caught Server Exception"
+          puts "#{e.class}: #{e.message}"
+          puts "\n#{e.original_backtrace}" if e.original_backtrace
+          exit(1)
+        elsif e.respond_to?(:response)
+          puts "!!! Caught Server Exception"
+          puts "HTTP CODE: #{e.http_code}"
+          puts "#{e.response.to_s}"
+          exit(1)
+        else
+          raise
+        end
+      end
+    end
+
+    def self.factory(type, database_url, remote_url, opts)
+      type = :resume if opts[:resume]
+      klass = case type
+        when :pull then Taps::Pull
+        when :push then Taps::Push
+        when :resume then eval(opts[:klass])
+        else raise "Unknown Operation Type -> #{type}"
+      end
+
+      klass.new(database_url, remote_url, opts)
+    end
+  end
+
+  class Pull < Operation
+    def file_prefix
+      "pull"
+    end
+
+    def to_hash
+      super.merge(:remote_tables_info => remote_tables_info)
+    end
+
+    def run
+      catch_errors do
+        unless resuming?
+          pull_schema if !skip_schema?
+          pull_indexes if indexes_first? && !skip_schema?
+        end
+        setup_signal_trap
+        pull_partial_data if resuming?
+        pull_data
+        pull_indexes if !indexes_first? && !skip_schema?
+        pull_reset_sequences
+      end
+    end
+
+    def pull_schema
+      puts "Receiving schema"
+
+      progress = ProgressBar.new('Schema', tables.size)
+      tables.each do |table_name, count|
+        schema_data = session_resource['pull/schema'].post({:table_name => table_name}, http_headers).to_s
+        log.debug "Table: #{table_name}\n#{schema_data}\n"
+        output = Taps::Utils.load_schema(database_url, schema_data)
+        output = output.to_s.strip
+        puts output unless output.empty?
+        progress.inc(1)
+      end
+      progress.finish
+    end
+
+    def pull_data
+      puts "Receiving data"
+
+      puts "#{tables.size} tables, #{format_number(record_count)} records"
+
+      tables.each do |table_name, count|
+        progress = ProgressBar.new(table_name.to_s, count)
+        stream = Taps::DataStream.factory(db, {
+          :chunksize => default_chunksize,
+          :table_name => table_name
+        })
+        pull_data_from_table(stream, progress)
+      end
+    end
+
+    def pull_partial_data
+      return if stream_state == {}
+
+      table_name = stream_state[:table_name]
+      record_count = tables[table_name.to_s]
+      puts "Resuming #{table_name}, #{format_number(record_count)} records"
+
+      progress = ProgressBar.new(table_name.to_s, record_count)
+      stream = Taps::DataStream.factory(db, stream_state)
+      pull_data_from_table(stream, progress)
+    end
+
+    def pull_data_from_table(stream, progress)
+      loop do
+        begin
+          if exiting?
+            store_session
+            exit 0
+          end
+
+          size = stream.fetch_remote(session_resource['pull/table'], http_headers)
+          break if stream.complete?
+          progress.inc(size) unless exiting?
+          stream.error = false
+          self.stream_state = stream.to_hash
+        rescue Taps::CorruptedData => e
+          puts "Corrupted Data Received #{e.message}, retrying..."
+          stream.error = true
+          next
+        end
+      end
+
+      progress.finish
+      completed_tables << stream.table_name.to_s
+      self.stream_state = {}
+    end
+
+    def tables
+      h = {}
+      remote_tables_info.each do |table_name, count|
+        next if completed_tables.include?(table_name.to_s)
+        h[table_name.to_s] = count
+      end
+      h
+    end
+
+    def record_count
+      @record_count ||= remote_tables_info.values.inject(0) { |a,c| a += c }
+    end
+
+    def remote_tables_info
+      opts[:remote_tables_info] ||= fetch_remote_tables_info
+    end
+
+    def fetch_remote_tables_info
+      retries = 0
+      max_retries = 10
+      begin
+        tables = ::OkJson.decode(session_resource['pull/table_names'].get(http_headers).to_s)
+      rescue RestClient::Exception
+        retries += 1
+        retry if retries <= max_retries
+        puts "Unable to fetch tables information from #{remote_url}. Please check the server log."
+        exit(1)
+      end
+
+      data = {}
+      apply_table_filter(tables).each do |table_name|
+        retries = 0
+        begin
+          count = Integer(session_resource['pull/table_count'].post({:table => table_name}, http_headers).to_s)
+          data[table_name] = count
+        rescue RestClient::Exception
+          retries += 1
+          retry if retries <= max_retries
+          puts "Unable to fetch tables information from #{remote_url}. Please check the server log."
+          exit(1)
+        end
+      end
+      data
+    end
+
+    def pull_indexes
+      puts "Receiving indexes"
+
+      idxs = ::OkJson.decode(session_resource['pull/indexes'].get(http_headers).to_s)
+
+      apply_table_filter(idxs).each do |table, indexes|
+        next unless indexes.size > 0
+        progress = ProgressBar.new(table, indexes.size)
+        indexes.each do |idx|
+          output = Taps::Utils.load_indexes(database_url, idx)
+          output = output.to_s.strip
+          puts output unless output.empty?
+          progress.inc(1)
+        end
+        progress.finish
+      end
+    end
+
+    def pull_reset_sequences
+      puts "Resetting sequences"
+
+      output = Taps::Utils.schema_bin(:reset_db_sequences, database_url)
+      output = output.to_s.strip
+      puts output unless output.empty?
+    end
+  end
+
+  class Push < Operation
+    def file_prefix
+      "push"
+    end
+
+    def to_hash
+      super.merge(:local_tables_info => local_tables_info)
+    end
+
+    def run
+      catch_errors do
+        unless resuming?
+          push_schema if !skip_schema?
+          push_indexes if indexes_first? && !skip_schema?
+        end
+        setup_signal_trap
+        push_partial_data if resuming?
+        push_data
+        push_indexes if !indexes_first? && !skip_schema?
+        push_reset_sequences
+      end
+    end
+
+    def push_indexes
+      idxs = ::OkJson.decode(Taps::Utils.schema_bin(:indexes_individual, database_url))
+
+      return unless idxs.size > 0
+
+      puts "Sending indexes"
+
+      apply_table_filter(idxs).each do |table, indexes|
+        next unless indexes.size > 0
+        progress = ProgressBar.new(table, indexes.size)
+        indexes.each do |idx|
+          session_resource['push/indexes'].post(idx, http_headers)
+          progress.inc(1)
+        end
+        progress.finish
+      end
+    end
+
+    def push_schema
+      puts "Sending schema"
+
+      progress = ProgressBar.new('Schema', tables.size)
+      tables.each do |table, count|
+        schema_data = Taps::Utils.schema_bin(:dump_table, database_url, table)
+        log.debug "Table: #{table}\n#{schema_data}\n"
+        session_resource['push/schema'].post(schema_data, http_headers)
+        progress.inc(1)
+      end
+      progress.finish
+    end
+
+    def push_reset_sequences
+      puts "Resetting sequences"
+
+      session_resource['push/reset_sequences'].post('', http_headers)
+    end
+
+    def push_partial_data
+      return if stream_state == {}
+
+      table_name = stream_state[:table_name]
+      record_count = tables[table_name.to_s]
+      puts "Resuming #{table_name}, #{format_number(record_count)} records"
+      progress = ProgressBar.new(table_name.to_s, record_count)
+      stream = Taps::DataStream.factory(db, stream_state)
+      push_data_from_table(stream, progress)
+    end
+
+    def push_data
+      puts "Sending data"
+
+      puts "#{tables.size} tables, #{format_number(record_count)} records"
+
+      tables.each do |table_name, count|
+        stream = Taps::DataStream.factory(db,
+          :table_name => table_name,
+          :chunksize => default_chunksize)
+        progress = ProgressBar.new(table_name.to_s, count)
+        push_data_from_table(stream, progress)
+      end
+    end
+
+    def push_data_from_table(stream, progress)
+      loop do
+        if exiting?
+          store_session
+          exit 0
+        end
+
+        row_size = 0
+        chunksize = stream.state[:chunksize]
+
+        begin
+          chunksize = Taps::Utils.calculate_chunksize(chunksize) do |c|
+            stream.state[:chunksize] = c.to_i
+            encoded_data, row_size, elapsed_time = nil
+            d1 = c.time_delta do
+              encoded_data, row_size, elapsed_time = stream.fetch
+            end
+            break if stream.complete?
+
+            data = nil
+            d2 = c.time_delta do
+              data = {
+                :state => stream.to_hash,
+                :checksum => Taps::Utils.checksum(encoded_data).to_s
+              }
+            end
+
+            begin
+              content, content_type = nil
+              d3 = c.time_delta do
+                content, content_type = Taps::Multipart.create do |r|
+                  r.attach :name => :encoded_data,
+                    :payload => encoded_data,
+                    :content_type => 'application/octet-stream'
+                  r.attach :name => :json,
+                    :payload => ::OkJson.encode(data),
+                    :content_type => 'application/json'
+                end
+              end
+              session_resource['push/table'].post(content, http_headers(:content_type => content_type))
+              self.stream_state = stream.to_hash
+            rescue => e
+              Taps::Utils.reraise_server_exception(e)
+            end
+
+            c.idle_secs = (d1 + d2 + d3)
+
+            elapsed_time
+          end
+        rescue Taps::CorruptedData => e
+          # retry the same data, it got corrupted somehow.
+          next
+        rescue Taps::DuplicatePrimaryKeyError => e
+          # verify the stream and retry it
+          stream = stream.verify_remote_stream(session_resource['push/verify_stream'], http_headers)
+          next
+        end
+        stream.state[:chunksize] = chunksize
+
+        progress.inc(row_size)
+
+        stream.increment(row_size)
+        break if stream.complete?
+      end
+
+      progress.finish
+      completed_tables << stream.table_name.to_s
+      self.stream_state = {}
+    end
+
+    def local_tables_info
+      opts[:local_tables_info] ||= fetch_local_tables_info
+    end
+
+    def tables
+      h = {}
+      local_tables_info.each do |table_name, count|
+        next if completed_tables.include?(table_name.to_s)
+        h[table_name.to_s] = count
+      end
+      h
+    end
+
+    def record_count
+      @record_count ||= local_tables_info.values.inject(0) { |a,c| a += c }
+    end
+
+    def fetch_local_tables_info
+      tables_with_counts = {}
+      db.tables.each do |table|
+        tables_with_counts[table] = db[table.to_sym.identifier].count
+      end
+      apply_table_filter(tables_with_counts)
+    end
+
+  end
+end