data_migrate 9.1.0 → 9.4.1

data/lib/data_migrate/database_tasks.rb

@@ -5,86 +5,142 @@ require "data_migrate/config"
 module DataMigrate
   ##
   # This class extends DatabaseTasks to add a schema_file method.
-  class DatabaseTasks
+  module DatabaseTasks
     extend ActiveRecord::Tasks::DatabaseTasks
-
-    class << self
-      def schema_file(_format = nil)
-        File.join(db_dir, "data_schema.rb")
-      end
-
-      def schema_file_type(_format = nil)
-        "data_schema.rb"
-      end
-
-      def dump_filename(spec_name, format = ActiveRecord::Base.schema_format)
-        filename = if spec_name == "primary"
-          schema_file_type(format)
+    extend self
+
+    # These method are only introduced in Rails 7.1
+    unless respond_to?(:with_temporary_connection_for_each)
+      def with_temporary_connection_for_each(env: ActiveRecord::Tasks::DatabaseTasks.env, name: nil, &block) # :nodoc:
+        if name
+          db_config = ActiveRecord::Base.configurations.configs_for(env_name: env, name: name)
+          with_temporary_connection(db_config, &block)
         else
-          "#{spec_name}_#{schema_file_type(format)}"
+          ActiveRecord::Base.configurations.configs_for(env_name: env, name: name).each do |db_config|
+            with_temporary_connection(db_config, &block)
+          end
         end
-
-        ENV["DATA_SCHEMA"] || File.join(db_dir, filename)
       end
 
