truncate_option to allow CASCADE mode alongside the default RESTRICT (#118)

* Add option to specify CASCADE and RESTRICT for Truncation strategy

* Add tests for truncate_option

---------

Co-authored-by: Naoko Nishimura <naoko.nishimura1018@gmail.com>

Author: pat

README.md

@@ -63,6 +63,8 @@ DatabaseCleaner[:active_record].strategy = DatabaseCleaner::ActiveRecord::Deleti
 
 * `:reset_ids` - Only valid for deletion strategy, when set to `true` resets ids to 1 after each table is cleaned.
 
+* `:truncate_option` - Only valid for PostgreSQL. Acceptable values are `:restrict` and `:cascade`. Default is `:restrict`
+
 ## Adapter configuration options
 
 `#db` defaults to the default ActiveRecord database, but can be specified manually in a few ways:

lib/database_cleaner/active_record/truncation.rb

@@ -1,18 +1,20 @@
 require "delegate"
+require 'database_cleaner/active_record/base'
 
 module DatabaseCleaner
   module ActiveRecord
     class Truncation < Base
       def initialize(opts={})
-        if !opts.empty? && !(opts.keys - [:only, :except, :pre_count, :cache_tables, :reset_ids]).empty?
-          raise ArgumentError, "The only valid options are :only, :except, :pre_count, :reset_ids, and :cache_tables. You specified #{opts.keys.join(',')}."
+        if !opts.empty? && !(opts.keys - [:only, :except, :pre_count, :cache_tables, :reset_ids, :truncate_option]).empty?
+          raise ArgumentError, "The only valid options are :only, :except, :pre_count, :reset_ids, :cache_tables and :truncate_option. You specified #{opts.keys.join(',')}."
         end
 
         @only = Array(opts[:only]).dup
         @except = Array(opts[:except]).dup
 
         @reset_ids = opts[:reset_ids]
         @pre_count = opts[:pre_count]
+        @truncate_option = opts[:truncate_option] || :restrict
         @cache_tables = opts.has_key?(:cache_tables) ? !!opts[:cache_tables] : true
       end
 
@@ -21,7 +23,7 @@ def clean
           if pre_count? && connection.respond_to?(:pre_count_truncate_tables)
             connection.pre_count_truncate_tables(tables_to_clean(connection))
           else
-            connection.truncate_tables(tables_to_clean(connection))
+            connection.truncate_tables(tables_to_clean(connection), { truncate_option: @truncate_option })
           end
         end
       end
@@ -101,7 +103,7 @@ def truncate_table(table_name)
           execute("DELETE FROM #{quote_table_name(table_name)}")
         end
 
-        def truncate_tables(tables)
+        def truncate_tables(tables, opts)
           tables.each { |t| truncate_table(t) }
         end
       end
@@ -151,7 +153,7 @@ def truncate_table(table_name)
           end
         end
 
-        def truncate_tables(tables)
+        def truncate_tables(tables, opts)
           tables.each { |t| truncate_table(t) }
         end
 
@@ -192,9 +194,10 @@ def database_tables
           tables_with_schema
         end
 
-        def truncate_tables(table_names)
+        def truncate_tables(table_names, opts)
           return if table_names.nil? || table_names.empty?
-          execute("TRUNCATE TABLE #{table_names.map{|name| quote_table_name(name)}.join(', ')} RESTART IDENTITY RESTRICT;")
+
+          execute("TRUNCATE TABLE #{table_names.map{|name| quote_table_name(name)}.join(', ')} RESTART IDENTITY #{opts[:truncate_option]};")
         end
 
         def pre_count_truncate_tables(tables)

spec/database_cleaner/active_record/truncation_spec.rb

@@ -100,6 +100,28 @@
           end
         end
 
+        describe "with truncate_option set to cascade" do
+          subject(:strategy) { described_class.new(truncate_option: :cascade) }
+
+          it "specifies cascade when truncating" do
+            User.create!({name: 1})
+
+            expect(strategy.send(:connection)).to receive(:truncate_tables).with(['users', 'agents'], {truncate_option: :cascade})
+            strategy.clean
+          end
+        end
+
+        describe "with no truncate_option set" do
+          subject(:strategy) { described_class.new }
+
+          it "specifies restrict when truncating" do
+            User.create!
+
+            expect(strategy.send(:connection)).to receive(:truncate_tables).with(['users', 'agents'], {truncate_option: :restrict})
+            strategy.clean
+          end
+        end
+
         context 'when :cache_tables is set to true' do
           subject(:strategy) { described_class.new(cache_tables: true) }