mongoriver 0.1.0 → 0.2.0

data/Rakefile

@@ -1,2 +1,16 @@
 #!/usr/bin/env rake
-require "bundler/gem_tasks"
+require 'bundler/setup'
+require 'bundler/gem_tasks'
+require 'rake/testtask'
+
+Rake::TestTask.new do |t|
+  t.test_files = FileList['test/test_*.rb']
+end
+
+Rake::TestTask.new(:'test-unit') do |t|
+  t.test_files = FileList['test/test_mongoriver.rb']
+end
+
+Rake::TestTask.new(:'test-connected') do |t|
+  t.test_files = FileList['test/test_*_connected.rb']
+end

data/bin/{optail → watch-oplog}

@@ -7,38 +7,25 @@ require 'bundler/setup'
 require 'mongoriver'
 
 module Mongoriver
-  class Mongocp < Streambed
+  class OplogWatcher < AbstractOutlet
     include Mongoriver::Logging
 
-    def initialize(upstreams, type, start_optime, pause)
-      super(upstreams, type)
-      @start_optime = start_optime
-      @pause = pause
+    def insert(db_name, collection_name, document)
+      log.info("got an insert for #{db_name}.#{collection_name}! #{document.inspect}")
     end
 
-    def pause
-      if @pause
-        $stderr.puts("Press enter to continue")
-        $stdin.readline
-      end
+    def remove(db_name, collection_name, document)
+      log.info("got a remove for #{db_name}.#{collection_name}! #{document.inspect}")
     end
 
-    def hook_optime
-      @start_optime
-    end
-
-    def hook_update_optime(ts, mandatory)
-    end
-
-    all_hooks.each do |name, _, opts|
-      next if name == :optime || name == :update_optime
-      define_method(hook_name(name)) {|*args| pause}
+    def update(db_name, collection_name, selector, updates)
+      log.info("got an update for #{db_name}.#{collection_name}! #{selector}, #{updates}")
     end
   end
 end
 
 def main
-  options = {:host => nil, :port => nil, :type => :slave, :optime => 0, :pause => true, :verbose => 0}
+  options = {:host => nil, :port => nil, :type => :direct, :optime => 0, :pause => true, :verbose => 0}
   optparse = OptionParser.new do |opts|
     opts.banner = "Usage: #{$0} [options]"
 
@@ -59,17 +46,9 @@ def main
       options[:port] = Integer(port)
     end
 
-    opts.on('-a', '--all', 'Allow connections even directly to a primary') do
-      options[:type] = :direct
-    end
-
     opts.on('-s OPTIME', '--start', 'Starting optime') do |optime|
       options[:optime] = Integer(optime)
     end
-
-    opts.on('-f', '--follow-automatically', "Don't prompt between ops") do
-      options[:pause] = false
-    end
   end
   optparse.parse!
 
@@ -86,8 +65,19 @@ def main
     log.level = Log4r::INFO
   end
 
-  runner = Mongoriver::Mongocp.new(["#{options[:host]}:#{options[:port]}"], options[:type], options[:optime], options[:pause])
-  runner.run
+  tailer = Mongoriver::Tailer.new(["#{options[:host]}:#{options[:port]}"], options[:type])
+  outlet = Mongoriver::OplogWatcher.new
+
+  stream = Mongoriver::Stream.new(tailer, outlet)
+
+  %w[TERM INT USR2].each do |sig|
+    Signal.trap(sig) do
+      log.info("Got SIG#{sig}. Preparing to exit...")
+      stream.stop
+    end
+  end
+
+  stream.run_forever
   return 0
 end
 

data/lib/mongoriver.rb

@@ -3,10 +3,12 @@ require 'log4r'
 
 module Mongoriver; end
 
+require 'mongoriver/version'
 require 'mongoriver/log'
+require 'mongoriver/assertions'
 
-require 'mongoriver/streambed'
 require 'mongoriver/tailer'
 require 'mongoriver/abstract_persistent_tailer'
 require 'mongoriver/persistent_tailer'
-require 'mongoriver/version'
+require 'mongoriver/abstract_outlet'
+require 'mongoriver/stream'

data/lib/mongoriver/abstract_outlet.rb

@@ -0,0 +1,20 @@
+module Mongoriver
+  class AbstractOutlet
+
+    # implement these methods in your subclass
+    def update_optime(timestamp); end
+
+    def insert(db_name, collection_name, document); end
+    def remove(db_name, collection_name, document); end
+    def update(db_name, collection_name, selector, update); end
+
+    def create_index(db_name, collection_name, index_key, options); end
+    def drop_index(db_name, collection_name, index_name); end
+
+    def create_collection(db_name, collection_name,  options); end
+    def drop_collection(db_name, collection_name); end
+    def rename_collection(db_name, old_collection_name, new_collection_name); end
+
+    def drop_database(db_name); end
+  end
+end

data/lib/mongoriver/assertions.rb

@@ -0,0 +1,9 @@
+module Mongoriver
+  module Assertions
+    class AssertionFailure < StandardError; end
+
+    def assert(condition, msg)
+      raise AssertionFailure.new(msg) unless condition
+    end
+  end
+end

