Extract `truncate` and `truncate_tables` into database statements

This is to easier make `truncate_tables` to bulk statements.

13 files changed, 42 insertions, 84 deletions

diff --git a/activerecord/lib/active_record/connection_adapters/abstract/database_statements.rb b/activerecord/lib/active_record/connection_adapters/abstract/database_statements.rb
index 3c872d6c1b..bfd1c8402c 100644
--- a/activerecord/lib/active_record/connection_adapters/abstract/database_statements.rb
+++ b/activerecord/lib/active_record/connection_adapters/abstract/database_statements.rb
@@ -144,7 +144,11 @@ module ActiveRecord
 
       # Executes the truncate statement.
       def truncate(table_name, name = nil)
-        raise NotImplementedError
+        execute "TRUNCATE TABLE #{quote_table_name(table_name)}", name
+      end
+
+      def truncate_tables(*table_names) # :nodoc:
+        table_names.each { |table_name| truncate(table_name) }
       end
 
       # Executes update +sql+ statement in the context of this connection using
diff --git a/activerecord/lib/active_record/connection_adapters/abstract_mysql_adapter.rb b/activerecord/lib/active_record/connection_adapters/abstract_mysql_adapter.rb
index 6eefda763f..5ceddf449c 100644
--- a/activerecord/lib/active_record/connection_adapters/abstract_mysql_adapter.rb
+++ b/activerecord/lib/active_record/connection_adapters/abstract_mysql_adapter.rb
@@ -274,10 +274,6 @@ module ActiveRecord
         show_variable "collation_database"
       end
 
-      def truncate(table_name, name = nil)
-        execute "TRUNCATE TABLE #{quote_table_name(table_name)}", name
-      end
-
       def table_comment(table_name) # :nodoc:
         scope = quoted_scope(table_name)
 
diff --git a/activerecord/lib/active_record/connection_adapters/postgresql/database_statements.rb b/activerecord/lib/active_record/connection_adapters/postgresql/database_statements.rb
index 41633872e2..208934385f 100644
--- a/activerecord/lib/active_record/connection_adapters/postgresql/database_statements.rb
+++ b/activerecord/lib/active_record/connection_adapters/postgresql/database_statements.rb
@@ -143,6 +143,12 @@ module ActiveRecord
           end
         end
 
+        def truncate_tables(*table_names) # :nodoc:
+          unless table_names.empty?
+            execute "TRUNCATE TABLE #{table_names.map(&method(:quote_table_name)).join(", ")}"
+          end
+        end
+
         # Begins a transaction.
         def begin_db_transaction
           execute "BEGIN"
diff --git a/activerecord/lib/active_record/connection_adapters/postgresql_adapter.rb b/activerecord/lib/active_record/connection_adapters/postgresql_adapter.rb
index 2905660499..29f764e8f4 100644
--- a/activerecord/lib/active_record/connection_adapters/postgresql_adapter.rb
+++ b/activerecord/lib/active_record/connection_adapters/postgresql_adapter.rb
@@ -259,10 +259,6 @@ module ActiveRecord
         @use_insert_returning = @config.key?(:insert_returning) ? self.class.type_cast_config_to_boolean(@config[:insert_returning]) : true
       end
 
-      def truncate(table_name, name = nil)
-        exec_query "TRUNCATE TABLE #{quote_table_name(table_name)}", name, []
-      end
-
       # Is this connection alive and ready for queries?
       def active?
         @lock.synchronize do
diff --git a/activerecord/lib/active_record/connection_adapters/sqlite3_adapter.rb b/activerecord/lib/active_record/connection_adapters/sqlite3_adapter.rb
index 3004caf82d..8ee7e4c763 100644
--- a/activerecord/lib/active_record/connection_adapters/sqlite3_adapter.rb
+++ b/activerecord/lib/active_record/connection_adapters/sqlite3_adapter.rb
@@ -155,10 +155,6 @@ module ActiveRecord
         @connection.close rescue nil
       end
 
-      def truncate(table_name, name = nil)
-        execute "DELETE FROM #{quote_table_name(table_name)}", name
-      end
-
       def supports_index_sort_order?
         true
       end
@@ -279,6 +275,10 @@ module ActiveRecord
         end
       end
 
