mosql 0.2.0 → 0.3.0

data/lib/mosql/streamer.rb

@@ -0,0 +1,227 @@
+module MoSQL
+  class Streamer
+    include MoSQL::Logging
+
+    BATCH = 1000
+
+    attr_reader :options, :tailer
+
+    NEW_KEYS = [:options, :tailer, :mongo, :sql, :schema]
+
+    def initialize(opts)
+      NEW_KEYS.each do |parm|
+        unless opts.key?(parm)
+          raise ArgumentError.new("Required argument `#{parm}' not provided to #{self.class.name}#new.")
+        end
+        instance_variable_set(:"@#{parm.to_s}", opts[parm])
+      end
+
+      @done    = false
+    end
+
+    def stop
+      @done = true
+    end
+
+    def import
+      if options[:reimport] || tailer.read_timestamp.seconds == 0
+        initial_import
+      end
+    end
+
+    def collection_for_ns(ns)
+      dbname, collection = ns.split(".", 2)
+      @mongo.db(dbname).collection(collection)
+    end
+
+    def unsafe_handle_exceptions(ns, obj)
+      begin
+        yield
+      rescue Sequel::DatabaseError => e
+        wrapped = e.wrapped_exception
+        if wrapped.result && options[:unsafe]
+          log.warn("Ignoring row (#{obj.inspect}): #{e}")
+        else
+          log.error("Error processing #{obj.inspect} for #{ns}.")
+          raise e
+        end
+      end
+    end
+
+    def bulk_upsert(table, ns, items)
+      begin
+        @schema.copy_data(table.db, ns, items)
+      rescue Sequel::DatabaseError => e
+        log.debug("Bulk insert error (#{e}), attempting invidual upserts...")
+        cols = @schema.all_columns(@schema.find_ns(ns))
+        items.each do |it|
+          h = {}
+          cols.zip(it).each { |k,v| h[k] = v }
+          unsafe_handle_exceptions(ns, h) do
+            @sql.upsert!(table, @schema.primary_sql_key_for_ns(ns), h)
+          end
+        end
+      end
+    end
+
+    def with_retries(tries=10)
+      tries.times do |try|
+        begin
+          yield
+        rescue Mongo::ConnectionError, Mongo::ConnectionFailure, Mongo::OperationFailure => e
+          # Duplicate key error
+          raise if e.kind_of?(Mongo::OperationFailure) && [11000, 11001].include?(e.error_code)
+          # Cursor timeout
+          raise if e.kind_of?(Mongo::OperationFailure) && e.message =~ /^Query response returned CURSOR_NOT_FOUND/
+          delay = 0.5 * (1.5 ** try)
+          log.warn("Mongo exception: #{e}, sleeping #{delay}s...")
+          sleep(delay)
+        end
+      end
+    end
+
+    def track_time
+      start = Time.now
+      yield
+      Time.now - start
+    end
+
+    def initial_import
+      @schema.create_schema(@sql.db, !options[:no_drop_tables])
+
+      unless options[:skip_tail]
+        start_ts = @mongo['local']['oplog.rs'].find_one({}, {:sort => [['$natural', -1]]})['ts']
+      end
+
+      @mongo.database_names.each do |dbname|
+        next unless spec = @schema.find_db(dbname)
+        log.info("Importing for Mongo DB #{dbname}...")
+        db = @mongo.db(dbname)
+        db.collections.select { |c| spec.key?(c.name) }.each do |collection|
+          ns = "#{dbname}.#{collection.name}"
+          import_collection(ns, collection)
+          exit(0) if @done
+        end
+      end
+
+      tailer.write_timestamp(start_ts) unless options[:skip_tail]
+    end
+
+    def did_truncate; @did_truncate ||= {}; end
+
+    def import_collection(ns, collection)
+      log.info("Importing for #{ns}...")
+      count = 0
+      batch = []
+      table = @sql.table_for_ns(ns)
+      unless options[:no_drop_tables] || did_truncate[table.first_source]
+        table.truncate
+        did_truncate[table.first_source] = true
+      end
+
+      start    = Time.now
+      sql_time = 0
+      collection.find(nil, :batch_size => BATCH) do |cursor|
+        with_retries do
+          cursor.each do |obj|
+            batch << @schema.transform(ns, obj)
+            count += 1
+
+            if batch.length >= BATCH
+              sql_time += track_time do
+                bulk_upsert(table, ns, batch)
+              end
+              elapsed = Time.now - start
+              log.info("Imported #{count} rows (#{elapsed}s, #{sql_time}s SQL)...")
+              batch.clear
+              exit(0) if @done
+            end
+          end
+        end
+      end
+
+      unless batch.empty?
+        bulk_upsert(table, ns, batch)
+      end
+    end
+
+    def optail
+      tailer.tail_from(options[:tail_from] ?
+                       BSON::Timestamp.new(options[:tail_from].to_i, 0) :
+                       nil)
+      until @done
+        tailer.stream(1000) do |op|
+          handle_op(op)
+        end
+      end
+    end
+
+    def sync_object(ns, _id)
+      primary_sql_key = @schema.primary_sql_key_for_ns(ns)
+      sqlid           = @sql.transform_one_ns(ns, { '_id' => _id })[primary_sql_key]
+      obj             = collection_for_ns(ns).find_one({:_id => _id})
+      if obj
+        unsafe_handle_exceptions(ns, obj) do
+          @sql.upsert_ns(ns, obj)
+        end
+      else
+        @sql.table_for_ns(ns).where(primary_sql_key.to_sym => sqlid).delete()
+      end
+    end
+
+    def handle_op(op)
+      log.debug("processing op: #{op.inspect}")
+      unless op['ns'] && op['op']
+        log.warn("Weird op: #{op.inspect}")
+        return
+      end
+
+      unless @schema.find_ns(op['ns'])
+        log.debug("Skipping op for unknown ns #{op['ns']}...")
+        return
+      end
+
+      ns = op['ns']
+      dbname, collection_name = ns.split(".", 2)
+
+      case op['op']
+      when 'n'
+        log.debug("Skipping no-op #{op.inspect}")
+      when 'i'
+        if collection_name == 'system.indexes'
+          log.info("Skipping index update: #{op.inspect}")
+        else
+          unsafe_handle_exceptions(ns, op['o'])  do
+            @sql.upsert_ns(ns, op['o'])
+          end
+        end
+      when 'u'
+        selector = op['o2']
+        update   = op['o']
+        if update.keys.any? { |k| k.start_with? '$' }
+          log.debug("resync #{ns}: #{selector['_id']} (update was: #{update.inspect})")
+          sync_object(ns, selector['_id'])
+        else
+          log.debug("upsert #{ns}: _id=#{selector['_id']}")
+
+          # The update operation replaces the existing object, but
+          # preserves its _id field, so grab the _id off of the
+          # 'query' field -- it's not guaranteed to be present on the
+          # update.
+          update = { '_id' => selector['_id'] }.merge(update)
+          unsafe_handle_exceptions(ns, update) do
+            @sql.upsert_ns(ns, update)
+          end
+        end
+      when 'd'
+        if options[:ignore_delete]
+          log.debug("Ignoring delete op on #{ns} as instructed.")
+        else
+          @sql.delete_ns(ns, op['o'])
+        end
+      else
+        log.info("Skipping unknown op #{op.inspect}")
+      end
+    end
+  end
+end