data/lib/mongoriver/stream.rb

@@ -0,0 +1,169 @@
+module Mongoriver
+  class Stream
+    include Mongoriver::Logging
+    include Mongoriver::Assertions
+
+    attr_accessor :tailer, :outlet
+
+    def initialize(tailer, outlet)
+      assert(tailer.is_a?(Tailer),
+             "tailer must be a subclass/instance of Tailer")
+      assert(outlet.is_a?(AbstractOutlet),
+             "outlet must be a subclass (or instance) of AbstractOutlet")
+
+      @tailer = tailer
+      @outlet = outlet
+      @stop = false
+      @stats = {}
+    end
+
+    def stats
+      @stats
+    end
+
+    def run_forever(starting_timestamp=nil)
+      if starting_timestamp
+        @tailer.tail_from(optime_from_ts(starting_timestamp))
+      else
+        @tailer.tail_from
+      end
+
+      until @stop
+        @tailer.stream do |op|
+          handle_op(op)
+        end
+      end
+    end
+
+    def stop
+      @stop = true
+      @tailer.stop
+    end
+
+    private
+
+    def optime_from_ts(timestamp)
+      if timestamp.is_a?(Integer)
+        if timestamp >= 0
+          BSON::Timestamp.new(timestamp, 0)
+        else
+          raise "Invalid optime: #{timestamp}"
+        end
+      else
+        raise "Unrecognized type #{timestamp.class} (#{timestamp.inspect}) " \
+              "for start_timestamp"
+      end
+    end
+
+    def trigger(name, *args)
+      signature = "#{name}(" + args.map { |arg| arg.inspect }.join(', ') + ")"
+      log.debug("triggering #{signature}")
+      @stats[name] ||= 0
+      @stats[name] += 1
+
+      @outlet.send(name, *args)
+    end
+
+    def parse_ns(ns)
+      ns.split('.', 2)
+    end
+
+    def handle_op(entry)
+      op = entry['op']
+      data = entry['o']
+      ns = entry['ns']
+
+      if op == 'n'
+        # This happens for initial rs.initiate() op, maybe others.
+        log.debug("Skipping no-op #{entry.inspect}")
+        return
+      end
+
+      db_name, collection_name = parse_ns(ns)
+      assert(db_name, "nil db name #{db_name.inspect} for #{entry.inspect}")
+
+      case op
+      when 'i'
+        handle_insert(db_name, collection_name, data)
+      when 'u'
+        selector = entry['o2']
+        trigger(:update, db_name, collection_name, selector, data)
+      when 'd'
+        trigger(:remove, db_name, collection_name, data)
+      when 'c'
+        assert(collection_name == '$cmd',
+               "Command collection name is #{collection_name.inspect} for " \
+               "#{entry.inspect}, but should be '$cmd'}")
+
+        handle_cmd(db_name, collection_name, data)
+      else
+        raise "Unrecognized op: #{op} (#{entry.inspect})"
+      end
+
+      optime = entry['ts']
+      trigger(:update_optime, optime.seconds)
+    end
+
+    def handle_insert(db_name, collection_name, data)
+      if collection_name == 'system.indexes'
+        handle_create_index(data)
+      else
+        trigger(:insert, db_name, collection_name, data)
+      end
+    end
+
+    def handle_create_index(spec)
+      db_name, collection_name = parse_ns(spec['ns'])
+      index_key = spec['key'].map { |field, dir| [field, dir.round] }
+      options = {}
+
+      spec.each do |key, value|
+        case key
+        when 'v'
+          unless value == 1
+            raise NotImplementedError.new("Only v=1 indexes are supported, " \
+                                          "not v=#{value.inspect}")
+          end
+        when 'ns', 'key', '_id' # do nothing
+        else
+          options[key.to_sym] = value
+        end
+      end
+
+      assert(options.include?(:name),
+             "No name defined for index spec #{spec.inspect}")
+
+      trigger(:create_index, db_name, collection_name, index_key, options)
+    end
+
+    def handle_cmd(db_name, collection_name, data)
+      if deleted_from_collection = data['deleteIndexes']
+        index_name = data['index']
+        trigger(:drop_index, db_name, deleted_from_collection, index_name)
+      elsif created_collection = data['create']
+        handle_create_collection(db_name, data)
+      elsif dropped_collection = data['drop']
+        trigger(:drop_collection, db_name, dropped_collection)
+      elsif old_collection_ns = data['renameCollection']
+        db_name, old_collection_name = parse_ns(old_collection_ns)
+        _, new_collection_name = parse_ns(data['to'])
+        trigger(:rename_collection, db_name, old_collection_name, new_collection_name)
+      elsif data['dropDatabase'] == 1
+        trigger(:drop_database, db_name)
+      else
+        raise "Unrecognized command #{data.inspect}"
+      end
+    end
+
+    def handle_create_collection(db_name, data)
+      collection_name = data.delete('create')
+
+      options = {}
+      data.each do |k, v|
+        options[k.to_sym] = (k == 'size') ? v.round : v
+      end
+
+      trigger(:create_collection, db_name, collection_name, options)
+    end
+  end
+end

data/lib/mongoriver/tailer.rb