+      def truncate(table_name, name = nil) # :nodoc:
+        execute "DELETE FROM #{quote_table_name(table_name)}", name
+      end
+
       def begin_db_transaction #:nodoc:
         log("begin transaction", nil) { @connection.transaction }
       end
diff --git a/activerecord/lib/active_record/tasks/database_tasks.rb b/activerecord/lib/active_record/tasks/database_tasks.rb
index a8433fa0db..5a46003732 100644
--- a/activerecord/lib/active_record/tasks/database_tasks.rb
+++ b/activerecord/lib/active_record/tasks/database_tasks.rb
@@ -185,14 +185,17 @@ module ActiveRecord
       def truncate_tables(configuration)
         ActiveRecord::Base.connected_to(database: { truncation: configuration }) do
           table_names = ActiveRecord::Base.connection.tables
-          internal_table_names = [
+          table_names -= [
             ActiveRecord::Base.schema_migrations_table_name,
             ActiveRecord::Base.internal_metadata_table_name
           ]
 
-          class_for_adapter(configuration["adapter"]).new(configuration).truncate_tables(*table_names.without(*internal_table_names))
+          ActiveRecord::Base.connection.disable_referential_integrity do
+            ActiveRecord::Base.connection.truncate_tables(*table_names)
+          end unless table_names.empty?
         end
       end
+      private :truncate_tables
 
       def truncate_all(environment = env)
         ActiveRecord::Base.configurations.configs_for(env_name: environment).each do |db_config|
diff --git a/activerecord/lib/active_record/tasks/mysql_database_tasks.rb b/activerecord/lib/active_record/tasks/mysql_database_tasks.rb
index f2b4ead98d..1c1b29b5e1 100644
--- a/activerecord/lib/active_record/tasks/mysql_database_tasks.rb
+++ b/activerecord/lib/active_record/tasks/mysql_database_tasks.rb
@@ -31,16 +31,6 @@ module ActiveRecord
         connection.recreate_database configuration["database"], creation_options
       end
 
-      def truncate_tables(*table_names)
-        return if table_names.empty?
-
-        ActiveRecord::Base.connection.disable_referential_integrity do
-          table_names.each do |table_name|
-            ActiveRecord::Base.connection.truncate(table_name)
-          end
-        end
-      end
-
       def charset
         connection.charset
       end
diff --git a/activerecord/lib/active_record/tasks/postgresql_database_tasks.rb b/activerecord/lib/active_record/tasks/postgresql_database_tasks.rb
index dc368eb97d..8acb11f75f 100644
--- a/activerecord/lib/active_record/tasks/postgresql_database_tasks.rb
+++ b/activerecord/lib/active_record/tasks/postgresql_database_tasks.rb
@@ -48,18 +48,6 @@ module ActiveRecord
         create true
       end
 
-      def truncate_tables(*table_names)
-        return if table_names.empty?
-
-        ActiveRecord::Base.connection.disable_referential_integrity do
-          quoted_table_names = table_names.map do |table_name|
-            ActiveRecord::Base.connection.quote_table_name(table_name)
-          end
-
-          ActiveRecord::Base.connection.execute "TRUNCATE TABLE #{quoted_table_names.join(", ")}"
-        end
-      end
-
       def structure_dump(filename, extra_flags)
         set_psql_env
 
diff --git a/activerecord/lib/active_record/tasks/sqlite_database_tasks.rb b/activerecord/lib/active_record/tasks/sqlite_database_tasks.rb
index cedbae6b7f..a82cea80ca 100644
--- a/activerecord/lib/active_record/tasks/sqlite_database_tasks.rb
+++ b/activerecord/lib/active_record/tasks/sqlite_database_tasks.rb
@@ -33,16 +33,6 @@ module ActiveRecord
         create
       end
 
-      def truncate_tables(*table_names)
-        return if table_names.empty?
-
-        ActiveRecord::Base.connection.disable_referential_integrity do
-          table_names.each do |table_name|
-            ActiveRecord::Base.connection.truncate(table_name)
-          end
-        end
-      end
-
       def charset
         connection.encoding
       end
diff --git a/activerecord/test/cases/adapter_test.rb b/activerecord/test/cases/adapter_test.rb
index d90003b0ba..d25faf21d8 100644
--- a/activerecord/test/cases/adapter_test.rb
+++ b/activerecord/test/cases/adapter_test.rb
@@ -452,6 +452,8 @@ module ActiveRecord
   class AdapterTestWithoutTransaction < ActiveRecord::TestCase
     self.use_transactional_tests = false
 
+    fixtures :posts, :authors, :author_addresses
+
     def setup
       @connection = ActiveRecord::Base.connection
     end
@@ -482,6 +484,26 @@ module ActiveRecord
       end
     end
 
+    def test_truncate
+      assert_operator @connection.query_value("SELECT COUNT(*) FROM posts"), :>, 0
+
+      @connection.truncate("posts")
+
+      assert_equal 0, @connection.query_value("SELECT COUNT(*) FROM posts")
+    end
+
+    def test_truncate_tables
+      assert_operator @connection.query_value("SELECT COUNT(*) FROM authors"), :>, 0
+      assert_operator @connection.query_value("SELECT COUNT(*) FROM author_addresses"), :>, 0
+
+      @connection.disable_referential_integrity do
+        @connection.truncate_tables("author_addresses", "authors")
+      end
+
+      assert_equal 0, @connection.query_value("SELECT COUNT(*) FROM authors")
+      assert_equal 0, @connection.query_value("SELECT COUNT(*) FROM author_addresses")
+    end
+
     # test resetting sequences in odd tables in PostgreSQL
     if ActiveRecord::Base.connection.respond_to?(:reset_pk_sequence!)
       require "models/movie"
diff --git a/activerecord/test/cases/adapters/mysql2/connection_test.rb b/activerecord/test/cases/adapters/mysql2/connection_test.rb
index 3103589186..9c6566106a 100644
--- a/activerecord/test/cases/adapters/mysql2/connection_test.rb
+++ b/activerecord/test/cases/adapters/mysql2/connection_test.rb
@@ -28,17 +28,6 @@ class Mysql2ConnectionTest < ActiveRecord::Mysql2TestCase
     end
   end
 
-  def test_truncate
-    rows = ActiveRecord::Base.connection.exec_query("select count(*) from comments")
-    count = rows.first.values.first
-    assert_operator count, :>, 0
-
-    ActiveRecord::Base.connection.truncate("comments")
-    rows = ActiveRecord::Base.connection.exec_query("select count(*) from comments")
-    count = rows.first.values.first
-    assert_equal 0, count
-  end
-
   def test_no_automatic_reconnection_after_timeout
     assert_predicate @connection, :active?
     @connection.update("set @@wait_timeout=1")
diff --git a/activerecord/test/cases/adapters/postgresql/connection_test.rb b/activerecord/test/cases/adapters/postgresql/connection_test.rb
index 40ab158c05..210758f462 100644
--- a/activerecord/test/cases/adapters/postgresql/connection_test.rb
+++ b/activerecord/test/cases/adapters/postgresql/connection_test.rb
@@ -25,14 +25,6 @@ module ActiveRecord
       super
     end
 
-    def test_truncate
-      count = ActiveRecord::Base.connection.execute("select count(*) from comments").first["count"].to_i
-      assert_operator count, :>, 0
-      ActiveRecord::Base.connection.truncate("comments")
-      count = ActiveRecord::Base.connection.execute("select count(*) from comments").first["count"].to_i
-      assert_equal 0, count
-    end
-
     def test_encoding
       assert_queries(1) do
         assert_not_nil @connection.encoding
diff --git a/activerecord/test/cases/adapters/sqlite3/connection_test.rb b/activerecord/test/cases/adapters/sqlite3/connection_test.rb
deleted file mode 100644
index 3dabc8766a..0000000000
--- a/activerecord/test/cases/adapters/sqlite3/connection_test.rb
+++ /dev/null
@@ -1,18 +0,0 @@
-# frozen_string_literal: true
-
-require "cases/helper"
-
-class SQLite3ConnectionTest < ActiveRecord::SQLite3TestCase
-  fixtures :comments
-
-  def test_truncate
-    rows = ActiveRecord::Base.connection.exec_query("select count(*) from comments")
-    count = rows.first.values.first
-    assert_operator count, :>, 0
-
-    ActiveRecord::Base.connection.truncate("comments")
-    rows = ActiveRecord::Base.connection.exec_query("select count(*) from comments")
-    count = rows.first.values.first
-    assert_equal 0, count
-  end
-end