data/lib/mosql/tailer.rb

@@ -25,12 +25,15 @@ module MoSQL
     end
 
     def write_timestamp(ts)
-      begin
-        @table.insert({:service => @service, :timestamp => ts.seconds})
-      rescue Sequel::DatabaseError => e
-        raise unless e.message =~ /duplicate key value violates unique constraint/
-        @table.where(:service => @service).update(:timestamp => ts.seconds)
+      unless @did_insert
+        begin
+          @table.insert({:service => @service, :timestamp => ts.seconds})
+        rescue Sequel::DatabaseError => e
+          raise unless MoSQL::SQLAdapter.duplicate_key_error?(e)
+        end
+        @did_insert = true
       end
+      @table.where(:service => @service).update(:timestamp => ts.seconds)
     end
   end
 end

data/lib/mosql/version.rb

@@ -1,3 +1,3 @@
 module MoSQL
-  VERSION = "0.2.0"
+  VERSION = "0.3.0"
 end

data/lib/mosql.rb

@@ -8,4 +8,5 @@ require 'mosql/version'
 require 'mosql/log'
 require 'mosql/sql'
 require 'mosql/schema'
+require 'mosql/streamer'
 require 'mosql/tailer'

data/test/functional/schema.rb

@@ -101,4 +101,46 @@ EOF
     table.insert(row)
     assert_equal(o['_id'].to_s, table.select.first[:_id])
   end