-      def check_schema_file(filename)
-        unless File.exist?(filename)
-          message = +%{#{filename} doesn't exist yet. Run `rake data:migrate` to create it, then try again.}
-          Kernel.abort message
+      def with_temporary_connection(db_config) # :nodoc:
+        with_temporary_pool(db_config) do |pool|
+          yield pool.connection
         end
       end
 
-      def pending_migrations
-        sort_migrations(
-          pending_schema_migrations,
-          pending_data_migrations
-        )
+      def migration_class # :nodoc:
+        ActiveRecord::Base
       end
 
-      def sort_migrations(*migrations)
-        migrations.flatten.sort { |a, b|  sort_string(a) <=> sort_string(b) }
+      def migration_connection # :nodoc:
+        migration_class.connection
       end
 
-      def sort_string migration
-        "#{migration[:version]}_#{migration[:kind] == :data ? 1 : 0}"
-      end
+      private def with_temporary_pool(db_config)
+        original_db_config = migration_class.connection_db_config
+        pool = migration_class.connection_handler.establish_connection(db_config)
 
-      def data_migrations_path
-        ::DataMigrate.config.data_migrations_path
+        yield pool
+      ensure
+        migration_class.connection_handler.establish_connection(original_db_config)
       end
+    end
 
-      def run_migration(migration, direction)
-        if migration[:kind] == :data
-          ::ActiveRecord::Migration.write("== %s %s" % ['Data', "=" * 71])
-          ::DataMigrate::DataMigrator.run(direction, data_migrations_path, migration[:version])
-        else
-          ::ActiveRecord::Migration.write("== %s %s" % ['Schema', "=" * 69])
-          ::DataMigrate::SchemaMigration.run(
-            direction,
-            ::DataMigrate::SchemaMigration.migrations_paths,
-            migration[:version]
-          )
+    def db_configs_with_versions
+      db_configs_with_versions = Hash.new { |h, k| h[k] = [] }
+
+      with_temporary_connection_for_each do |conn|
+        db_config = conn.pool.db_config
+        if db_config.primary?
+          versions_to_run = DataMigrate::DatabaseTasks.pending_data_migrations.map { |m| m[:version] }
+          target_version = ActiveRecord::Tasks::DatabaseTasks.target_version
+
+          versions_to_run.each do |version|
+            next if target_version && target_version != version
+            db_configs_with_versions[version] << DatabaseConfigurationWrapper.new(db_config)
+          end
         end
       end
 
-      def schema_dump_path(db_config, format = ActiveRecord.schema_format)
-        return ENV["DATA_SCHEMA"] if ENV["DATA_SCHEMA"]
+      db_configs_with_versions
+    end
+
+    def schema_file(_format = nil)
+      File.join(db_dir, "data_schema.rb")
+    end
+
+    def schema_file_type(_format = nil)
+      "data_schema.rb"
+    end
+
+    # This method is removed in Rails 7.0
+    def dump_filename(spec_name, format = ActiveRecord::Base.schema_format)
+      filename = if spec_name == "primary"
+        schema_file_type(format)
+      else
+        "#{spec_name}_#{schema_file_type(format)}"
+      end
 
-        # We only require a schema.rb file for the primary database
-        return unless db_config.primary?
+      ENV["DATA_SCHEMA"] || File.join(db_dir, filename)
+    end
 
-        super.gsub(/(_)?schema\.rb\z/, '\1data_schema.rb')
+    def check_schema_file(filename)
+      unless File.exist?(filename)
+        message = +%{#{filename} doesn't exist yet. Run `rake data:migrate` to create it, then try again.}
+        Kernel.abort message
       end
+    end
+
+    def pending_migrations
+      sort_migrations(
+        pending_schema_migrations,
+        pending_data_migrations
+      )
+    end
+
+    def sort_migrations(*migrations)
+      migrations.flatten.sort { |a, b|  sort_string(a) <=> sort_string(b) }
+    end
+
+    def sort_string migration
+      "#{migration[:version]}_#{migration[:kind] == :data ? 1 : 0}"
+    end
 
-      # Override this method from `ActiveRecord::Tasks::DatabaseTasks`
-      # to ensure that the sha saved in ar_internal_metadata table
-      # is from the original schema.rb file
-      def schema_sha1(file)
-        super(file.gsub(/data_schema.rb\z/, 'schema.rb'))
+    def data_migrations_path
+      ::DataMigrate.config.data_migrations_path
+    end
+
+    def run_migration(migration, direction)
+      if migration[:kind] == :data
+        ::ActiveRecord::Migration.write("== %s %s" % ['Data', "=" * 71])
+        ::DataMigrate::DataMigrator.run(direction, data_migrations_path, migration[:version])
+      else
+        ::ActiveRecord::Migration.write("== %s %s" % ['Schema', "=" * 69])
+        ::DataMigrate::SchemaMigration.run(
+          direction,
+          ::DataMigrate::SchemaMigration.migrations_paths,
+          migration[:version]
+        )
       end
     end
 
-    def self.forward(step = 1)
+    def schema_dump_path(db_config, format = ActiveRecord.schema_format)
+      return ENV["DATA_SCHEMA"] if ENV["DATA_SCHEMA"]
+
+      # We only require a schema.rb file for the primary database
+      return unless db_config.primary?
+
+      File.join(File.dirname(ActiveRecord::Tasks::DatabaseTasks.schema_dump_path(db_config, format)), schema_file_type)
+    end
+
+    # Override this method from `ActiveRecord::Tasks::DatabaseTasks`
+    # to ensure that the sha saved in ar_internal_metadata table
+    # is from the original schema.rb file
+    def schema_sha1(file)
+      ActiveRecord::Tasks::DatabaseTasks.schema_dump_path(ActiveRecord::Base.configurations.configs_for(env_name: ActiveRecord::Tasks::DatabaseTasks.env, name: "primary"))
+    end
+
+    def forward(step = 1)
       DataMigrate::DataMigrator.create_data_schema_table
       migrations = pending_migrations.reverse.pop(step).reverse
       migrations.each do | pending_migration |
@@ -98,22 +154,55 @@ module DataMigrate
       end
     end
 
-    def self.pending_data_migrations
+    def pending_data_migrations
       data_migrations = DataMigrate::DataMigrator.migrations(data_migrations_path)
-      sort_migrations(DataMigrate::DataMigrator.new(:up, data_migrations ).
-        pending_migrations.map {|m| { version: m.version, name: m.name, kind: :data }})
+      data_migrator = DataMigrate::RailsHelper.data_migrator(:up, data_migrations)
+      sort_migrations(
+        data_migrator.pending_migrations.map { |m| { version: m.version, name: m.name, kind: :data } }
+        )
     end
 
-    def self.pending_schema_migrations
+    def pending_schema_migrations
       ::DataMigrate::SchemaMigration.pending_schema_migrations
     end
 
-    def self.past_migrations(sort = nil)
-      data_versions = DataMigrate::DataSchemaMigration.table_exists? ? DataMigrate::DataSchemaMigration.normalized_versions : []
-      schema_versions = ActiveRecord::SchemaMigration.normalized_versions
+    def past_migrations(sort = nil)
+      data_versions = DataMigrate::RailsHelper.data_schema_migration.table_exists? ? DataMigrate::RailsHelper.data_schema_migration.normalized_versions : []
+      schema_versions = DataMigrate::RailsHelper.schema_migration.normalized_versions
       migrations = data_versions.map { |v| { version: v.to_i, kind: :data } } + schema_versions.map { |v| { version: v.to_i, kind: :schema } }
 
       sort&.downcase == "asc" ? sort_migrations(migrations) : sort_migrations(migrations).reverse
     end
+
+    def self.migrate_with_data
+      DataMigrate::DataMigrator.create_data_schema_table
+
+      ActiveRecord::Migration.verbose = ENV["VERBOSE"] ? ENV["VERBOSE"] == "true" : true
+
+      db_configs = ActiveRecord::Base.configurations.configs_for(env_name: ActiveRecord::Tasks::DatabaseTasks.env)
+
+      schema_mapped_versions = ActiveRecord::Tasks::DatabaseTasks.db_configs_with_versions(db_configs)
+      data_mapped_versions = DataMigrate::DatabaseTasks.db_configs_with_versions
+
+      mapped_versions = schema_mapped_versions.merge(data_mapped_versions) do |_key, schema_db_configs, data_db_configs|
+        schema_db_configs + data_db_configs
+      end
+
+      mapped_versions.sort.each do |version, db_configs|
+        db_configs.each do |db_config|
+          if is_data_migration = db_config.is_a?(DataMigrate::DatabaseConfigurationWrapper)
+            db_config = db_config.db_config
+          end
+
+          DataMigrate::DatabaseTasks.with_temporary_connection(db_config) do
+            if is_data_migration
+              DataMigrate::DataMigrator.run(:up, DataMigrate::DatabaseTasks.data_migrations_path, version)
+            else
+              ActiveRecord::Tasks::DatabaseTasks.migrate(version)
+            end
+          end
+        end
+      end
+    end
   end
 end

data/lib/data_migrate/migration_context.rb

@@ -11,7 +11,8 @@ module DataMigrate
                               migrations
                             end
 
-      DataMigrator.new(:up, selected_migrations, target_version).migrate
+      data_migrator = DataMigrate::RailsHelper.data_migrator(:up, selected_migrations, DataMigrate::RailsHelper.schema_migration, DataMigrate::RailsHelper.internal_metadata, target_version)
+      data_migrator.migrate
     end
 
     def down(target_version = nil)
@@ -22,11 +23,13 @@ module DataMigrate
           migrations
         end
 
-      DataMigrator.new(:down, selected_migrations, target_version).migrate
+      data_migrator = DataMigrate::RailsHelper.data_migrator(:down, selected_migrations, DataMigrate::RailsHelper.schema_migration, DataMigrate::RailsHelper.internal_metadata, target_version)
+      data_migrator.migrate
     end
 
     def run(direction, target_version)
-      DataMigrator.new(direction, migrations, target_version).run
+      data_migrator = DataMigrate::RailsHelper.data_migrator(direction, migrations, DataMigrate::RailsHelper.schema_migration, DataMigrate::RailsHelper.internal_metadata, target_version)
+      data_migrator.run
     end
 
     def current_version
@@ -40,12 +43,12 @@ module DataMigrate
     end
 
     def migrations_status
-      db_list = DataSchemaMigration.normalized_versions
+      db_list = DataMigrate::RailsHelper.data_schema_migration.normalized_versions
 
       file_list = migration_files.map do |file|
         version, name, scope = parse_migration_filename(file)
         raise ActiveRecord::IllegalMigrationNameError.new(file) unless version
-        version = ActiveRecord::SchemaMigration.normalize_migration_number(version)
+        version = DataMigrate::RailsHelper.schema_migration.normalize_migration_number(version)
         status = db_list.delete(version) ? "up" : "down"
         [status, version, (name + scope).humanize]
       end.compact
@@ -60,15 +63,15 @@ module DataMigrate
     private
 
     def get_all_versions
-      if DataMigrate::DataSchemaMigration.table_exists?
-        DataSchemaMigration.normalized_versions.map(&:to_i)
+      if DataMigrate::RailsHelper.data_schema_migration.table_exists?
+        DataMigrate::RailsHelper.data_schema_migration.normalized_versions.map(&:to_i)
       else
         []
       end
     end
 
     def move(direction, steps)
-      migrator = DataMigrator.new(direction, migrations)
+      migrator = DataMigrate::RailsHelper.data_migrator(direction, migrations)
 
       if current_version != 0 && !migrator.current_migration
         raise ActiveRecord::UnknownMigrationVersionError.new(current_version)

data/lib/data_migrate/rails_helper.rb

@@ -0,0 +1,79 @@
+module DataMigrate
+  class RailsHelper
+    class << self
+      def rails_version_equal_to_or_higher_than_7_1
+        return @equal_to_or_higher_than_7_1 if defined?(@equal_to_or_higher_than_7_1)
+
+        @equal_to_or_higher_than_7_1 = Gem::Dependency.new("railties", ">= 7.1.0.alpha").match?("railties", Gem.loaded_specs["railties"].version, true)
+      end
+
+      def rails_version_equal_to_or_higher_than_7_0
+        return @rails_version_equal_to_or_higher_than_7_0 if defined?(@rails_version_equal_to_or_higher_than_7_0)
+
+        @rails_version_equal_to_or_higher_than_7_0 = Gem::Dependency.new("railties", ">= 7.0").match?("railties", Gem.loaded_specs["railties"].version, true)
+      end
+
+      def internal_metadata
+        if rails_version_equal_to_or_higher_than_7_1
+          ActiveRecord::Base.connection.internal_metadata
+        else
+          ActiveRecord::InternalMetadata
+        end
+      end
+
+      def schema_migration
+        if rails_version_equal_to_or_higher_than_7_1
+          ActiveRecord::Base.connection.schema_migration
+        else
+          ActiveRecord::SchemaMigration
+        end
+      end
+
+      def schema_migration_versions
+        if rails_version_equal_to_or_higher_than_7_1
+          schema_migration.versions
+        else
+          schema_migration.all.pluck(:version)
+        end
+      end
+
+      def schema_create_version(version)
+        if rails_version_equal_to_or_higher_than_7_1
+          schema_migration.create_version(version)
+        else
+          schema_migration.create(version: version)
+        end
+      end
+
+      def data_schema_delete_version(version)
+        if rails_version_equal_to_or_higher_than_7_1
+          data_schema_migration.delete_version(version)
+        else
+          data_schema_migration.where(version: version.to_s).delete_all
+        end
+      end
+
+      def data_schema_migration
+        if rails_version_equal_to_or_higher_than_7_1
+          DataMigrate::DataSchemaMigration.new(ActiveRecord::Tasks::DatabaseTasks.migration_connection)
+        else
+          DataMigrate::DataSchemaMigration
+        end
+      end
+
+      def data_migrator(
+        direction,
+        migrations,
+        schema_migration = DataMigrate::RailsHelper.schema_migration,
+        internal_metadata = DataMigrate::RailsHelper.internal_metadata,
+        target_version = nil
+      )
+        if rails_version_equal_to_or_higher_than_7_1
+          DataMigrate::DataMigrator.new(direction, migrations, schema_migration, internal_metadata, target_version)
+        else
+          DataMigrate::DataMigrator.new(direction, migrations, schema_migration, target_version)
+        end
+      end
+    end
+  end
+end

data/lib/data_migrate/schema_dumper.rb

@@ -26,7 +26,7 @@ module DataMigrate
 
     def initialize(connection)
       @connection = connection
-      all_versions =  DataSchemaMigration.normalized_versions
+      all_versions =  DataMigrate::RailsHelper.data_schema_migration.normalized_versions
 
       @version = begin
                     all_versions.max

data/lib/data_migrate/schema_migration.rb

@@ -7,14 +7,14 @@ module DataMigrate
     def self.pending_schema_migrations
       all_migrations = DataMigrate::MigrationContext.new(migrations_paths).migrations
       sort_migrations(
-        ActiveRecord::Migrator.new(:up, all_migrations, ActiveRecord::Base.connection.schema_migration).
+        ActiveRecord::Migrator.new(:up, all_migrations, DataMigrate::RailsHelper.schema_migration, DataMigrate::RailsHelper.internal_metadata).
         pending_migrations.
         map {|m| { version: m.version, kind: :schema }}
       )
     end
 
     def self.run(direction, migration_paths, version)
-      ActiveRecord::MigrationContext.new(migration_paths, ActiveRecord::Base.connection.schema_migration).run(direction, version)
+      ActiveRecord::MigrationContext.new(migration_paths, DataMigrate::RailsHelper.schema_migration).run(direction, version)
     end
 
     def self.sort_migrations(set1, set2 = nil)
@@ -24,9 +24,10 @@ module DataMigrate
 
     def self.migrations_paths
       spec_name = DataMigrate.config.spec_name
-      if spec_name && Gem::Dependency.new("railties", "~> 7.0").match?("railties", Gem.loaded_specs["railties"].version)
+      # The positional argument true is to include pre-release versions, such as 7.1.0.alpha
+      if spec_name && Gem::Dependency.new("railties", ">= 7.0").match?("railties", Gem.loaded_specs["railties"].version, true)
         ActiveRecord::Base.configurations.configs_for(env_name: Rails.env, name: spec_name).migrations_paths
-      elsif spec_name && Gem::Dependency.new("railties", "~> 6.0").match?("railties", Gem.loaded_specs["railties"].version)
+      elsif spec_name && Gem::Dependency.new("railties", "~> 6.1").match?("railties", Gem.loaded_specs["railties"].version)
         ActiveRecord::Base.configurations.configs_for(env_name: Rails.env, spec_name: spec_name).migrations_paths
       else
         Rails.application.config.paths["db/migrate"].to_a

data/lib/data_migrate/status_service.rb

@@ -24,11 +24,11 @@ module DataMigrate
     private
 
     def table_name
-      DataMigrate::DataSchemaMigration.table_name
+      DataMigrate::RailsHelper.data_schema_migration.table_name
     end
 
     def output(stream)
-      unless DataMigrate::DataSchemaMigration.table_exists?
+      unless DataMigrate::RailsHelper.data_schema_migration.table_exists?
         stream.puts "Data migrations table does not exist yet."
         return
       end
@@ -46,7 +46,7 @@ module DataMigrate
     def database_name
       if Gem::Dependency.new("railties", "~> 7.0").match?("railties", Gem.loaded_specs["railties"].version)
         ActiveRecord::Base.connection_db_config.configuration_hash[:database]
-      elsif Gem::Dependency.new("railties", "~> 6.0").match?("railties", Gem.loaded_specs["railties"].version)
+      elsif Gem::Dependency.new("railties", "~> 6.1").match?("railties", Gem.loaded_specs["railties"].version)
         ActiveRecord::Base.connection_config[:database]
       end
     end

data/lib/data_migrate/tasks/data_migrate_tasks.rb

@@ -50,9 +50,9 @@ module DataMigrate
 
       def status_with_schema
         db_list_data = ActiveRecord::Base.connection.select_values(
-          "SELECT version FROM #{DataMigrate::DataSchemaMigration.table_name}"
+          "SELECT version FROM #{DataMigrate::RailsHelper.data_schema_migration.table_name}"
         )
-        db_list_schema = ActiveRecord::SchemaMigration.all.pluck(:version)
+        db_list_schema = DataMigrate::RailsHelper.schema_migration_versions
         file_list = []
 
         Dir.foreach(File.join(Rails.root, migrations_paths)) do |file|
@@ -96,7 +96,7 @@ module DataMigrate
       def database_name
         if Gem::Dependency.new("railties", "~> 7.0").match?("railties", Gem.loaded_specs["railties"].version)
           ActiveRecord::Base.connection_db_config.database
-        elsif Gem::Dependency.new("railties", "~> 6.0").match?("railties", Gem.loaded_specs["railties"].version)
+        elsif Gem::Dependency.new("railties", "~> 6.1").match?("railties", Gem.loaded_specs["railties"].version)
           ActiveRecord::Base.connection_config[:database]
         end
       end

data/lib/data_migrate/test.rb

@@ -0,0 +1,14 @@
+module Base
+  extend self
+
+  def foo
+    puts "Base#foo called"
+  end
+end
+
+module Child
+  extend Base
+  extend self
+
+  puts "foo: #{respond_to?(:foo)}"
+end

data/lib/data_migrate/version.rb

@@ -1,3 +1,3 @@
 module DataMigrate
-  VERSION = "9.1.0".freeze
+  VERSION = "9.4.1".freeze
 end

data/lib/data_migrate.rb

@@ -1,5 +1,6 @@
 # frozen_string_literal: true
 
+require File.join(File.dirname(__FILE__), "data_migrate", "rails_helper")
 require File.join(File.dirname(__FILE__), "data_migrate", "data_migrator")
 require File.join(File.dirname(__FILE__), "data_migrate", "data_schema_migration")
 require File.join(File.dirname(__FILE__), "data_migrate", "data_schema")
@@ -9,9 +10,9 @@ require File.join(File.dirname(__FILE__), "data_migrate", "status_service")
 require File.join(File.dirname(__FILE__), "data_migrate", "migration_context")
 require File.join(File.dirname(__FILE__), "data_migrate", "railtie")
 require File.join(File.dirname(__FILE__), "data_migrate", "tasks/data_migrate_tasks")
-require File.join(File.dirname(__FILE__), "data_migrate", "legacy_migrator")
 require File.join(File.dirname(__FILE__), "data_migrate", "config")
 require File.join(File.dirname(__FILE__), "data_migrate", "schema_migration")
+require File.join(File.dirname(__FILE__), "data_migrate", "database_configurations_wrapper")
 
 module DataMigrate
   def self.root

data/spec/data_migrate/data_migrator_spec.rb

@@ -3,7 +3,8 @@
 require "spec_helper"
 
 describe DataMigrate::DataMigrator do
-  let(:subject) { DataMigrate::DataMigrator }
+  let(:described_class) { DataMigrate::DataMigrator }
+
   let(:db_config) do
     {
       adapter: "sqlite3",
@@ -13,8 +14,8 @@ describe DataMigrate::DataMigrator do
 
   before do
     ActiveRecord::Base.establish_connection(db_config)
-    ::ActiveRecord::SchemaMigration.create_table
-    DataMigrate::DataSchemaMigration.create_table
+    DataMigrate::RailsHelper.schema_migration.create_table
+    DataMigrate::RailsHelper.data_schema_migration.create_table
   end
 
   after do
@@ -23,12 +24,14 @@ describe DataMigrate::DataMigrator do
   end
 
   describe ".load_migrated" do
+    let(:migrator) { DataMigrate::RailsHelper.data_migrator(:up, []) }
+
     it "loads migrated versions" do
-      DataMigrate::DataSchemaMigration.create(version: 20090000000000)
-      ::ActiveRecord::SchemaMigration.create(version: 20100000000000)
-      DataMigrate::DataSchemaMigration.create(version: 20110000000000)
-      ::ActiveRecord::SchemaMigration.create(version: 20120000000000)
-      migrated = subject.new(:up, []).load_migrated
+      DataMigrate::RailsHelper.data_schema_migration.create_version(20090000000000)
+      DataMigrate::RailsHelper.schema_create_version(20100000000000)
+      DataMigrate::RailsHelper.data_schema_migration.create_version(20110000000000)
+      DataMigrate::RailsHelper.schema_create_version(20120000000000)
+      migrated = migrator.load_migrated
       expect(migrated.count).to eq 2
       expect(migrated).to include 20090000000000
       expect(migrated).to include 20110000000000
@@ -38,7 +41,7 @@ describe DataMigrate::DataMigrator do
   describe :create_data_schema_table do
     it "creates the data_migrations table" do
       ActiveRecord::Migration.drop_table("data_migrations") rescue nil
-      subject.create_data_schema_table
+      described_class.create_data_schema_table
       expect(
         ActiveRecord::Base.connection.table_exists?("data_migrations")
       ).to eq true
@@ -47,7 +50,7 @@ describe DataMigrate::DataMigrator do
 
   describe "#migrations_status" do
     it "returns all migrations statuses" do
-      status = subject.migrations_status
+      status = described_class.migrations_status
       expect(status.length).to eq 2
       expect(status.first).to eq ["down", "20091231235959", "Some name"]
       expect(status.second).to eq ["down", "20171231235959", "Super update"]
@@ -57,19 +60,19 @@ describe DataMigrate::DataMigrator do
   describe :match do
     context "when the file does not match" do
       it "returns nil" do
-        expect(subject.match("not_a_data_migration_file")).to be_nil
+        expect(described_class.match("not_a_data_migration_file")).to be_nil
       end
     end
 
     context "when the file doesn't end in .rb" do
       it "returns nil" do
-        expect(subject.match("20091231235959_some_name.rb.un~")).to be_nil
+        expect(described_class.match("20091231235959_some_name.rb.un~")).to be_nil
       end
     end
 
     context "when the file matches" do
       it "returns a valid MatchData object" do
-        match_data = subject.match("20091231235959_some_name.rb")
+        match_data = described_class.match("20091231235959_some_name.rb")
 
         expect(match_data[0]).to eq "20091231235959_some_name.rb"
         expect(match_data[1]).to eq "20091231235959"

data/spec/data_migrate/data_schema_migration_spec.rb

@@ -1,16 +1,33 @@
 require 'spec_helper'
 
 describe DataMigrate::DataSchemaMigration do
-  let(:subject) { DataMigrate::DataSchemaMigration }
-  describe :table_name do
-    it "returns correct table name" do
-      expect(subject.table_name).to eq("data_migrations")
+  if DataMigrate::RailsHelper.rails_version_equal_to_or_higher_than_7_1
+    let(:connection) { double(:connection) }
+    let(:subject) { DataMigrate::DataSchemaMigration.new(connection) }
+
+    describe :table_name do
+      it "returns correct table name" do
+        expect(subject.table_name).to eq("data_migrations")
+      end
+    end
+
+    describe :index_name do
+      it "returns correct primary key name" do
+        expect(subject.primary_key).to eq("version")
+      end
+    end
+  else
+    let(:subject) { DataMigrate::DataSchemaMigration }
+    describe :table_name do
+      it "returns correct table name" do
+        expect(subject.table_name).to eq("data_migrations")
+      end
     end
-  end
 
-  describe :index_name do
-    it "returns correct primary key name" do
-      expect(subject.primary_key).to eq("version")
+    describe :index_name do
+      it "returns correct primary key name" do
+        expect(subject.primary_key).to eq("version")
+      end
     end
   end
 end

data/spec/data_migrate/data_spec.rb

@@ -55,7 +55,7 @@ describe DataMigrate::Data do
 
         sql_select = <<-SQL
           SELECT version
-          FROM #{DataMigrate::DataSchemaMigration.table_name}
+          FROM #{DataMigrate::RailsHelper.data_schema_migration.table_name}
         SQL
 
         db_list_data = ActiveRecord::Base.connection.