From 2cb4051bf11f2911d4298fa3406c0f66a0af7eb4 Mon Sep 17 00:00:00 2001
From: Geremia Taglialatela <tagliala.dev@gmail.com>
Date: Tue, 28 Nov 2023 17:07:44 +0100
Subject: [PATCH] Add Rails 7.2 compatibility

Also improves 7.1 compatibility

Close #252

Ref:
- rails/rails@009c7e7
- rails/rails#50064
---
 .../connection_adapters/chronomodel_adapter.rb    |  8 ++++----
 lib/chrono_model.rb                               |  4 ++++
 lib/chrono_model/adapter.rb                       | 15 +++++++++++++++
 spec/aruba/rake_task_spec.rb                      |  5 ++++-
 4 files changed, 27 insertions(+), 5 deletions(-)

diff --git a/lib/active_record/connection_adapters/chronomodel_adapter.rb b/lib/active_record/connection_adapters/chronomodel_adapter.rb
index 9626377a..1a71f127 100644
--- a/lib/active_record/connection_adapters/chronomodel_adapter.rb
+++ b/lib/active_record/connection_adapters/chronomodel_adapter.rb
@@ -3,9 +3,10 @@
 require 'chrono_model'
 
 module ActiveRecord
+  # TODO: Remove when dropping Rails < 7.2 compatibility
   module ConnectionHandling
     def chronomodel_adapter_class
-      ConnectionAdapters::PostgreSQLAdapter
+      ChronoModel::Adapter
     end
 
     # Install the new adapter in ActiveRecord. This approach is required because
@@ -13,6 +14,8 @@ def chronomodel_adapter_class
     # super() in overridden Module methods.
     #
     def chronomodel_connection(config) # :nodoc:
+      return chronomodel_adapter_class.new(config) if ActiveRecord::VERSION::STRING >= '7.1'
+
       conn_params = config.symbolize_keys
 
       conn_params.delete_if { |_, v| v.nil? }
@@ -29,9 +32,6 @@ def chronomodel_connection(config) # :nodoc:
 
       adapter = ChronoModel::Adapter.new(conn, logger, conn_params, config)
 
-      # Rails 7.2.0, see ifad/chronomodel#236
-      adapter.connect! if adapter.respond_to?(:connect!)
-
       unless adapter.chrono_supported?
         raise ChronoModel::Error, 'Your database server is not supported by ChronoModel. ' \
                                   'Currently, only PostgreSQL >= 9.3 is supported.'
diff --git a/lib/chrono_model.rb b/lib/chrono_model.rb
index 2a2a583c..3ed9101b 100644
--- a/lib/chrono_model.rb
+++ b/lib/chrono_model.rb
@@ -64,3 +64,7 @@ def self.history_models
 
   raise 'In order to use ChronoModel, set `config.active_record.schema_format` to `:sql`'
 end
+
+if ActiveRecord::ConnectionAdapters.respond_to?(:register)
+  ActiveRecord::ConnectionAdapters.register 'chronomodel', 'ChronoModel::Adapter', 'chrono_model/adapter'
+end
diff --git a/lib/chrono_model/adapter.rb b/lib/chrono_model/adapter.rb
index 519a9416..354db72f 100644
--- a/lib/chrono_model/adapter.rb
+++ b/lib/chrono_model/adapter.rb
@@ -33,6 +33,21 @@ class Adapter < ActiveRecord::ConnectionAdapters::PostgreSQLAdapter
     # The schema holding historical data
     HISTORY_SCHEMA  = 'history'
 
+    if ActiveRecord::VERSION::STRING >= '7.1'
+      def initialize(*)
+        super
+
+        connect!
+
+        unless chrono_supported?
+          raise ChronoModel::Error, 'Your database server is not supported by ChronoModel. ' \
+                                    'Currently, only PostgreSQL >= 9.3 is supported.'
+        end
+
+        chrono_setup!
+      end
+    end
+
     # Returns true whether the connection adapter supports our
     # implementation of temporal tables. Currently, Chronomodel
     # is supported starting with PostgreSQL 9.3 (90300 in PostgreSQL's
diff --git a/spec/aruba/rake_task_spec.rb b/spec/aruba/rake_task_spec.rb
index 8cf63519..034357d1 100644
--- a/spec/aruba/rake_task_spec.rb
+++ b/spec/aruba/rake_task_spec.rb
@@ -13,7 +13,10 @@
   describe 'bundle exec rake -T' do
     subject { last_command_started }
 
-    before { run_command_and_stop('bundle exec rake -T') }
+    before do
+      copy_db_config
+      run_command_and_stop('bundle exec rake -T')
+    end
 
     it { is_expected.to have_output(load_schema_task(as_regexp: true)) }
   end