+
+  describe 'special fields' do
+  SPECIAL_MAP = <<EOF
+---
+db:
+  collection:
+    :meta:
+      :table: special
+    :columns:
+      - _id: TEXT
+      - mosql_updated:
+        :source: $timestamp
+        :type: timestamp
+EOF
+
+    before do
+      @specialmap = MoSQL::Schema.new(YAML.load(SPECIAL_MAP))
+
+      @sequel.drop_table?(:special)
+      @specialmap.create_schema(@sequel)
+    end
+
+    it 'sets a default on the column' do
+      @sequel[:special].insert({_id: 'a'})
+      row = @sequel[:special].select.first
+      assert_instance_of(Time, row[:mosql_updated])
+    end
+
+    it 'Can populate $timestamp on COPY' do
+      objects = [
+                 {'_id' => "a"},
+                 {'_id' => "b"}
+                ]
+      before = @sequel.select(Sequel.function(:NOW)).first[:now]
+      @specialmap.copy_data(@sequel, 'db.collection',
+                            objects.map { |o| @specialmap.transform('db.collection', o) } )
+      after = @sequel.select(Sequel.function(:NOW)).first[:now]
+      rows = @sequel[:special].select.sort_by { |r| r[:_id] }
+      assert_instance_of(Time, rows[0][:mosql_updated])
+      assert(rows[0][:mosql_updated] > before && rows[0][:mosql_updated] < after)
+    end
+  end
 end

data/test/functional/streamer.rb

