Merge pull request #5 from pmk1c/foreign-key-constraints

Make this usable with foreign key constraints

lib/csv_db.rb

@@ -16,7 +16,7 @@ def self.extension
   end
 
   class Load < SerializationHelper::Load
-    def self.load_documents(io, truncate = true)
+    def self.load_documents(io)
       tables = {}
       curr_table = nil
       io.each do |line|
@@ -34,7 +34,7 @@ def self.load_documents(io, truncate = true)
       end
 
       tables.each_pair do |table_name, contents|
-        load_table(table_name, contents, truncate)
+        load_table(table_name, contents)
       end
     end
   end
@@ -72,7 +72,7 @@ def self.dump_table_records(io, table)
         end
       end
     end
-    
+
   end
 
 

lib/json_db.rb

@@ -62,11 +62,11 @@ def self.dump_table_records(io, table)
   end
 
   class Load < SerializationHelper::Load
-    def self.load_documents(io, truncate = true)
+    def self.load_documents(io)
         JSON.load(io).tap do |json|
           json.keys.each do |table_name|
             next if json[table_name].nil?
-            load_table(table_name, json[table_name], truncate)
+            load_table(table_name, json[table_name])
           end
         end
     end

lib/serialization_helper.rb

@@ -62,28 +62,39 @@ def reenable_logger
       ActiveRecord::Base.logger = @@old_logger
     end
   end
-  
+
   class Load
     def self.load(io, truncate = true)
       ActiveRecord::Base.connection.transaction do
-        load_documents(io, truncate)
+        defer_fk_constraints do
+          truncate_all if truncate
+          load_documents(io)
+        end
       end
     end
 
-    def self.truncate_table(table)
-      begin
-        ActiveRecord::Base.connection.execute("TRUNCATE #{SerializationHelper::Utils.quote_table(table)}")
-      rescue Exception
-        ActiveRecord::Base.connection.execute("DELETE FROM #{SerializationHelper::Utils.quote_table(table)}")
+    def self.truncate_all
+      quoted_tables = tables.map do |table|
+        SerializationHelper::Utils.quote_table(table)
+      end
+      case database_type
+      when :postgresql
+        ActiveRecord::Base.connection.execute("TRUNCATE #{quoted_tables.join(',')} CASCADE")
+      when :mysql
+        quoted_tables.each do |quoted_table|
+          ActiveRecord::Base.connection.execute("TRUNCATE #{quoted_table}")
+        end
       end
     end
 
-    def self.load_table(table, data, truncate = true)
+    def self.tables
+      ActiveRecord::Base.connection.tables
+    end
+
+
+    def self.load_table(table, data)
       return if table == 'ar_internal_metadata'
       column_names = data['columns']
-      if truncate
-        truncate_table(table)
-      end
       load_records(table, column_names, data['records'])
       reset_pk_sequence!(table)
     end
@@ -95,7 +106,7 @@ def self.load_records(table, column_names, records, records_per_page=1000)
       columns = column_names.map{|cn| ActiveRecord::Base.connection.columns(table).detect{|c| c.name == cn}}
       quoted_column_names = column_names.map { |column| ActiveRecord::Base.connection.quote_column_name(column) }.join(',')
       quoted_table_name = SerializationHelper::Utils.quote_table(table)
-      
+
       0.step(records.count-1, records_per_page) do |offset|
         all_quoted_values = records[offset, records_per_page].map do |record|
           '(' + record.zip(columns).map{|c| ActiveRecord::Base.connection.quote(c.first, c.last)}.join(',') + ')'
@@ -108,9 +119,39 @@ def self.reset_pk_sequence!(table_name)
       if ActiveRecord::Base.connection.respond_to?(:reset_pk_sequence!)
         ActiveRecord::Base.connection.reset_pk_sequence!(table_name)
       end
-    end    
+    end
 
-
+    def self.defer_fk_constraints(&block)
+      case database_type
+      when :postgresql
+        # make all fk constraints deferrable
+        fk_constraints = []
+        tables.each do |table|
+          fk_constraints_on_table = ActiveRecord::Base.connection.foreign_keys(table)
+          fk_constraints_on_table.each do |fk_constraint|
+            quoted_table_name = SerializationHelper::Utils.quote_table(table)
+            ActiveRecord::Base.connection.execute("ALTER TABLE #{quoted_table_name} ALTER CONSTRAINT #{fk_constraint.name} DEFERRABLE INITIALLY IMMEDIATE")
+          end
+          fk_constraints += fk_constraints_on_table
+        end
+        # defer all fk constraints
+        ActiveRecord::Base.connection.execute("SET CONSTRAINTS #{fk_constraints.collect(&:name).join(',')} DEFERRED")
+        yield block
+      when :mysql
+        ActiveRecord::Base.connection.execute("SET foreign_key_checks = 0")
+        yield block
+        ActiveRecord::Base.connection.execute("SET foreign_key_checks = 1")
+      end
+    end
+
+    def self.database_type
+      case ActiveRecord::Base.connection.class.name
+      when 'ActiveRecord::ConnectionAdapters::PostgreSQLAdapter'
+        :postgresql
+      when 'ActiveRecord::ConnectionAdapters::Mysql2Adapter', 'ActiveRecord::ConnectionAdapters::MysqlAdapter'
+        :mysql
+      end
+    end
   end
 
   module Utils

lib/tasks/yaml_db_tasks.rake

@@ -47,6 +47,14 @@ namespace :db do
     desc "Load gzipped contents from stdin"
     task :load_gzipped_stdin => :environment do
       gz = Zlib::GzipReader.new($stdin.dup)
+      # workaround for ruby < 2.3.0
+      unless gz.respond_to? :external_encoding
+        class << gz
+          def external_encoding
+            Encoding::UTF_8
+          end
+        end
+      end
       SerializationHelper::Base.new(helper).load_from_io gz
       gz.close
     end

lib/yaml_db.rb

@@ -56,11 +56,11 @@ def self.table_record_header(io)
   end
 
   class Load < SerializationHelper::Load
-    def self.load_documents(io, truncate = true)
+    def self.load_documents(io)
         YAML.load_documents(io) do |ydoc|
           ydoc.keys.each do |table_name|
             next if ydoc[table_name].nil?
-            load_table(table_name, ydoc[table_name], truncate)
+            load_table(table_name, ydoc[table_name])
           end
         end
     end