lib/mosql/cli.rb in mosql-0.1.0 vs lib/mosql/cli.rb in mosql-0.1.1

@@ -84,10 +84,14 @@
         end
 
         opts.on("--reimport", "Force a data re-import") do
           @options[:reimport] = true
         end
+
+        opts.on("--no-drop-tables", "Don't drop the table if it exists during the initial import") do
+          @options[:no_drop_tables] = true
+        end
       end
 
       optparse.parse!(@args)
 
       log = Log4r::Logger.new('Stripe')
@@ -181,11 +185,11 @@
       yield
       Time.now - start
     end
 
     def initial_import
-      @schemamap.create_schema(@sql.db, true)
+      @schemamap.create_schema(@sql.db, !options[:no_drop_tables])
 
       start_ts = @mongo['local']['oplog.rs'].find_one({}, {:sort => [['$natural', -1]]})['ts']
 
       want_dbs = @schemamap.all_mongo_dbs & @mongo.database_names
       want_dbs.each do |dbname|
@@ -205,11 +209,11 @@
     def import_collection(ns, collection)
       log.info("Importing for #{ns}...")
       count = 0
       batch = []
       table = @sql.table_for_ns(ns)
-      table.truncate
+      table.truncate unless options[:no_drop_tables]
 
       start    = Time.now
       sql_time = 0
       collection.find(nil, :batch_size => BATCH) do |cursor|
         with_retries do
@@ -286,17 +290,23 @@
         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=#{update['_id']}")
+          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)
           @sql.upsert_ns(ns, update)
         end
       when 'd'
         if options[:ignore_delete]
           log.debug("Ignoring delete op on #{ns} as instructed.")
         else
-          @sql.table_for_ns(ns).where(:_id => op['o']['_id']).delete
+          @sql.delete_ns(ns, op['o'])
         end
       else
         log.info("Skipping unknown op #{op.inspect}")
       end
     end