@@ -0,0 +1,213 @@
+require File.join(File.dirname(__FILE__), '_lib.rb')
+require 'mosql/cli'
+
+class MoSQL::Test::Functional::StreamerTest < MoSQL::Test::Functional
+  def build_streamer
+    MoSQL::Streamer.new(:mongo => mongo,
+                        :tailer => nil,
+                        :options => {},
+                        :sql => @adapter,
+                        :schema => @map)
+  end
+
+  describe 'with a basic schema' do
+    TEST_MAP = <<EOF
+---
+mosql_test:
+  collection:
+    :meta:
+      :table: sqltable
+    :columns:
+      - _id: TEXT
+      - var: INTEGER
+  renameid:
+    :meta:
+      :table: sqltable2
+    :columns:
+      - id:
+        :source: _id
+        :type: TEXT
+      - goats: INTEGER
+EOF
+
+    before do
+      @map = MoSQL::Schema.new(YAML.load(TEST_MAP))
+      @adapter = MoSQL::SQLAdapter.new(@map, sql_test_uri)
+
+      @sequel.drop_table?(:sqltable)
+      @sequel.drop_table?(:sqltable2)
+      @map.create_schema(@sequel)
+
+      @streamer = build_streamer
+    end
+
+    it 'handle "u" ops without _id' do
+      o = { '_id' => BSON::ObjectId.new, 'var' => 17 }
+      @adapter.upsert_ns('mosql_test.collection', o)
+
+      @streamer.handle_op({ 'ns' => 'mosql_test.collection',
+                            'op' => 'u',
+                            'o2' => { '_id' => o['_id'] },
+                            'o'  => { 'var' => 27 }
+                          })
+      assert_equal(27, sequel[:sqltable].where(:_id => o['_id'].to_s).select.first[:var])
+    end
+
+    it 'handle "d" ops with BSON::ObjectIds' do
+      o = { '_id' => BSON::ObjectId.new, 'var' => 17 }
+      @adapter.upsert_ns('mosql_test.collection', o)
+
+      @streamer.handle_op({ 'ns' => 'mosql_test.collection',
+                            'op' => 'd',
+                            'o' => { '_id' => o['_id'] },
+                          })
+      assert_equal(0, sequel[:sqltable].where(:_id => o['_id'].to_s).count)
+    end
+
+    it 'handle "u" ops with $set and BSON::ObjectIDs' do
+      o = { '_id' => BSON::ObjectId.new, 'var' => 17 }
+      @adapter.upsert_ns('mosql_test.collection', o)
+
+      # $set's are currently a bit of a hack where we read the object
+      # from the db, so make sure the new object exists in mongo
+      connect_mongo['mosql_test']['collection'].insert(o.merge('var' => 100),
+                                                       :w => 1)
+
+      @streamer.handle_op({ 'ns' => 'mosql_test.collection',
+                            'op' => 'u',
+                            'o2' => { '_id' => o['_id'] },
+                            'o'  => { '$set' => { 'var' => 100 } },
+                          })
+      assert_equal(100, sequel[:sqltable].where(:_id => o['_id'].to_s).select.first[:var])
+    end
+
+    it 'handle "u" ops with $set and a renamed _id' do
+      o = { '_id' => BSON::ObjectId.new, 'goats' => 96 }
+      @adapter.upsert_ns('mosql_test.renameid', o)
+
+      # $set's are currently a bit of a hack where we read the object
+      # from the db, so make sure the new object exists in mongo
+      connect_mongo['mosql_test']['renameid'].insert(o.merge('goats' => 0),
+                                                     :w => 1)
+
+      @streamer.handle_op({ 'ns' => 'mosql_test.renameid',
+                            'op' => 'u',
+                            'o2' => { '_id' => o['_id'] },
+                            'o'  => { '$set' => { 'goats' => 0 } },
+                          })
+      assert_equal(0, sequel[:sqltable2].where(:id => o['_id'].to_s).select.first[:goats])
+    end
+
+    it 'handles "d" ops with a renamed id' do
+      o = { '_id' => BSON::ObjectId.new, 'goats' => 1 }
+      @adapter.upsert_ns('mosql_test.renameid', o)
+
+      @streamer.handle_op({ 'ns' => 'mosql_test.renameid',
+                            'op' => 'd',
+                            'o' => { '_id' => o['_id'] },
+                          })
+      assert_equal(0, sequel[:sqltable2].where(:id => o['_id'].to_s).count)
+    end
+
+    describe '.bulk_upsert' do
+      it 'inserts multiple rows' do
+        objs = [
+                { '_id' => BSON::ObjectId.new, 'var' => 0 },
+                { '_id' => BSON::ObjectId.new, 'var' => 1 },
+                { '_id' => BSON::ObjectId.new, 'var' => 3 },
+               ].map { |o| @map.transform('mosql_test.collection', o) }
+
+        @streamer.bulk_upsert(sequel[:sqltable], 'mosql_test.collection',
+                              objs)
+
+        assert(sequel[:sqltable].where(:_id => objs[0].first, :var => 0).count)
+        assert(sequel[:sqltable].where(:_id => objs[1].first, :var => 1).count)
+        assert(sequel[:sqltable].where(:_id => objs[2].first, :var => 3).count)
+      end
+
+      it 'upserts' do
+        _id = BSON::ObjectId.new
+        objs = [
+                { '_id' => _id, 'var' => 0 },
+                { '_id' => BSON::ObjectId.new, 'var' => 1 },
+                { '_id' => BSON::ObjectId.new, 'var' => 3 },
+               ].map { |o| @map.transform('mosql_test.collection', o) }
+
+        @streamer.bulk_upsert(sequel[:sqltable], 'mosql_test.collection',
+                              objs)
+
+        newobjs = [
+                   { '_id' => _id, 'var' => 117 },
+                   { '_id' => BSON::ObjectId.new, 'var' => 32 },
+                  ].map { |o| @map.transform('mosql_test.collection', o) }
+        @streamer.bulk_upsert(sequel[:sqltable], 'mosql_test.collection',
+                              newobjs)
+
+
+        assert(sequel[:sqltable].where(:_id => newobjs[0].first, :var => 117).count)
+        assert(sequel[:sqltable].where(:_id => newobjs[1].first, :var => 32).count)
+      end
+
+      describe 'when working with --unsafe' do
+        it 'raises on error by default' do
+          assert_raises(Sequel::DatabaseError) do
+            @streamer.handle_op({ 'ns' => 'mosql_test.collection',
+                                  'op' => 'u',
+                                  'o2' => { '_id' => 'a' },
+                                  'o'  => { 'var' => 1 << 70 },
+                                })
+          end
+        end
+
+        it 'does not raises on error with :unsafe' do
+          @streamer.options[:unsafe] = true
+          @streamer.handle_op({ 'ns' => 'mosql_test.collection',
+                                'op' => 'u',
+                                'o2' => { '_id' => 'a' },
+                                'o'  => { 'var' => 1 << 70 },
+                              })
+          assert_equal(0, sequel[:sqltable].where(:_id => 'a').count)
+        end
+      end
+    end
+  end
+
+  describe 'when dealing with aliased dbs' do
+  ALIAS_MAP = <<EOF
+---
+test:
+  :meta:
+    :alias: test_[0-9]+
+  collection:
+    :meta:
+      :table: sqltable
+    :columns:
+      - _id: TEXT
+      - var: INTEGER
+EOF
+    before do
+      @map = MoSQL::Schema.new(YAML.load(ALIAS_MAP))
+      @adapter = MoSQL::SQLAdapter.new(@map, sql_test_uri)
+
+      @sequel.drop_table?(:sqltable)
+      @map.create_schema(@sequel)
+
+      @streamer = build_streamer
+    end
+
+    it 'imports from all dbs' do
+      ids = (1.upto(4)).map { BSON::ObjectId.new }
+      ids.each_with_index do |_id, i|
+        collection = mongo["test_#{i}"]['collection']
+        collection.drop
+        collection.insert({:_id => _id, :var => i}, :w => 1)
+      end
+
+      @streamer.options[:skip_tail] = true
+      @streamer.initial_import
+
+      sqlobjs = @sequel[:sqltable].select.to_a
+      assert_equal(ids.map(&:to_s).sort, sqlobjs.map { |o| o[:_id] }.sort)
+    end
+  end
+end

