Restore Rails edge compatibility

lib/chrono_model/adapter.rb

@@ -56,7 +56,7 @@ def chrono_setup!
       chrono_upgrade_warning
     end
 
-    # Runs primary_key, indexes and default_sequence_name in the
+    # Runs primary_keys, indexes and default_sequence_name in the
     # temporal schema, as the table there defined is the source for
     # this information.
     #
@@ -70,27 +70,21 @@ def chrono_setup!
     #
     # NOTE: These methods are dynamically defined, see the source.
     #
-    def primary_key(table_name); end
-
-    %i[primary_key indexes default_sequence_name].each do |method|
-      define_method(method) do |*args|
-        table_name = args.first
-        return super(*args) unless is_chrono?(table_name)
-
-        on_schema(TEMPORAL_SCHEMA, recurse: :ignore) { super(*args) }
+    %i[primary_keys indexes default_sequence_name].each do |method|
+      define_method(method) do |table_name, *args|
+        chrono_lookup(table_name, TEMPORAL_SCHEMA) { |name| super(name, *args) }
       end
     end
 
-    # Runs column_definitions in the temporal schema, as the table there
-    # defined is the source for this information.
+    # Runs columns in the temporal schema, as the table there defined is
+    # the source for this information.
     #
     # The default search path is included however, since the table
     # may reference types defined in other schemas, which result in their
     # names becoming schema qualified, which will cause type resolutions to fail.
-    def column_definitions(table_name)
-      return super unless is_chrono?(table_name)
-
-      on_schema("#{TEMPORAL_SCHEMA},#{schema_search_path}", recurse: :ignore) { super }
+    #
+    def columns(table_name)
+      chrono_lookup(table_name, "#{TEMPORAL_SCHEMA},#{schema_search_path}") { |name| super(name) }
     end
 
     # Evaluates the given block in the temporal schema.
@@ -131,10 +125,7 @@ def on_schema(schema, recurse: :follow)
       # there is no way to know which path will be restored when the
       # transaction ends.
       #
-      transaction_aborted =
-        chrono_connection.transaction_status == PG::Connection::PQTRANS_INERROR
-
-      if transaction_aborted && Thread.current['recursions'] == 1
+      if chrono_connection.transaction_status == PG::Connection::PQTRANS_INERROR
         @schema_search_path = nil
       else
         self.schema_search_path = old_path
@@ -180,6 +171,28 @@ def chrono_connection
       @chrono_connection ||= @raw_connection || @connection
     end
 
+    # Reads schema metadata through the given block, evaluated in the
+    # given schema for temporal tables and as-is for plain ones.
+    #
+    # Rails 8.2 schema readers accept an Array of table names too and
+    # return a Hash keyed by table name: temporal and plain tables are
+    # then read separately, in one query each, and merged back in the
+    # requested order.
+    #
+    def chrono_lookup(table_name, schema)
+      unless table_name.is_a?(Array)
+        return yield(table_name) unless is_chrono?(table_name)
+
+        return on_schema(schema, recurse: :ignore) { yield(table_name) }
+      end
+
+      chrono, plain = table_name.partition { |name| is_chrono?(name) }
+
+      result = yield(plain)
+      result.merge!(on_schema(schema, recurse: :ignore) { yield(chrono) }) if chrono.any?
+      result.slice(*table_name.map(&:to_s))
+    end
+
     # Counts the number of recursions in a thread local variable
     #
     def count_recursions # yield

spec/chrono_model/adapter/base_spec.rb

@@ -34,6 +34,68 @@
     end
   end
 
+  describe '.columns' do
+    subject { adapter.columns(table).find { |column| column.name == 'test' } }
+
+    context 'with temporal tables' do
+      include_context 'with temporal tables'
+
+      it { is_expected.to have_attributes(default: 'default-value', null: false) }
+    end
+
+    context 'with plain tables' do
+      include_context 'with plain tables'
+
+      it { is_expected.to have_attributes(default: 'default-value', null: false) }
+    end
+  end
+
+  describe '.primary_key' do
+    subject { adapter.primary_key(table) }
+
+    context 'with temporal tables' do
+      include_context 'with temporal tables'
+
+      it { is_expected.to eq 'id' }
+    end
+
+    context 'with plain tables' do
+      include_context 'with plain tables'
+
+      it { is_expected.to eq 'id' }
+    end
+  end
+
+  describe 'reading many tables at once', if: ActiveRecord::VERSION::STRING >= '8.2' do
+    include_context 'with temporal tables'
+
+    before do
+      adapter.create_table 'plain_table', &columns
+      adapter.add_index 'plain_table', :foo
+      adapter.on_temporal_schema { adapter.add_index table, :foo }
+    end
+
+    after { adapter.drop_table 'plain_table' }
+
+    let(:tables) { ['plain_table', table] }
+
+    it 'reads temporal tables in the temporal schema, in the requested order' do
+      defaults = adapter.columns(tables).transform_values { |cols| cols.to_h { |c| [c.name, c.default] } }
+
+      expect(defaults.keys).to eq tables
+      expect(defaults[table]).to eq(defaults['plain_table']).and include('test' => 'default-value')
+    end
+
+    it { expect(adapter.primary_keys(tables)).to eq('plain_table' => ['id'], table => ['id']) }
+
+    it {
+      expect(adapter.indexes(tables).transform_values { |indexes| indexes.map(&:name) })
+        .to eq('plain_table' => ['index_plain_table_on_foo'], table => ['index_test_table_on_foo'])
+    }
+
+    it { expect(adapter.columns([])).to eq({}) }
+  end
+
   describe '.on_schema' do
     subject(:on_schema) { adapter }
 
@@ -76,9 +138,18 @@
 
         it {
           expect { on_schema }
-            .to raise_error(/current transaction is aborted/)
-            .and(change { adapter.instance_variable_get(:@schema_search_path) })
+            .to raise_error(ActiveRecord::StatementInvalid, /syntax error/)
+            .and(change { adapter.instance_variable_get(:@schema_search_path) }.to(nil))
         }
+
+        it 'restores the search path after rollback' do
+          expect { on_schema }.to raise_error(ActiveRecord::StatementInvalid)
+
+          adapter.execute 'ROLLBACK'
+
+          expect(adapter.schema_search_path).to eq '"$user", public'
+          expect(adapter).to be_in_schema(:default)
+        end
       end
     end