From ef4bf94e2c962e0a5873eb0597bfb2115914f7cd Mon Sep 17 00:00:00 2001
From: Eddie Lebow <Eddie.Lebow@phishme.com>
Date: Wed, 26 Jan 2022 15:36:01 -0500
Subject: [PATCH] Accept `_` integer notation in VERSION arg to database tasks

Also correct unrelated typo in DatabaseTasks test description.

on-behalf-of: @Cofense <oss@cofense.com>
---
 activerecord/CHANGELOG.md                      |  4 ++++
 activerecord/lib/active_record/migration.rb    |  7 +++++++
 .../lib/active_record/tasks/database_tasks.rb  |  2 +-
 .../test/cases/tasks/database_tasks_test.rb    | 18 ++++++++++++++----
 4 files changed, 26 insertions(+), 5 deletions(-)

diff --git a/activerecord/CHANGELOG.md b/activerecord/CHANGELOG.md
index 63b43ace4f3..0288880b10e 100644
--- a/activerecord/CHANGELOG.md
+++ b/activerecord/CHANGELOG.md
@@ -1,3 +1,7 @@
+*   Permit underscores in the VERSION argument to database rake tasks.
+
+    *Eddie Lebow*
+
 *   Reversed the order of `INSERT` statements in `structure.sql` dumps
 
     This should decrease the likelihood of merge conflicts. New migrations
diff --git a/activerecord/lib/active_record/migration.rb b/activerecord/lib/active_record/migration.rb
index 374ee2edf96..1f026374dd6 100644
--- a/activerecord/lib/active_record/migration.rb
+++ b/activerecord/lib/active_record/migration.rb
@@ -575,6 +575,13 @@ module ActiveRecord
 
     MigrationFilenameRegexp = /\A([0-9]+)_([_a-z0-9]*)\.?([_a-z0-9]*)?\.rb\z/ # :nodoc:
 
+    def self.valid_version_format?(version_string) # :nodoc:
+      [
+        MigrationFilenameRegexp,
+        /\A\d(_?\d)*\z/ # integer with optional underscores
+      ].any? { |pattern| pattern.match?(version_string) }
+    end
+
     # This class is used to verify that all migrations have been run before
     # loading a web page if <tt>config.active_record.migration_error</tt> is set to :page_load
     class CheckPending
diff --git a/activerecord/lib/active_record/tasks/database_tasks.rb b/activerecord/lib/active_record/tasks/database_tasks.rb
index 5e783908a92..a68f72cb3fa 100644
--- a/activerecord/lib/active_record/tasks/database_tasks.rb
+++ b/activerecord/lib/active_record/tasks/database_tasks.rb
@@ -305,7 +305,7 @@ module ActiveRecord
       end
 
       def check_target_version
-        if target_version && !(Migration::MigrationFilenameRegexp.match?(ENV["VERSION"]) || /\A\d+\z/.match?(ENV["VERSION"]))
+        if target_version && !Migration.valid_version_format?(ENV["VERSION"])
           raise "Invalid format of target version: `VERSION=#{ENV['VERSION']}`"
         end
       end
diff --git a/activerecord/test/cases/tasks/database_tasks_test.rb b/activerecord/test/cases/tasks/database_tasks_test.rb
index c69524b0f3a..70c378c6680 100644
--- a/activerecord/test/cases/tasks/database_tasks_test.rb
+++ b/activerecord/test/cases/tasks/database_tasks_test.rb
@@ -1075,6 +1075,10 @@ module ActiveRecord
       e = assert_raise(RuntimeError) { ActiveRecord::Tasks::DatabaseTasks.migrate }
       assert_match(/Invalid format of target version/, e.message)
 
+      ENV["VERSION"] = "1__1"
+      e = assert_raise(RuntimeError) { ActiveRecord::Tasks::DatabaseTasks.migrate }
+      assert_match(/Invalid format of target version/, e.message)
+
       ENV["VERSION"] = "1_name"
       e = assert_raise(RuntimeError) { ActiveRecord::Tasks::DatabaseTasks.migrate }
       assert_match(/Invalid format of target version/, e.message)
@@ -1415,13 +1419,16 @@ module ActiveRecord
       version = ENV["VERSION"]
 
       ENV["VERSION"] = "0"
-      assert_equal ENV["VERSION"].to_i, ActiveRecord::Tasks::DatabaseTasks.target_version
+      assert_equal 0, ActiveRecord::Tasks::DatabaseTasks.target_version
 
       ENV["VERSION"] = "42"
-      assert_equal ENV["VERSION"].to_i, ActiveRecord::Tasks::DatabaseTasks.target_version
+      assert_equal 42, ActiveRecord::Tasks::DatabaseTasks.target_version
 
       ENV["VERSION"] = "042"
-      assert_equal ENV["VERSION"].to_i, ActiveRecord::Tasks::DatabaseTasks.target_version
+      assert_equal 42, ActiveRecord::Tasks::DatabaseTasks.target_version
+
+      ENV["VERSION"] = "2000_01_01_000042"
+      assert_equal 20000101000042, ActiveRecord::Tasks::DatabaseTasks.target_version
     ensure
       ENV["VERSION"] = version
     end
@@ -1436,7 +1443,7 @@ module ActiveRecord
       ENV["VERSION"] = version
     end
 
-    def test_check_target_version_does_not_raise_error_if_version_is_not_setted
+    def test_check_target_version_does_not_raise_error_if_version_is_not_set
       version = ENV.delete("VERSION")
       assert_nothing_raised { ActiveRecord::Tasks::DatabaseTasks.check_target_version }
     ensure
@@ -1489,6 +1496,9 @@ module ActiveRecord
       ENV["VERSION"] = "001"
       assert_nothing_raised { ActiveRecord::Tasks::DatabaseTasks.check_target_version }
 
+      ENV["VERSION"] = "1_001"
+      assert_nothing_raised { ActiveRecord::Tasks::DatabaseTasks.check_target_version }
+
       ENV["VERSION"] = "001_name.rb"
       assert_nothing_raised { ActiveRecord::Tasks::DatabaseTasks.check_target_version }
     ensure