From b75df790b184e7869ecb683e561e53e7a8289ecf Mon Sep 17 00:00:00 2001
From: Geremia Taglialatela <tagliala.dev@gmail.com>
Date: Thu, 15 Jun 2023 13:17:07 +0200
Subject: [PATCH] Register database tasks properly

Also use class for Rails < 5.2 and class name for Rails >= 5.2 to match
AR's behavior

Fix #151
---
 lib/chrono_model.rb         |  2 +-
 lib/chrono_model/railtie.rb | 28 ++++++++++++++++------------
 2 files changed, 17 insertions(+), 13 deletions(-)

diff --git a/lib/chrono_model.rb b/lib/chrono_model.rb
index 73923430..73562ede 100644
--- a/lib/chrono_model.rb
+++ b/lib/chrono_model.rb
@@ -31,7 +31,7 @@ def self.history_models
   end
 end
 
-if defined?(Rails)
+if defined?(Rails::Railtie)
   require 'chrono_model/railtie'
 end
 
diff --git a/lib/chrono_model/railtie.rb b/lib/chrono_model/railtie.rb
index 94b74bec..951555b2 100644
--- a/lib/chrono_model/railtie.rb
+++ b/lib/chrono_model/railtie.rb
@@ -1,6 +1,10 @@
+# frozen_string_literal: true
+
 require 'active_record/tasks/chronomodel_database_tasks'
+
 module ChronoModel
   class Railtie < ::Rails::Railtie
+    TASKS_CLASS = ActiveRecord::Tasks::ChronomodelDatabaseTasks
 
     def task_config
       if Rails.version < '6.1'
@@ -10,17 +14,18 @@ def task_config
       end
     end
 
+    # Register our database tasks under our adapter name
+    if Rails.version < '5.2'
+      ActiveRecord::Tasks::DatabaseTasks.register_task(/chronomodel/, TASKS_CLASS)
+    else
+      ActiveRecord::Tasks::DatabaseTasks.register_task(/chronomodel/, TASKS_CLASS.to_s)
+    end
+
     rake_tasks do
       if Rails.application.config.active_record.schema_format != :sql
         raise 'In order to use ChronoModel, config.active_record.schema_format must be :sql!'
       end
 
-      tasks_class = ActiveRecord::Tasks::ChronomodelDatabaseTasks
-
-      # Register our database tasks under our adapter name
-      #
-      ActiveRecord::Tasks::DatabaseTasks.register_task(/chronomodel/, tasks_class)
-
       if Rails.version < '6.1'
         # Make schema:dump and schema:load invoke structure:dump and structure:load
         Rake::Task['db:schema:dump'].clear.enhance(['environment']) do
@@ -32,19 +37,18 @@ def task_config
         end
       end
 
-      desc "Dumps database into db/data.NOW.sql or file specified via DUMP="
+      desc 'Dumps database into db/data.NOW.sql or file specified via DUMP='
       task 'db:data:dump' => :environment do
         target = ENV['DUMP'] || Rails.root.join('db', "data.#{Time.now.to_f}.sql")
-        tasks_class.new(task_config).data_dump(target)
+        TASKS_CLASS.new(task_config).data_dump(target)
       end
 
-      desc "Loads database dump from file specified via DUMP="
+      desc 'Loads database dump from file specified via DUMP='
       task 'db:data:load' => :environment do
         source = ENV['DUMP'].presence or
-          raise ArgumentError, "Invoke as rake db:data:load DUMP=/path/to/data.sql"
-        tasks_class.new(task_config).data_load(source)
+          raise ArgumentError, 'Invoke as rake db:data:load DUMP=/path/to/data.sql'
+        TASKS_CLASS.new(task_config).data_load(source)
       end
     end
-
   end
 end