@@ -3,14 +3,17 @@ module Mongoriver
     include Mongoriver::Logging
 
     attr_reader :upstream_conn
+    attr_reader :oplog
 
-    def initialize(upstreams, type)
+    def initialize(upstreams, type, oplog = "oplog.rs")
       @upstreams = upstreams
       @type = type
+      @oplog = oplog
       # This number seems high
       @conn_opts = {:op_timeout => 86400}
 
       @cursor = nil
+      @stop = false
 
       connect_upstream
     end
@@ -29,7 +32,7 @@ module Mongoriver
         opts = @conn_opts.merge(:slave_ok => true)
         host, port = parse_direct_upstream
         @upstream_conn = Mongo::Connection.new(host, port, opts)
-        raise "Server at #{@upstream_conn.host}:#{@upstream_conn.port} is the primary -- if you're ok with that, check why your wrapper is passing :direct rather than :slave (HINT: try passing a -a to scripts like optail or mongocp)" if @type == :slave && @upstream_conn.primary?
+        raise "Server at #{@upstream_conn.host}:#{@upstream_conn.port} is the primary -- if you're ok with that, check why your wrapper is passing :direct rather than :slave" if @type == :slave && @upstream_conn.primary?
         ensure_upstream_replset!
       when :existing
         raise "Must pass in a single existing Mongo::Connection with :existing" unless @upstreams.length == 1 && @upstreams[0].respond_to?(:db)
@@ -61,7 +64,7 @@ module Mongoriver
     end
 
     def oplog_collection
-      @upstream_conn.db('local').collection('oplog.rs')
+      @upstream_conn.db('local').collection(oplog)
     end
 
     def tail_from(ts, opts = {})
@@ -81,20 +84,26 @@ module Mongoriver
       end
     end
 
-    def stop
-      @cursor.close if @cursor
-      @cursor = nil
-    end
-
     def stream(limit=nil)
       count = 0
-      while @cursor.has_next?
+      while !@stop && @cursor.has_next?
         count += 1
         break if limit && count >= limit
+
         yield @cursor.next
       end
 
       return @cursor.has_next?
     end
+
+    def stop
+      @stop = true
+    end
+
+    def close
+      @cursor.close if @cursor
+      @cursor = nil
+      @stop = false
+    end
   end
 end

data/lib/mongoriver/version.rb

@@ -1,3 +1,3 @@
 module Mongoriver
-  VERSION = "0.1.0"
+  VERSION = "0.2.0"
 end

data/mongoriver.gemspec

@@ -19,4 +19,8 @@ Gem::Specification.new do |gem|
   gem.add_runtime_dependency('mongo', '>= 1.7')
   gem.add_runtime_dependency('bson_ext')
   gem.add_runtime_dependency('log4r')
+
+  gem.add_development_dependency('rake')
+  gem.add_development_dependency('minitest')
+  gem.add_development_dependency('mocha', '>= 0.13')
 end

data/test/test_mongoriver.rb

@@ -0,0 +1,65 @@
+require 'mongoriver'
+require 'mongo'
+require 'minitest/autorun'
+require 'mocha/setup'
+
+describe 'Mongoriver::Stream' do
+  def create_op(op)
+    ts = Time.now.to_i
+    {'ts'=>BSON::Timestamp.new(ts, 0), 'h'=>1234, 'v'=>1, 'ns'=>'foo.bar'}.merge(op)
+  end
+
+  before do
+    conn = stub(:db => nil)
+    @tailer = Mongoriver::Tailer.new([conn], :existing)
+    @outlet = Mongoriver::AbstractOutlet.new
+    @stream = Mongoriver::Stream.new(@tailer, @outlet)
+
+    @outlet.expects(:update_optime).at_least_once
+  end
+
+  it 'triggers insert' do
+    @outlet.expects(:insert).once.with('foo', 'bar', {'_id' => 'baz'})
+    @stream.send(:handle_op, create_op({'op'=>'i', 'o'=>{'_id'=>'baz'}}))
+  end
+
+  it 'triggers update' do
+    @outlet.expects(:update).once.with('foo', 'bar', {'_id' => 'baz'}, {'a' => 'b'})
+    @stream.send(:handle_op, create_op({'op'=>'u', 'o2'=>{'_id'=>'baz'}, 'o'=>{'a'=>'b'}}))
+  end
+
+  it 'triggers remove' do
+    @outlet.expects(:remove).once.with('foo', 'bar', {'_id' => 'baz'})
+    @stream.send(:handle_op, create_op({'op'=>'d', 'b'=>true, 'o'=>{'_id'=>'baz'}}))
+  end
+
+  it 'triggers create_collection' do
+    @outlet.expects(:create_collection).once.with('foo', 'bar', {:capped => true, :size => 10})
+    @stream.send(:handle_op, create_op({'op'=>'c', 'ns'=>'foo.$cmd', 'o'=>{'create'=>'bar', 'capped'=>true, 'size'=>10.0}}))
+  end
+
+  it 'triggers drop_collection' do
+    @outlet.expects(:drop_collection).once.with('foo', 'bar')
+    @stream.send(:handle_op, create_op({'op'=>'c', 'ns'=>'foo.$cmd', 'o'=>{'drop'=>'bar'}}))
+  end
+
+  it 'triggers rename_collection' do
+    @outlet.expects(:rename_collection).once.with('foo', 'bar', 'bar_2')
+    @stream.send(:handle_op, create_op({'op'=>'c', 'ns'=>'admin.$cmd', 'o'=>{'renameCollection'=>'foo.bar', 'to'=>'foo.bar_2'}}))
+  end
+
+  it 'triggers create_index' do
+    @outlet.expects(:create_index).once.with('foo', 'bar', [['baz', 1]], {:name => 'baz_1'})
+    @stream.send(:handle_op, create_op({'op'=>'i', 'ns'=>'foo.system.indexes', 'o'=>{'_id'=>'index_id', 'ns'=>'foo.bar', 'key'=>{'baz'=>1.0}, 'name'=>'baz_1'}}))
+  end
+
+  it 'triggers drop_index' do
+    @outlet.expects(:drop_index).once.with('foo', 'bar', 'baz_1')
+    @stream.send(:handle_op, create_op({'op'=>'c', 'ns'=>'foo.$cmd', 'o'=>{'deleteIndexes'=>'bar', 'index'=>'baz_1'}}))
+  end
+
+  it 'triggers drop_database' do
+    @outlet.expects(:drop_database).once.with('foo')
+    @stream.send(:handle_op, create_op({'op'=>'c', 'ns'=>'foo.$cmd', 'o'=>{'dropDatabase'=>1.0}}))
+  end
+end

