Remove permanent connection checkouts

joshuay03 · joshuay03 · commit a417e60f6a75 · 2025-06-30T14:19:33.000+10:00
diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
@@ -6,31 +6,27 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        include:
-          - ruby: 3.4
-            gemfile: gemfiles/activerecord71.gemfile
-            postgres: 16
-          - ruby: 3.4
-            gemfile: gemfiles/activerecord71.gemfile
-            postgres: 15
-          - ruby: 3.4
-            gemfile: gemfiles/activerecord71.gemfile
-            postgres: 14
-          - ruby: 3.4
-            gemfile: gemfiles/activerecord71.gemfile
-            postgres: 13
-          - ruby: 3.3
-            gemfile: gemfiles/activerecord71.gemfile
-            postgres: 16
-          - ruby: 3.3
-            gemfile: gemfiles/activerecord71.gemfile
-            postgres: 15
-          - ruby: 3.3
-            gemfile: gemfiles/activerecord71.gemfile
-            postgres: 14
-          - ruby: 3.3
-            gemfile: gemfiles/activerecord71.gemfile
-            postgres: 13
+        ruby: ["2.7", "3.0", "3.1", "3.2", "3.3", "3.4", head]
+        postgres: [13, 14, 15, 16]
+        gemfile:
+          [
+            gemfiles/activerecord60.gemfile,
+            gemfiles/activerecord61.gemfile,
+            gemfiles/activerecord70.gemfile,
+            gemfiles/activerecord71.gemfile,
+            gemfiles/activerecord72.gemfile,
+            gemfiles/activerecord80.gemfile,
+          ]
+        exclude:
+          # Rails 7.2 requires Ruby 3.1 or higher
+          - ruby: "2.7"
+            gemfile: gemfiles/activerecord72.gemfile
+          - ruby: "2.7"
+            gemfile: gemfiles/activerecord80.gemfile
+          - ruby: "3.0"
+            gemfile: gemfiles/activerecord72.gemfile
+          - ruby: "3.0"
+            gemfile: gemfiles/activerecord80.gemfile
     env:
       BUNDLE_GEMFILE: ${{ matrix.gemfile }}
     steps:
diff --git a/Gemfile b/Gemfile
@@ -5,7 +5,7 @@ gemspec
 gem "minitest", ">= 5"
 gem "rake"
 
-gem "activerecord", "~> 7.0.0"
+gem "activerecord"
 gem "combustion"
 gem "pg"
 gem "pg_query"
diff --git a/gemfiles/activerecord60.gemfile b/gemfiles/activerecord60.gemfile
@@ -8,4 +8,4 @@ gem "rake"
 gem "activerecord", "~> 6.0.0"
 gem "combustion"
 gem "pg"
-gem "pg_query", "~> 2"
+gem "pg_query"
diff --git a/gemfiles/activerecord61.gemfile b/gemfiles/activerecord61.gemfile
@@ -8,4 +8,4 @@ gem "rake"
 gem "activerecord", "~> 6.1.0"
 gem "combustion"
 gem "pg"
-gem "pg_query", "~> 2"
+gem "pg_query"
diff --git a/gemfiles/activerecord70.gemfile b/gemfiles/activerecord70.gemfile
@@ -0,0 +1,11 @@
+source "https://rubygems.org"
+
+gemspec path: ".."
+
+gem "minitest", ">= 5"
+gem "rake"
+
+gem "activerecord", "~> 7.0.0"
+gem "combustion"
+gem "pg"
+gem "pg_query"
diff --git a/gemfiles/activerecord71.gemfile b/gemfiles/activerecord71.gemfile
@@ -5,7 +5,7 @@ gemspec path: ".."
 gem "minitest", ">= 5"
 gem "rake"
 
-gem "activerecord", github: "rails/rails"
+gem "activerecord", "~> 7.1.0"
 gem "combustion"
 gem "pg"
 gem "pg_query"
