From f6b4e1017f44066a49d1e53f436ab54e35d67090 Mon Sep 17 00:00:00 2001
From: Naoko Nishimura <naoko.nishimura1018@gmail.com>
Date: Fri, 20 Sep 2024 18:30:13 -0700
Subject: [PATCH] Add option to specify CASCADE and RESTRICT for Truncation
 strategy

---
 README.md                                       |  2 ++
 .../active_record/truncation.rb                 | 17 ++++++++++-------
 2 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/README.md b/README.md
index d1bc665..c05de91 100644
--- a/README.md
+++ b/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:
diff --git a/lib/database_cleaner/active_record/truncation.rb b/lib/database_cleaner/active_record/truncation.rb
index 2e044a0..eac95ae 100644
--- a/lib/database_cleaner/active_record/truncation.rb
+++ b/lib/database_cleaner/active_record/truncation.rb
@@ -1,11 +1,12 @@
 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
@@ -13,6 +14,7 @@ def initialize(opts={})
 
         @reset_ids = opts[:reset_ids]
         @pre_count = opts[:pre_count]
+        @truncate_option = opts[:cascade] || :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)