Fixed connection leasing for Active Record 7.2+

ankane · Oct 7, 2024 · 8fdf313 · 8fdf313
1 parent b9f316c
commit 8fdf313
Show file tree

Hide file tree

Showing 5 changed files with 27 additions and 11 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,3 +1,7 @@
+## 6.5.1 (unreleased)
+
+- Fixed connection leasing for Active Record 7.2+
+
 ## 6.5.0 (2024-10-01)
 
 - Added support for Active Record 8

diff --git a/lib/groupdate/adapters/base_adapter.rb b/lib/groupdate/adapters/base_adapter.rb
@@ -3,7 +3,7 @@ module Adapters
     class BaseAdapter
       attr_reader :period, :column, :day_start, :week_start, :n_seconds
 
-      def initialize(relation, column:, period:, time_zone:, time_range:, week_start:, day_start:, n_seconds:)
+      def initialize(relation, column:, period:, time_zone:, time_range:, week_start:, day_start:, n_seconds:, adapter_name: nil)
         @relation = relation
         @column = column
         @period = period
@@ -12,6 +12,7 @@ def initialize(relation, column:, period:, time_zone:, time_range:, week_start:,
         @week_start = week_start
         @day_start = day_start
         @n_seconds = n_seconds
+        @adapter_name = adapter_name
 
         if ActiveRecord::VERSION::MAJOR >= 7
           if ActiveRecord.default_timezone == :local

diff --git a/lib/groupdate/adapters/postgresql_adapter.rb b/lib/groupdate/adapters/postgresql_adapter.rb
@@ -23,7 +23,7 @@ def group_clause
           when :week
             ["(DATE_TRUNC('day', #{day_start_column} - INTERVAL '1 day' * ((? + EXTRACT(DOW FROM #{day_start_column})::integer) % 7)) + INTERVAL ?)::date", time_zone, day_start_interval, 13 - week_start, time_zone, day_start_interval, day_start_interval]
           when :custom
-            if @relation.connection.adapter_name == "Redshift"
+            if @adapter_name == "Redshift"
               ["TIMESTAMP 'epoch' + (FLOOR(EXTRACT(EPOCH FROM #{column}::timestamp) / ?) * ?) * INTERVAL '1 second'", n_seconds, n_seconds]
             else
               ["TO_TIMESTAMP(FLOOR(EXTRACT(EPOCH FROM #{column}::timestamptz) / ?) * ?)", n_seconds, n_seconds]

diff --git a/lib/groupdate/magic.rb b/lib/groupdate/magic.rb
@@ -180,12 +180,13 @@ def cast_result(result, multiple_groups)
       end
 
       def time_zone_support?(relation)
-        if relation.connection.adapter_name.match?(/mysql/i)
-          # need to call klass for Rails < 5.2
-          sql = relation.klass.send(:sanitize_sql_array, ["SELECT CONVERT_TZ(NOW(), '+00:00', ?)", time_zone.tzinfo.name])
-          !relation.connection.select_all(sql).to_a.first.values.first.nil?
-        else
-          true
+        relation.connection_pool.with_connection do |connection|
+          if connection.adapter_name.match?(/mysql|trilogy/i)
+            sql = relation.send(:sanitize_sql_array, ["SELECT CONVERT_TZ(NOW(), '+00:00', ?)", time_zone.tzinfo.name])
+            !connection.select_all(sql).to_a.first.values.first.nil?
+          else
+            true
+          end
         end
       end
 
@@ -203,7 +204,7 @@ def check_nils(result, multiple_groups, relation)
       def self.generate_relation(relation, field:, **options)
         magic = Groupdate::Magic::Relation.new(**options)
 
-        adapter_name = relation.connection.adapter_name
+        adapter_name = relation.connection_pool.with_connection { |c| c.adapter_name }
         adapter = Groupdate.adapters[adapter_name]
         raise Groupdate::Error, "Connection adapter not supported: #{adapter_name}" unless adapter
 
@@ -221,7 +222,8 @@ def self.generate_relation(relation, field:, **options)
             time_range: magic.time_range,
             week_start: magic.week_start,
             day_start: magic.day_start,
-            n_seconds: magic.n_seconds
+            n_seconds: magic.n_seconds,
+            adapter_name: adapter_name
           ).generate
 
         # add Groupdate info
@@ -251,7 +253,7 @@ def validate_column(column)
         def resolve_column(relation, column)
           node = relation.send(:relation).send(:arel_columns, [column]).first
           node = Arel::Nodes::SqlLiteral.new(node) if node.is_a?(String)
-          relation.connection.visitor.accept(node, Arel::Collectors::SQLString.new).value
+          relation.connection_pool.with_connection { |c| c.visitor.accept(node, Arel::Collectors::SQLString.new).value }
         end
       end
 

diff --git a/test/database_test.rb b/test/database_test.rb
@@ -236,6 +236,15 @@ def test_n_duration
     assert_equal({}, User.group_by_second(:created_at, n: 2.minutes).count)
   end
 
+  def test_connection_leasing
+    ActiveRecord::Base.connection_handler.clear_active_connections!
+    assert_nil ActiveRecord::Base.connection_pool.active_connection?
+    ActiveRecord::Base.connection_pool.with_connection do
+      User.group_by_day(:created_at).count
+    end
+    assert_nil ActiveRecord::Base.connection_pool.active_connection?
+  end
+
   private
 
   def this_year