data/test/test_mongoriver_connected.rb

@@ -0,0 +1,83 @@
+require 'mongoriver'
+require 'mongo'
+require 'minitest/autorun'
+require 'mocha/setup'
+
+# Connected tests: run these with eg MONGO_SERVER=localhost:27017
+
+MONGO_SERVER = ENV['MONGO_SERVER'] || 'localhost:27017'
+
+def connect
+  begin
+    host, port = MONGO_SERVER.split(':', 2)
+    Mongo::Connection.new(host, port)
+  rescue Mongo::ConnectionFailure
+    nil
+  end
+end
+
+describe 'connected tests' do
+  before do
+    @mongo = connect
+    skip unless @mongo
+  end
+
+  describe 'Mongoriver::Stream' do
+    before do
+      @tailer = Mongoriver::Tailer.new([MONGO_SERVER], :direct)
+      @outlet = Mongoriver::AbstractOutlet.new
+
+      @stream = Mongoriver::Stream.new(@tailer, @outlet)
+
+      @tail_from = @tailer.most_recent_timestamp.seconds + 1
+      sleep(1)
+    end
+
+    it 'triggers the correct ops in the correct order' do
+      db = 'test'
+      collection = 'test'
+      doc = {'_id' => 'foo', 'bar' => 'baz'}
+      updated_doc = doc.dup.merge('bar' => 'qux')
+      index_keys = [['bar', 1]]
+
+      @outlet.expects(:update_optime).at_least_once
+
+      op_sequence = sequence('op_sequence')
+
+      @outlet.expects(:insert).once.with(db, collection, doc).in_sequence(op_sequence)
+      @outlet.expects(:update).once.with(db, collection, {'_id' => 'foo'}, updated_doc).in_sequence(op_sequence)
+      @outlet.expects(:remove).once.with(db, collection, {'_id' => 'foo'}).in_sequence(op_sequence)
+
+      @outlet.expects(:create_index).once.with(db, collection, index_keys, {:name => 'bar_1'}).in_sequence(op_sequence)
+      @outlet.expects(:drop_index).once.with(db, collection, 'bar_1').in_sequence(op_sequence)
+
+      @outlet.expects(:rename_collection).once.with(db, collection, collection+'_foo').in_sequence(op_sequence)
+      @outlet.expects(:drop_collection).once.with(db, collection+'_foo').in_sequence(op_sequence)
+      @outlet.expects(:drop_database).once.with(db) { @stream.stop }.in_sequence(op_sequence)
+
+      coll = @mongo[db][collection]
+      coll.insert(doc)
+      coll.update({'_id' => 'foo'}, doc.merge('bar' => 'qux'))
+      coll.remove({'_id' => 'foo'})
+
+      name = coll.ensure_index(index_keys)
+      coll.drop_index(name)
+
+      @mongo[db].rename_collection(collection, collection+'_foo')
+      @mongo[db].drop_collection(collection+'_foo')
+      @mongo.drop_database(db)
+
+      @stream.run_forever(@tail_from)
+    end
+
+    it 'passes options to create_collection' do
+      @outlet.expects(:create_collection).once.with('test', 'test', {:capped => true, :size => 10}) { @stream.stop }
+      @outlet.expects(:update_optime).at_least_once.with(anything) { @stream.stop }
+
+      @mongo['test'].create_collection('test', :capped => true, :size => 10)
+      @mongo.drop_database('test')
+
+      @stream.run_forever(@tail_from)
+    end
+  end
+end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: mongoriver
 version: !ruby/object:Gem::Version
-  version: 0.1.0
+  version: 0.2.0
   prerelease: 
 platform: ruby
 authors:
