Association inclusion, limit depth, unscoping, disable nullify, … (#13)

Co-authored-by: Polina Gurtovaya <zloymult@gmail.com>
Co-authored-by: Andrey Novikov <envek@envek.name>

README.md

@@ -52,14 +52,22 @@ EvilSeed.configure do |config|
     #
     # Association path is a dot-delimited string of association chain starting from model itself:
     # example: "forum.users.questions"
-    root.exclude(/\btracking_pixels\b/, 'forum.popular_questions')
+    root.exclude(/\btracking_pixels\b/, 'forum.popular_questions', /\Aforum\.parent\b/)
+
+    # Include back only certain associations
+    root.include(/\Aforum(\.parent(\.questions(\.answers)?)?)?\z/)
 
     # It's possible to limit the number of included into dump has_many and has_one records for every association
     # Note that belongs_to records for all not excluded associations are always dumped to keep referential integrity.
     root.limit_associations_size(100)
 
     # Or for certain association only
     root.limit_associations_size(10, 'forum.questions')
+
+    # Limit the depth of associations to be dumped from the root level
+    # All traverses through has_many, belongs_to, etc are counted
+    # So forum.subforums.subforums.questions.answers will be 5 levels deep
+    root.limit_deep(10)
   end
 
   # Everything you can pass to +where+ method will work as constraints:
@@ -95,6 +103,20 @@ EvilSeed.configure do |config|
   # This will remove the columns even if the model is not a root node and is
   # dumped via an association.
   config.ignore_columns("Profile", :name)
+
+  # Disable foreign key nullification for records that are not included in the dump
+  # By default, EvilSeed will nullify foreign keys for records that are not included in the dump
+  config.dont_nullify = true
+
+  # Unscope relations to include soft-deleted records etc
+  # This is useful when you want to include all records, including those that are hidden by default
+  # By default, EvilSeed will abide default scope of models
+  config.unscoped = true
+
+  # Verbose mode will print out the progress of the dump to the console along with writing the file
+  # By default, verbose mode is off
+  config.verbose = true
+  config.verbose_sql = true
 end
 ```
 

lib/evil_seed/configuration.rb

@@ -7,10 +7,14 @@
 module EvilSeed
   # This module holds configuration for creating dump: which models and their constraints
   class Configuration
-    attr_accessor :record_dumper_class
+    attr_accessor :record_dumper_class, :verbose, :verbose_sql, :unscoped, :dont_nullify
 
     def initialize
       @record_dumper_class = RecordDumper
+      @verbose = false
+      @verbose_sql = false
+      @unscoped = false
+      @dont_nullify = false
       @ignored_columns = Hash.new { |h, k| h[k] = [] }
     end
 
@@ -19,7 +23,7 @@ def roots
     end
 
     def root(model, *constraints)
-      new_root = Root.new(model, *constraints)
+      new_root = Root.new(model, dont_nullify, *constraints)
       yield new_root if block_given?
       roots << new_root
     end

lib/evil_seed/configuration/root.rb

@@ -5,16 +5,19 @@ class Configuration
     # Configuration for dumping some root model and its associations
     class Root
       attr_reader :model, :constraints
-      attr_reader :total_limit, :association_limits
-      attr_reader :exclusions
+      attr_reader :total_limit, :association_limits, :deep_limit, :dont_nullify
+      attr_reader :exclusions, :inclusions
 
       # @param model       [String]       Name of the model class to dump
       # @param constraints [String, Hash] Everything you can feed into +where+ to limit number of records
-      def initialize(model, *constraints)
+      def initialize(model, dont_nullify, *constraints)
         @model = model
         @constraints = constraints
         @exclusions = []
+        @inclusions = []
         @association_limits = {}
+        @deep_limit = nil
+        @dont_nullify = dont_nullify
       end
 
       # Exclude some of associations from the dump
@@ -23,6 +26,12 @@ def exclude(*association_patterns)
         @exclusions += association_patterns
       end
 
+      # Include some excluded associations back to the dump
+      # @param association_patterns Array<String, Regex> Patterns to exclude associated models from dump
+      def include(*association_patterns)
+        @inclusions += association_patterns
+      end
+
       # Limit number of records in all (if pattern is not provided) or given  associations to include into dump
       # @param limit               [Integer]       Maximum number of records in associations to include into dump
       # @param association_pattern [String, Regex] Pattern to limit number of records for certain associated models
@@ -34,8 +43,22 @@ def limit_associations_size(limit, association_pattern = nil)
         end
       end
 
+      # Limit deepenes of associations to include into dump
+      # @param limit               [Integer]       Maximum level to recursively dive into associations
+      def limit_deep(limit)
+        @deep_limit = limit
+      end
+
+      def do_not_nullify(nullify_flag)
+        @dont_nullify = nullify_flag
+      end
+
       def excluded?(association_path)
-        exclusions.any? { |exclusion| exclusion.match(association_path) }
+        exclusions.any? { |exclusion| association_path.match(exclusion) } #.match(association_path) }
+      end
+
+      def included?(association_path)
+        inclusions.any? { |inclusion| association_path.match(inclusion) } #.match(association_path) }
       end
     end
   end

lib/evil_seed/record_dumper.rb

@@ -65,11 +65,17 @@ def insert_statement
 
     def write!(attributes)
       # Remove non-insertable columns from attributes
-      attributes = attributes.slice(*insertable_column_names)
+      attributes = prepare(attributes.slice(*insertable_column_names))
+
+      if configuration.verbose_sql
+        puts("-- #{relation_dumper.association_path}\n")
+        puts(@tuples_written.zero? ? insert_statement : ",\n")
+        puts("  (#{attributes.join(', ')})")
+      end
 
       @output.write("-- #{relation_dumper.association_path}\n") && @header_written = true unless @header_written
       @output.write(@tuples_written.zero? ? insert_statement : ",\n")
-      @output.write("  (#{prepare(attributes).join(', ')})")
+      @output.write("  (#{attributes.join(', ')})")
       @tuples_written += 1
       @output.write(";\n") && @tuples_written = 0 if @tuples_written == MAX_TUPLES_PER_INSERT_STMT
     end
@@ -84,7 +90,7 @@ def prepare(attributes)
       attributes.map do |key, value|
         type = model_class.attribute_types[key]
         model_class.connection.quote(type.serialize(value))
-      end
+      end.flatten.compact
     end
   end
 end

lib/evil_seed/relation_dumper.rb

@@ -25,13 +25,15 @@ class RelationDumper
 
     attr_reader :relation, :root_dumper, :model_class, :association_path, :search_key, :identifiers, :nullify_columns,
                 :belongs_to_reflections, :has_many_reflections, :foreign_keys, :loaded_ids, :to_load_map,
-                :record_dumper, :inverse_reflection, :table_names, :options
+                :record_dumper, :inverse_reflection, :table_names, :options,
+                :current_deep, :verbose
 
-    delegate :root, :configuration, :total_limit, :loaded_map, to: :root_dumper
+    delegate :root, :configuration, :dont_nullify, :total_limit, :deep_limit, :loaded_map, to: :root_dumper
 
     def initialize(relation, root_dumper, association_path, **options)
       @relation               = relation
       @root_dumper            = root_dumper
+      @verbose                = configuration.verbose
       @identifiers            = options[:identifiers]
       @to_load_map            = Hash.new { |h, k| h[k] = [] }
       @foreign_keys           = Hash.new { |h, k| h[k] = [] }
@@ -46,15 +48,23 @@ def initialize(relation, root_dumper, association_path, **options)
       @belongs_to_reflections = setup_belongs_to_reflections
       @has_many_reflections   = setup_has_many_reflections
       @options                = options
+      @current_deep           = association_path.split('.').size
+      @dont_nullify           = dont_nullify
+
+      puts("- #{association_path}") if verbose
     end
 
     # Generate dump and write it into +io+
     # @return [Array<IO>] List of dump IOs for separate tables in order of dependencies (belongs_to are first)
     def call
       dump!
-      belongs_to_dumps = dump_belongs_to_associations!
-      has_many_dumps   = dump_has_many_associations!
-      [belongs_to_dumps, record_dumper.result, has_many_dumps].flatten.compact
+      if deep_limit and current_deep > deep_limit
+        [record_dumper.result].flatten.compact
+      else
+        belongs_to_dumps = dump_belongs_to_associations!
+        has_many_dumps   = dump_has_many_associations!
+        [belongs_to_dumps, record_dumper.result, has_many_dumps].flatten.compact
+      end
     end
 
     private
@@ -64,16 +74,22 @@ def dump!
       model_class.ignored_columns += Array(configuration.ignored_columns_for(model_class.sti_name))
       model_class.send(:reload_schema_from_cache) if ActiveRecord.version < Gem::Version.new("6.1.0.rc1") # See https://github.com/rails/rails/pull/37581
       if identifiers.present?
+        puts("  # #{search_key} => #{identifiers}") if verbose
         # Don't use AR::Base#find_each as we will get error on Oracle if we will have more than 1000 ids in IN statement
         identifiers.in_groups_of(MAX_IDENTIFIERS_IN_IN_STMT).each do |ids|
-          fetch_attributes(relation.where(search_key => ids.compact)).each do |attributes|
+          attrs = fetch_attributes(relation.where(search_key => ids.compact))
+          puts(" -- dumped #{attrs.size}") if verbose
+          attrs.each do |attributes|
             next unless check_limits!
             dump_record!(attributes)
           end
         end
       else
+        puts("  # #{relation.count}") if verbose
         relation.in_batches do |relation|
-          fetch_attributes(relation).each do |attributes|
+          attrs = fetch_attributes(relation)
+          puts(" -- dumped #{attrs.size}") if verbose
+          attrs.each do |attributes|
             next unless check_limits!
             dump_record!(attributes)
           end
@@ -84,8 +100,10 @@ def dump!
     end
 
     def dump_record!(attributes)
-      nullify_columns.each do |nullify_column|
-        attributes[nullify_column] = nil
+      unless dont_nullify
+        nullify_columns.each do |nullify_column|
+          attributes[nullify_column] = nil
+        end
       end
       return unless record_dumper.call(attributes)
       foreign_keys.each do |reflection_name, fk_column|
@@ -140,7 +158,11 @@ def check_limits!
     end
 
     def build_relation(reflection)
-      relation = reflection.klass.all
+      if configuration.unscoped
+        relation = reflection.klass.unscoped
+      else
+        relation = reflection.klass.all
+      end
       relation = relation.instance_eval(&reflection.scope) if reflection.scope
       relation = relation.where(reflection.type => model_class.to_s) if reflection.options[:as] # polymorphic
       relation
@@ -149,23 +171,32 @@ def build_relation(reflection)
     def setup_belongs_to_reflections
       model_class.reflect_on_all_associations(:belongs_to).reject do |reflection|
         next false if reflection.options[:polymorphic] # TODO: Add support for polymorphic belongs_to
+        included = root.included?("#{association_path}.#{reflection.name}")
         excluded = root.excluded?("#{association_path}.#{reflection.name}") || reflection.name == inverse_reflection
-        if excluded
-          nullify_columns << reflection.foreign_key if model_class.column_names.include?(reflection.foreign_key)
+        if excluded and not included
+          if model_class.column_names.include?(reflection.foreign_key)
+            puts(" -- excluded #{reflection.foreign_key}") if verbose
+            nullify_columns << reflection.foreign_key
+          end
         else
           foreign_keys[reflection.name] = reflection.foreign_key
           table_names[reflection.name]  = reflection.table_name
         end
-        excluded
+        excluded and not included
       end
     end
 
     # This method returns only direct has_one and has_many reflections. For HABTM it returns intermediate has_many
     def setup_has_many_reflections
+      puts(" -- reflections #{model_class._reflections.keys}") if verbose
       model_class._reflections.select do |_reflection_name, reflection|
         next false if model_class.primary_key.nil?
+
         next false if reflection.is_a?(ActiveRecord::Reflection::ThroughReflection)
-        %i[has_one has_many].include?(reflection.macro) && !root.excluded?("#{association_path}.#{reflection.name}")
+
+        included = root.included?("#{association_path}.#{reflection.name}")
+        excluded = root.excluded?("#{association_path}.#{reflection.name}") || reflection.name == inverse_reflection
+        %i[has_one has_many].include?(reflection.macro) && !(excluded and not included)
       end.map(&:second)
     end
   end

lib/evil_seed/root_dumper.rb

@@ -5,7 +5,7 @@
 module EvilSeed
   # This module collects dumps generation for root and all it's dependencies
   class RootDumper
-    attr_reader :root, :dumper, :model_class, :total_limit, :association_limits
+    attr_reader :root, :dumper, :model_class, :total_limit, :deep_limit, :dont_nullify, :association_limits
 
     delegate :loaded_map, :configuration, to: :dumper
 
@@ -14,6 +14,8 @@ def initialize(root, dumper)
       @dumper = dumper
       @to_load_map = {}
       @total_limit = root.total_limit
+      @deep_limit = root.deep_limit
+      @dont_nullify = root.dont_nullify
       @association_limits = root.association_limits.dup
 
       @model_class = root.model.constantize
@@ -24,6 +26,7 @@ def initialize(root, dumper)
     def call
       association_path = model_class.model_name.singular
       relation = model_class.all
+      relation = relation.unscoped if configuration.unscoped
       relation = relation.where(*root.constraints) if root.constraints.any? # without arguments returns not a relation
       RelationDumper.new(relation, self, association_path).call
     end

test/db/models.rb

@@ -49,6 +49,8 @@ class Answer < ActiveRecord::Base
   has_many :votes, as: :votable
 
   has_many :voters, through: :votes, source: :user
+
+  default_scope { where(deleted_at: nil) }
 end
 
 class TrackingPixel < ActiveRecord::Base

test/db/schema.rb

@@ -40,6 +40,7 @@ def create_schema!
       t.boolean    :best, default: false
       t.text       :text
       t.references :author
+      t.datetime   :deleted_at
       t.timestamps null: false
     end
     add_foreign_key :answers, :users, column: :author_id, on_delete: :nullify

test/db/seeds.rb

@@ -43,6 +43,12 @@
   author: User.find_by!(login: 'bob'),
 )
 
+answer = question.answers.create!(
+  text:   'Oops, I was wrong',
+  author: User.find_by!(login: 'eva'),
+  deleted_at: Time.current,
+)
+
 answer.votes.create!(user: User.find_by!(login: 'eva'))
 
 question_attrs = %w[first second third fourth fifth].map { |name| { name: name, forum: forums.first } }

test/evil_seed/dumper_test.rb

@@ -42,6 +42,7 @@ def test_it_dumps_tree_structures_with_foreign_keys
       assert_match(/'fourth'/, result)
       assert_match(/'fifth'/, result)
       assert result.index(/'One'/) < result.index(/'Descendant forum'/)
+      refute_match(/'Oops, I was wrong'/, result)
     end
 
     def test_limits_being_applied
@@ -59,5 +60,22 @@ def test_limits_being_applied
       assert_match(/'fourth'/, result)
       refute_match(/'fifth'/, result)
     end
+
+    def test_it_applies_unscoping_and_inclusions
+      configuration = EvilSeed::Configuration.new
+      configuration.root('Forum', name: 'Descendant forum') do |root|
+        root.include(/forum(\.parent(\.questions(\.answers)?)?)?\z/)
+        root.exclude(/.\..+/)
+      end
+      configuration.unscoped = true
+
+      io = StringIO.new
+      EvilSeed::Dumper.new(configuration).call(io)
+      result = io.string
+      File.write(File.join('tmp', "#{__method__}.sql"), result)
+      assert io.closed?
+      assert_match(/'Descendant forum'/, result)
+      assert_match(/'Oops, I was wrong'/, result)
+    end
   end
 end