diff --git a/gemfiles/activerecord72.gemfile b/gemfiles/activerecord72.gemfile
@@ -0,0 +1,11 @@
+source "https://rubygems.org"
+
+gemspec path: ".."
+
+gem "minitest", ">= 5"
+gem "rake"
+
+gem "activerecord", "~> 7.2.0"
+gem "combustion"
+gem "pg"
+gem "pg_query"
diff --git a/gemfiles/activerecord80.gemfile b/gemfiles/activerecord80.gemfile
@@ -0,0 +1,11 @@
+source "https://rubygems.org"
+
+gemspec path: ".."
+
+gem "minitest", ">= 5"
+gem "rake"
+
+gem "activerecord", "~> 8.0.0"
+gem "combustion"
+gem "pg"
+gem "pg_query"
diff --git a/lib/pghero.rb b/lib/pghero.rb
@@ -64,7 +64,7 @@ class << self
       :query_stats_available?, :query_stats_enabled?, :query_stats_extension_enabled?, :query_stats_readable?,
       :rds_stats, :read_iops_stats, :aws_region, :relation_sizes, :replica?, :replication_lag, :replication_lag_stats,
       :reset_query_stats, :reset_stats, :running_queries, :aws_secret_access_key, :sequence_danger, :sequences, :settings,
-      :slow_queries, :space_growth, :ssl_used?, :stats_connection, :suggested_indexes, :suggested_indexes_by_query,
+      :slow_queries, :space_growth, :ssl_used?, :with_stats_connection, :suggested_indexes, :suggested_indexes_by_query,
       :suggested_indexes_enabled?, :system_stats_enabled?, :table_caching, :table_hit_rate, :table_stats,
       :total_connections, :transaction_id_danger, :unused_indexes, :unused_tables, :write_iops_stats
 
diff --git a/lib/pghero/database.rb b/lib/pghero/database.rb
@@ -126,7 +126,7 @@ def connection_model
         # rough check for Postgres adapter
         # keep this message generic so it's useful
         # when empty url set in Docker image pghero.yml
-        unless @connection_model.connection.adapter_name =~ /postg/i
+        unless @connection_model.adapter_class::ADAPTER_NAME =~ /postg/i
           raise Error, "Invalid connection URL"
         end
         @adapter_checked = true
diff --git a/lib/pghero/methods/basic.rb b/lib/pghero/methods/basic.rb
@@ -31,17 +31,18 @@ def server_version_num
       end
 
       def quote_ident(value)
-        connection.quote_column_name(value)
+        with_connection do |connection|
+          connection.quote_column_name(value)
+        end
       end
 
       private
 
       def select_all(sql, conn: nil, query_columns: [], cast_values: false)
-        conn ||= connection
         # squish for logs
         retries = 0
         begin
-          result = uncast_result = conn.select_all(add_source(squish(sql)))
+          result = uncast_result = with_connection { |conn| conn.select_all(add_source(squish(sql))) }
           if cast_values
             # ActiveRecord::Result#cast_values turns PostgreSQL arrays into Ruby arrays, etc.
             # But it turns ActiveRecord::Result into Array of results, which is why we keep
@@ -87,7 +88,9 @@ def select_all(sql, conn: nil, query_columns: [], cast_values: false)
       end
 
       def select_all_stats(sql, **options)
-        select_all(sql, **options, conn: stats_connection)
+        with_stats_connection do |stats_connection|
+          select_all(sql, **options, conn: stats_connection)
+        end
       end
 
       def select_all_size(sql)
@@ -103,19 +106,23 @@ def select_one(sql, conn: nil)
       end
 
       def select_one_stats(sql)
-        select_one(sql, conn: stats_connection)
+        with_stats_connection do |stats_connection|
+          select_one(sql, conn: stats_connection)
+        end
       end
 
       def execute(sql)
-        connection.execute(add_source(sql))
+        with_connection do |connection|
+          connection.execute(add_source(sql))
+        end
       end
 
-      def connection
-        connection_model.connection
+      def with_connection(&block)
+        connection_model.connection_pool.with_connection(&block)
       end
 
-      def stats_connection
-        ::PgHero::Stats.connection
+      def with_stats_connection(&block)
+        ::PgHero::Stats.connection_pool.with_connection(&block)
       end
 
       def squish(str)
@@ -127,15 +134,21 @@ def add_source(sql)
       end
 
       def quote(value)
-        connection.quote(value)
+        with_connection do |connection|
+          connection.quote(value)
+        end
       end
 
       def quote_table_name(value)
-        connection.quote_table_name(value)
+        with_connection do |connection|
+          connection.quote_table_name(value)
+        end
       end
 
       def quote_column_name(value)