@@ -9,7 +9,7 @@ authors:
 autorequire: 
 bindir: bin
 cert_chain: []
-date: 2013-02-05 00:00:00.000000000 Z
+date: 2013-05-07 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: mongo
@@ -59,12 +59,59 @@ dependencies:
     - - ! '>='
       - !ruby/object:Gem::Version
         version: '0'
+- !ruby/object:Gem::Dependency
+  name: rake
+  requirement: !ruby/object:Gem::Requirement
+    none: false
+    requirements:
+    - - ! '>='
+      - !ruby/object:Gem::Version
+        version: '0'
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    none: false
+    requirements:
+    - - ! '>='
+      - !ruby/object:Gem::Version
+        version: '0'
+- !ruby/object:Gem::Dependency
+  name: minitest
+  requirement: !ruby/object:Gem::Requirement
+    none: false
+    requirements:
+    - - ! '>='
+      - !ruby/object:Gem::Version
+        version: '0'
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    none: false
+    requirements:
+    - - ! '>='
+      - !ruby/object:Gem::Version
+        version: '0'
+- !ruby/object:Gem::Dependency
+  name: mocha
+  requirement: !ruby/object:Gem::Requirement
+    none: false
+    requirements:
+    - - ! '>='
+      - !ruby/object:Gem::Version
+        version: '0.13'
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    none: false
+    requirements:
+    - - ! '>='
+      - !ruby/object:Gem::Version
+        version: '0.13'
 description: Some tools and libraries to simplify tailing the mongod oplog
 email:
 - gdb@gregbrockman.com
 executables:
-- mongocp
-- optail
+- watch-oplog
 extensions: []
 extra_rdoc_files: []
 files:
@@ -73,16 +120,19 @@ files:
 - LICENSE
 - README.md
 - Rakefile
-- bin/mongocp
-- bin/optail
+- bin/watch-oplog
 - lib/mongoriver.rb
+- lib/mongoriver/abstract_outlet.rb
 - lib/mongoriver/abstract_persistent_tailer.rb
+- lib/mongoriver/assertions.rb
 - lib/mongoriver/log.rb
 - lib/mongoriver/persistent_tailer.rb
-- lib/mongoriver/streambed.rb
+- lib/mongoriver/stream.rb
 - lib/mongoriver/tailer.rb
 - lib/mongoriver/version.rb
 - mongoriver.gemspec
+- test/test_mongoriver.rb
+- test/test_mongoriver_connected.rb
 homepage: ''
 licenses: []
 post_install_message: 
@@ -107,4 +157,7 @@ rubygems_version: 1.8.23
 signing_key: 
 specification_version: 3
 summary: monogdb oplog-tailing utilities.
-test_files: []
+test_files:
+- test/test_mongoriver.rb
+- test/test_mongoriver_connected.rb
+has_rdoc: 

data/bin/mongocp

