From 796091024cd6fc01f274311e0d9e7185f450663a Mon Sep 17 00:00:00 2001
From: Andrew Kane <andrew@ankane.org>
Date: Mon, 7 Oct 2024 15:09:30 -0700
Subject: [PATCH] Replaced connection with with_connection

---
 lib/pghero/database.rb                  |  2 +-
 lib/pghero/methods/basic.rb             | 32 +++++++++++++------------
 lib/pghero/methods/suggested_indexes.rb |  4 +++-
 3 files changed, 21 insertions(+), 17 deletions(-)

diff --git a/lib/pghero/database.rb b/lib/pghero/database.rb
index 2b70dd85..ebe28d27 100644
--- 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.match?(/postg/i)
+        unless @connection_model.connection_db_config.adapter.to_s.match?(/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
index 0b3b1d63..6b94fe59 100644
--- a/lib/pghero/methods/basic.rb
+++ b/lib/pghero/methods/basic.rb
@@ -31,13 +31,18 @@ def server_version_num
       end
 
       def quote_ident(value)
-        connection.quote_column_name(value)
+        with_connection { |c| c.quote_column_name(value) }
       end
 
       private
 
-      def select_all(sql, conn: nil, query_columns: [])
-        conn ||= connection
+      def select_all(sql, stats: false, query_columns: [])
+        with_connection(stats: stats) do |conn|
+          select_all_leased(sql, conn: conn, query_columns: query_columns)
+        end
+      end
+
+      def select_all_leased(sql, conn:, query_columns:)
         # squish for logs
         retries = 0
         begin
@@ -81,7 +86,7 @@ def select_all(sql, conn: nil, query_columns: [])
       end
 
       def select_all_stats(sql, **options)
-        select_all(sql, **options, conn: stats_connection)
+        select_all(sql, **options, stats: true)
       end
 
       def select_all_size(sql)
@@ -97,15 +102,12 @@ def select_one(sql)
       end
 
       def execute(sql)
-        connection.execute(add_source(sql))
-      end
-
-      def connection
-        connection_model.connection
+        with_connection { |c| c.execute(add_source(sql)) }
       end
 
-      def stats_connection
-        ::PgHero::Stats.connection
+      def with_connection(stats: false, &block)
+        model = stats ? ::PgHero::Stats : connection_model
+        model.connection_pool.with_connection(&block)
       end
 
       def squish(str)
@@ -117,15 +119,15 @@ def add_source(sql)
       end
 
       def quote(value)
-        connection.quote(value)
+        with_connection { |c| c.quote(value) }
       end
 
       def quote_table_name(value)
-        connection.quote_table_name(value)
+        with_connection { |c| c.quote_table_name(value) }
       end
 
       def quote_column_name(value)
-        connection.quote_column_name(value)
+        with_connection { |c| c.quote_column_name(value) }
       end
 
       def unquote(part)
@@ -146,7 +148,7 @@ def with_transaction(lock_timeout: nil, statement_timeout: nil, rollback: false)
       end
 
       def table_exists?(table)
-        stats_connection.table_exists?(table)
+        with_connection(stats: true) { |c| c.table_exists?(table) }
       end
     end
   end
diff --git a/lib/pghero/methods/suggested_indexes.rb b/lib/pghero/methods/suggested_indexes.rb
index 9d17560e..1fd7e23c 100644
--- 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