data/test/unit/lib/mosql/schema.rb

@@ -12,6 +12,7 @@ db:
         :source: _id
         :type: TEXT
       - var: INTEGER
+      - str: TEXT
   with_extra_props:
     :meta:
       :table: sqltable2
@@ -88,16 +89,17 @@ EOF
   it 'creates a SQL schema with the right fields' do
     db = {}
     stub_1 = stub()
-    stub_1.expects(:column).with('id', 'TEXT')
-    stub_1.expects(:column).with('var', 'INTEGER')
+    stub_1.expects(:column).with('id', 'TEXT', {})
+    stub_1.expects(:column).with('var', 'INTEGER', {})
+    stub_1.expects(:column).with('str', 'TEXT', {})
     stub_1.expects(:column).with('_extra_props').never
     stub_1.expects(:primary_key).with([:id])
     stub_2 = stub()
-    stub_2.expects(:column).with('id', 'TEXT')
+    stub_2.expects(:column).with('id', 'TEXT', {})
     stub_2.expects(:column).with('_extra_props', 'TEXT')
     stub_2.expects(:primary_key).with([:id])
     stub_3 = stub()
-    stub_3.expects(:column).with('_id', 'TEXT')
+    stub_3.expects(:column).with('_id', 'TEXT', {})
     stub_3.expects(:column).with('_extra_props').never
     stub_3.expects(:primary_key).with([:_id])
     (class << db; self; end).send(:define_method, :create_table?) do |tbl, &blk|