@@ -1,250 +0,0 @@
-#!/usr/bin/env ruby
-require 'logger'
-require 'optparse'
-
-require 'rubygems'
-require 'bundler/setup'
-require 'mongoriver'
-
-module Mongoriver
-  class Mongocp < Streambed
-    include Mongoriver::Logging
-
-    def initialize(upstreams, type, downstream, prefix)
-      super(upstreams, type)
-      @downstream = downstream
-      @prefix = prefix
-      connect_downstream
-    end
-
-    def hook_optime
-      if optime = optime_collection.find_one(:_id => @prefix)
-        optime['ts']
-      else
-        nil
-      end
-    end
-
-    def hook_update_optime(ts, mandatory)
-      optime_collection.update({:_id => @prefix}, {'$set' => {:ts => ts}}, :upsert => true) if mandatory || rand(20) == 0
-    end
-
-    def hook_initial_sync_index(db_name, collection_name, index_key, options)
-      collection = downstream_collection(db_name, collection_name)
-      index_hash = BSON::OrderedHash.new
-      index_key.each {|k,v| index_hash[k] = v}
-      collection.send(:generate_indexes, index_hash, nil, options)
-    end
-
-    def hook_initial_sync_record_batch(db_name, collection_name, records)
-      collection = downstream_collection(db_name, collection_name)
-      bulk_insert(collection, records)
-    end
-
-    # TODO: should probably do the same key checking nonsense as the above
-    def hook_stream_insert(db_name, collection_name, object)
-      collection = downstream_collection(db_name, collection_name)
-      wrap_errors(collection, object['_id']) do
-        # Only needed if safe mode is set in the driver. Note that the
-        # argument here for oplog idempotency in the case of unique
-        # keys is kind of interesting. I believe I can prove
-        # idempotency as long as Mongo has no insert order-dependent
-        # unique indexes (which I believe is true) and that you do all
-        # your object updates as upserts.
-        allow_dupkeys do
-          collection.insert(object)
-        end
-      end
-    end
-
-    def hook_stream_update(db_name, collection_name, selector, update)
-      collection = downstream_collection(db_name, collection_name)
-      wrap_errors(collection, selector['_id']) do
-        collection.update(selector, update, :upsert => true)
-      end
-    end
-
-    def hook_stream_remove(db_name, collection_name, object)
-      collection = downstream_collection(db_name, collection_name)
-      wrap_errors(collection, object['_id']) do
-        collection.remove(object)
-      end
-    end
-
-    def hook_stream_create_collection(db_name, create)
-      db = downstream_db(db_name)
-      wrap_errors(db, create) do
-        db.create_collection(create)
-      end
-    end
-
-    # "Error renaming collection: #<BSON::OrderedHash:0x83869e34 {\"errmsg\"=>\"exception: source namespace does not exist\", \"code\"=>10026, \"ok\"=>0.0}>"
-    #
-    # Possibly need the same thing if the destination already exists
-    def hook_stream_rename_collection(db_name, source, target)
-      db = downstream_db(db_name)
-      wrap_errors(db, "#{source} -> #{target}") do
-        begin
-          db.rename_collection(source, target)
-        rescue Mongo::MongoDBError => e
-          if e.message =~ /Error renaming collection: .*exception: source namespace does not exist"/
-            log.warn("Ignoring rename of non-existent collection #{source} -> #{target}: #{e} (expected when replaying part of the oplog)")
-          elsif e.message =~ /Error renaming collection: .*exception: target namespace exists"/
-            log.warn("Ignoring rename of #{source} to existing collection #{target}: #{e} (expected when replaying part of the oplog)")
-          else
-            raise
-          end
-        end
-      end
-    end
-
-    def hook_stream_drop_index(db_name, collection_name, index_name)
-      collection = downstream_collection(db_name, collection_name)
-      wrap_errors(collection, index_name) do
-        begin
-          collection.drop_index(index_name)
-        rescue Mongo::MongoDBError => e
-          raise
-          if e.message =~ /index not found/
-            log.warn("Ignoring drop of non-existent index #{index_name.inspect}: #{e} (expected when replaying part of the oplog)")
-          else
-            raise
-          end
-        end
-      end
-    end
-
-    def hook_stream_drop_collection(db_name, dropped)
-      db = downstream_db(db_name)
-      wrap_errors(db, dropped) do
-        db.drop_collection(dropped)
-      end
-    end
-
-    def hook_stream_drop_database(db_name)
-      db = downstream_db(db_name)
-      wrap_errors(db, db_name) do
-        db.command(:dropDatabase => 1)
-      end
-    end
-
-    private
-
-    def allow_dupkeys(&blk)
-      begin
-        blk.call
-      rescue Mongo::OperationFailure => e
-        if e.error_code == 11000
-          log.warn("Ignoring unique index violation: #{e} (expected when replaying part of the oplog)")
-        else
-          raise
-        end
-      end
-    end
-
-    def bulk_insert(collection, docs)
-      begin
-        # Use the internal insert_documents method because it lets us
-        # disable key verification
-        collection.send(:insert_documents, docs, collection.name, false)
-      rescue Mongo::MongoRubyError => e
-        log.error("#{ns}: Caught error on batch insert", e)
-        docs.each do |doc|
-          wrap_errors(collection, doc['_id']) do
-            collection.send(:insert_documents, [doc], collection.name, false)
-          end
-        end
-      end
-    end
-
-    def wrap_errors(collection_or_db, object, &blk)
-      begin
-        blk.call
-      rescue Mongo::MongoRubyError => e
-        if collecton_or_db.kind_of?(Mongo::Collection)
-          ns = "#{collection_or_db.db.name}.#{collection_or_db.name}"
-        else
-          ns = collection_or_db.db.name
-        end
-        log.error("#{ns}: Unknown error for #{object}", e)
-      end
-    end
-
-    def downstream_db(db_name)
-      prefixed = "#{@prefix}_#{db_name}"
-      @downstream_conn.db(prefixed)
-    end
-
-    def downstream_collection(db_name, collection_name)
-      downstream_db(db_name).collection(collection_name)
-    end
-
-    def optime_collection
-      @optime_collection ||= @downstream_conn.db('_mongocp').collection('optime')
-    end
-
-    def connect_downstream
-      host, port = @tailer.parse_host_spec(@downstream)
-      @downstream_conn = Mongo::Connection.new(host, port, :safe => true)
-    end
-  end
-end
-
-def main
-  options = {:host => nil, :port => nil, :type => :slave, :verbose => 0}
-  optparse = OptionParser.new do |opts|
-    opts.banner = "Usage: #{$0} [options]"
-
-    opts.on('-v', '--verbosity', 'Verbosity of debugging output') do
-      options[:verbose] += 1
-    end
-
-    opts.on('-h', '--help', 'Display this message') do
-      puts opts
-      exit(1)
-    end
-
-    opts.on('--help', 'Display this message') do
-      puts opts
-      exit(1)
-    end
-
-    opts.on('-h HOST', '--host', 'Upstream host to connect to') do |host|
-      options[:host] = host
-    end
-
-    opts.on('-p PORT', '--port', 'Upstream host to connect to') do |port|
-      options[:port] = Integer(port)
-    end
-
-    opts.on('-a', '--all', 'Allow connections even directly to a primary') do
-      options[:type] = :direct
-    end
-  end
-  optparse.parse!
-
-  if ARGV.length != 0
-    puts optparse
-    return 1
-  end
-
-  log = Log4r::Logger.new('Stripe')
-  log.outputters = Log4r::StdoutOutputter.new(STDERR)
-  if options[:verbose] >= 1
-    log.level = Log4r::DEBUG
-  else
-    log.level = Log4r::INFO
-  end
-  runner = Mongoriver::Mongocp.new(["#{options[:host]}:#{options[:port]}"], options[:type], 'localhost:5001', 'test')
-  runner.run
-  return 0
-end
-
-if $0 == __FILE__
-  ret = main
-  begin
-    exit(ret)
-  rescue TypeError
-    exit(0)
-  end
-end