-        connection.quote_column_name(value)
+        with_connection do |connection|
+          connection.quote_column_name(value)
+        end
       end
 
       def unquote(part)
@@ -156,7 +169,9 @@ def with_transaction(lock_timeout: nil, statement_timeout: nil, rollback: false)
       end
 
       def table_exists?(table)
-        stats_connection.table_exists?(table)
+        with_stats_connection do |stats_connection|
+          stats_connection.table_exists?(table)
+        end
       end
     end
   end
diff --git a/lib/pghero/methods/suggested_indexes.rb b/lib/pghero/methods/suggested_indexes.rb
@@ -79,7 +79,9 @@ def autoindex(create: false)
         suggested_indexes.each do |index|
           p index
           if create
-            connection.execute("CREATE INDEX CONCURRENTLY ON #{quote_table_name(index[:table])} (#{index[:columns].map { |c| quote_column_name(c) }.join(",")})")
+            with_connection do |connection|
+              connection.execute("CREATE INDEX CONCURRENTLY ON #{quote_table_name(index[:table])} (#{index[:columns].map { |c| quote_column_name(c) }.join(",")})")
+            end
           end
         end
       end
diff --git a/test/queries_test.rb b/test/queries_test.rb
@@ -8,7 +8,7 @@ def test_running_queries
   def test_filter_data
     query = "SELECT pg_sleep(1)"
     # TODO manually checkout connection if needed
-    t = Thread.new { ActiveRecord::Base.connection.execute(query) }
+    t = Thread.new { ActiveRecord::Base.lease_connection.execute(query) }
     sleep(0.5)
 
     assert_equal query, database.running_queries.first[:query]
diff --git a/test/query_stats_test.rb b/test/query_stats_test.rb
@@ -40,7 +40,7 @@ def test_reset_instance_query_stats_database
     skip unless gte12?
 
     assert database.reset_query_stats
-    ActiveRecord::Base.connection.select_all("SELECT 1")
+    ActiveRecord::Base.lease_connection.select_all("SELECT 1")
     assert database.query_stats.any? { |qs| qs[:query] == "SELECT $1" }
 
     assert database.reset_instance_query_stats(database: database.database_name)
@@ -62,7 +62,7 @@ def test_reset_query_stats_user
     skip unless gte12?
 
     assert database.reset_query_stats
-    ActiveRecord::Base.connection.select_all("SELECT 1")
+    ActiveRecord::Base.lease_connection.select_all("SELECT 1")
     assert database.query_stats.any? { |qs| qs[:query] == "SELECT $1" }
 
     assert database.reset_query_stats(user: database.current_user)
@@ -84,8 +84,8 @@ def test_reset_query_stats_query_hash
     skip unless gte12?
 
     assert database.reset_query_stats
-    ActiveRecord::Base.connection.select_all("SELECT 1")
-    ActiveRecord::Base.connection.select_all("SELECT 1 + 1")
+    ActiveRecord::Base.lease_connection.select_all("SELECT 1")
+    ActiveRecord::Base.lease_connection.select_all("SELECT 1 + 1")
 
     assert database.query_stats.any? { |qs| qs[:query] == "SELECT $1" }
     assert database.query_stats.any? { |qs| qs[:query] == "SELECT $1 + $2" }
diff --git a/test/test_helper.rb b/test/test_helper.rb
@@ -31,6 +31,10 @@ def with_explain_timeout(value)
   config.active_record.logger = logger
 end
 
+if ActiveRecord.respond_to?(:permanent_connection_checkout)
+  ActiveRecord.permanent_connection_checkout = :disallowed
+end
+
 class City < ActiveRecord::Base
 end
 
@@ -47,7 +51,7 @@ class User < ActiveRecord::Base
     }
   end
 State.insert_all!(states)
-ActiveRecord::Base.connection.execute("ANALYZE states")
+ActiveRecord::Base.lease_connection.execute("ANALYZE states")
 
 users =
   5000.times.map do |i|
@@ -67,4 +71,4 @@ class User < ActiveRecord::Base
     }
   end
 User.insert_all!(users)
-ActiveRecord::Base.connection.execute("ANALYZE users")
+ActiveRecord::Base.lease_connection.execute("ANALYZE users")