@@ -118,8 +120,8 @@ EOF
 
   describe 'when transforming' do
     it 'transforms rows' do
-      out = @map.transform('db.collection', {'_id' => "row 1", 'var' => 6})
-      assert_equal(["row 1", 6], out)
+      out = @map.transform('db.collection', {'_id' => "row 1", 'var' => 6, 'str' => 'a string'})
+      assert_equal(["row 1", 6, 'a string'], out)
     end
 
     it 'Includes extra props' do
@@ -130,9 +132,21 @@ EOF
     end
 
     it 'gets all_columns right' do
-      assert_equal(['id', 'var'], @map.all_columns(@map.find_ns('db.collection')))
+      assert_equal(['id', 'var', 'str'], @map.all_columns(@map.find_ns('db.collection')))
       assert_equal(['id', '_extra_props'], @map.all_columns(@map.find_ns('db.with_extra_props')))
     end
+
+    it 'stringifies symbols' do
+      out = @map.transform('db.collection', {'_id' => "row 1", 'str' => :stringy})
+      assert_equal(["row 1", nil, 'stringy'], out)
+    end
+
+    it 'changes NaN to null in extra_props' do
+      out = @map.transform('db.with_extra_props', {'_id' => 7, 'nancy' => 0.0/0.0})
+      extra = JSON.parse(out[1])
+      assert(extra.key?('nancy'))
+      assert_equal(nil, extra['nancy'])
+    end
   end
 
   describe 'when copying data' do
@@ -195,4 +209,76 @@ EOF
             {'a' => { 'c' => 4 }})
     end
   end
+
+  describe 'when handling a map with aliases' do
+  ALIAS_MAP = <<EOF
+---
+db:
+  :meta:
+    :alias: db_[0-9]+
+  collection:
+    :meta:
+      :table: sqltable
+    :columns:
+      - _id: TEXT
+      - var: INTEGER
+EOF
+    before do
+      @map = MoSQL::Schema.new(YAML.load(ALIAS_MAP))
+    end
+
+    it 'can look up collections by aliases' do
+      ns = @map.find_ns("db.collection")
+      assert_equal(ns, @map.find_ns("db_00.collection"))
+      assert_equal(ns, @map.find_ns("db_01.collection"))
+    end
+
+    it 'caches negative lookups' do
+      assert_equal(nil, @map.find_ns("nosuchdb.foo"))
+      assert(@map.instance_variable_get(:@map).key?("nosuchdb"))
+    end
+
+    it 'can do lookups after a negative cache' do
+      @map.find_ns("nosuchdb.foo")
+      assert_nil(@map.find_ns("otherdb.collection"))
+    end
+  end
+
+  describe 'parsing magic source values' do
+  OTHER_MAP = <<EOF
+---
+db:
+  collection:
+    :meta:
+      :table: a_table
+    :columns:
+      - _id: TEXT
+      - mosql_created:
+        :source: $timestamp
+        :type: timestamp
+  invalid:
+    :meta:
+      :table: invalid
+    :columns:
+      - _id: TEXT
+      - magic:
+        :source: $magic
+        :type: timestamp
+EOF
+
+    before do
+      @othermap = MoSQL::Schema.new(YAML.load(OTHER_MAP))
+    end
+
+    it 'translates $timestamp' do
+      r = @othermap.transform('db.collection', { '_id' => 'a' })
+      assert_equal(['a', Sequel.function(:now)], r)
+    end
+
+    it 'rejects unknown specials' do
+      assert_raises(MoSQL::SchemaError) do
+        r = @othermap.transform('db.invalid', { '_id' => 'a' })
+      end
+    end
+  end
 end