data/lib/mongoriver/streambed.rb

@@ -1,299 +0,0 @@
-module Mongoriver
-  class Streambed
-    include Mongoriver::Logging
-
-    attr_reader :stats
-
-    class AssertionFailure < StandardError; end
-
-    def assert(condition, msg)
-      raise AssertionFailure.new(msg) unless condition
-    end
-
-    def initialize(upstreams, type)
-      @tailer = Mongoriver::Tailer.new(upstreams, type)
-      @record_fetch_batch_size = 1024
-      @record_sync_batch_size = 256
-      @stats = Hash.new(0)
-    end
-
-    def run
-      self.class.validate_hooks!
-
-      unless ts = starting_optime
-        ts = @tailer.most_recent_timestamp
-        initial_sync
-        hook_update_optime(ts, true)
-      end
-
-      tail_from(ts)
-    end
-
-    def self.my_hooks
-      @hooks ||= []
-    end
-
-    def self.all_hooks
-      hooks = my_hooks
-      if superclass <= Streambed
-        hooks + superclass.all_hooks
-      else
-        hooks
-      end
-    end
-
-    def self.validate_hooks!
-      errors = []
-      all_hooks.each do |name, args, opts|
-        method = self.instance_method(hook_name(name))
-        signature = "#{method.name}(#{args.join(', ')})"
-        if method.owner == Streambed && !opts[:default]
-          errors << "Must provide implementation of #{signature}"
-        end
-      end
-
-      raise "You need to fix the following hook errors:
-
-  #{errors.join("\n  ")}" if errors.length > 0
-    end
-
-    def self.hook_name(name)
-      "hook_#{name}"
-    end
-
-    def self.hook(name, args=[], opts={})
-      if default = opts[:default]
-        target = hook_name(default)
-        implementation = Proc.new do |*args, &blk|
-          send(target, *args, &blk)
-        end
-      else
-        implementation = Proc.new do
-          raise NotImplementedError.new("Override in subclass")
-        end
-      end
-
-      define_method(hook_name(name), implementation)
-      my_hooks << [name, args, opts]
-    end
-
-    hook :optime
-    hook :update_optime, [:ts, :mandatory]
-    hook :initial_sync_index, [:db_name, :collection_name, :index_key, :options]
-    hook :initial_sync_record_batch, [:db_name, :collection_name, :records]
-    hook :stream_insert, [:db_name, :collection_name, :object]
-    hook :stream_update, [:db_name, :collection_name, :selector, :update]
-    hook :stream_remove, [:db_name, :collection_name, :object]
-    # Not usually a difference between the initial index creation and
-    # creating it while streaming ops.
-    hook :stream_create_index, [:db_name, :collection_name, :index_key, :options], :default => :initial_sync_index
-    # This seems to be called while doing a mapreduce.
-    hook :stream_create_collection, [:db_name, :create]
-    # This also seems to be called while doing a mapreduce. Note that
-    # I think mongo has a concept of temporary table, which I should
-    # look into, and renameCollection has some temporary table option.
-    hook :stream_rename_collection, [:db_name, :source, :target]
-    hook :stream_drop_index, [:db_name, :collection_name, :index_name]
-    hook :stream_drop_collection, [:db_name, :dropped]
-    hook :stream_drop_database, [:db_name]
-
-    private
-
-    def starting_optime
-      case time = hook_optime
-      when Integer
-        if time >= 0
-          BSON::Timestamp.new(time, 0)
-        elsif time == -1
-          @tailer.most_recent_timestamp
-        else
-          raise "Invalid optime: #{time}"
-        end
-      when BSON::Timestamp, nil
-        time
-      else
-        raise "Unrecognized type #{time.class} (#{time.inspect}) for start time"
-      end
-    end
-
-    def initial_sync
-      initial_sync_all_indexes
-      initial_sync_all_records
-    end
-
-    def initial_sync_all_indexes
-      log.info("Beginning initial sync of indexes")
-      syncable_databases.each {|db| initial_sync_indexes_for_db(db)}
-      log.info("Done initial sync of indexes")
-    end
-
-    def initial_sync_indexes_for_db(db)
-      db.collection('system.indexes').find.each do |index|
-        options = extract_options_from_index_spec(index)
-        index_key = index['key'].to_a
-
-        ns = index['ns']
-        db_name, collection_name = parse_ns(ns)
-        assert(db_name == db.name, "Index db name #{db_name.inspect} differs from current db name #{db.name.inspect}")
-
-        log.info("#{ns}: Initial sync of index #{options[:name]}")
-        hook_initial_sync_index(db_name, collection_name, index_key, options)
-      end
-    end
-
-    def initial_sync_all_records
-      log.info("Beginning initial sync of records")
-      syncable_databases.each {|db| initial_sync_records_for_db(db)}
-      log.info("Done initial sync of records")
-    end
-
-    def initial_sync_records_for_db(db)
-      syncable_collections(db).each do |collection|
-        initial_sync_records_for_collection(collection)
-      end
-    end
-
-    def initial_sync_records_for_collection(collection)
-      db_name = collection.db.name
-      collection_name = collection.name
-      ns = "#{db_name}.#{collection_name}"
-
-      log.info("#{ns}: Starting record initial sync")
-
-      records = []
-      collection.find({}, :batch_size => @record_fetch_batch_size, :timeout => false, :sort => [['$natural', 1]]) do |cursor|
-        while cursor.has_next?
-          records << cursor.next
-          if records.length > @record_sync_batch_size
-            # TODO: add better logging than this
-            log.info("#{ns}: Running sync of batch of #{records.length} records")
-            hook_initial_sync_record_batch(db_name, collection_name, records)
-            records = []
-          end
-        end
-      end
-      log.info("#{ns}: Finishing sync with a batch of #{records.length} records")
-      hook_initial_sync_record_batch(db_name, collection_name, records)
-
-      log.info("#{ns}: Finished record initial sync")
-    end
-
-    # This should be fine to instantiate all at once, since
-    # database_names returns all the dbs as strings anyway
-    def syncable_databases
-      @tailer.upstream_conn.database_names.map do |db_name|
-        next if db_name == 'local'
-        @tailer.upstream_conn.db(db_name)
-      end.compact
-    end
-
-    def syncable_collections(db)
-      db.collection_names.map do |collection_name|
-        next if collection_name.start_with?('system.')
-        db.collection(collection_name)
-      end.compact
-    end
-
-    def extract_options_from_index_spec(index)
-      options = {}
-      index.each do |key, value|
-        case key
-        when 'v'
-          raise NotImplementedError.new("Only v=1 indexes are supported at the moment, not v=#{value.inspect}") unless value == 1
-        when 'ns', 'key'
-        else
-          options[key.to_sym] = value
-        end
-      end
-
-      assert(options.include?(:name), "No name defined for index spec #{index.inspect}")
-      options
-    end
-
-    def stream_op(entry)
-      op = entry['op']
-      data = entry['o']
-      ns = entry['ns']
-
-      if op == 'n'
-        # This happens for initial rs.initiate() op, maybe others.
-        log.info("Skipping no-op #{entry.inspect}")
-        return
-      end
-
-      db_name, collection_name = parse_ns(ns)
-      assert(db_name, "Nil db name #{db_name.inspect} for #{entry.inspect}")
-
-      case op
-      when 'i'
-        if collection_name == 'system.indexes'
-          record(ns, entry, :create_index)
-          index_db_name, index_collection_name = parse_ns(data['ns'])
-          index_key = data['key'].to_a
-          options = extract_options_from_index_spec(data)
-          hook_stream_create_index(index_db_name, index_collection_name, index_key, options)
-        else
-          record(ns, entry, :insert)
-          hook_stream_insert(db_name, collection_name, data)
-        end
-      when 'u'
-        record(ns, entry, :update)
-        hook_stream_update(db_name, collection_name, entry['o2'], data)
-      when 'd'
-        record(ns, entry, :remove)
-        hook_stream_remove(db_name, collection_name, data)
-      when 'c'
-        assert(collection_name == '$cmd', "Command collection name is #{collection_name.inspect} for #{entry.inspect}")
-        if deleted_from = data['deleteIndexes']
-          record(ns, entry, :drop_index)
-          index = data['index']
-          hook_stream_drop_index(db_name, deleted_from, index)
-        elsif dropped = data['drop']
-          record(ns, entry, :drop_collection)
-          hook_stream_drop_collection(db_name, dropped)
-        elsif dropped = data['dropDatabase']
-          record(ns, entry, :drop_database)
-          hook_stream_drop_database(db_name)
-        elsif source = data['renameCollection']
-          record(ns, entry, :rename_collection)
-          target = data['to']
-          hook_stream_rename_collection(db_name, source, target)
-        elsif create = data['create']
-          record(ns, entry, :create)
-          hook_stream_create_collection(db_name, create)
-        else
-          raise "Unrecognized command #{data.inspect}"
-        end
-      else
-        raise "Unrecognized op: #{op} (#{entry.inspect})"
-      end
-
-      optime = entry['ts']
-      hook_update_optime(optime, false)
-    end
-
-    def tail_from(ts)
-      begin
-        @tailer.tail_from(ts)
-        loop do
-          @tailer.stream do |op|
-            stream_op(op)
-          end
-        end
-      ensure
-        @tailer.stop
-      end
-    end
-
-    def record(ns, entry, type)
-      stats[type] += 1
-      log.debug("#{ns}: #{type.inspect} #{entry.inspect}")
-    end
-
-    protected
-
-    def parse_ns(ns)
-      ns.split('.', 2)
-    end
-  end
-end