From 9df59fe400b5cbd4cd6ff2c8fc7bd9a7f53dd76c Mon Sep 17 00:00:00 2001
From: pawurb <contact@pawelurbanek.com>
Date: Thu, 21 Nov 2024 13:26:19 +0100
Subject: [PATCH] Fix custom URL connection

---
 lib/rails-pg-extras.rb         | 10 +++++++++-
 lib/rails_pg_extras/version.rb |  2 +-
 2 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/lib/rails-pg-extras.rb b/lib/rails-pg-extras.rb
index eeeb914..027ddf9 100644
--- a/lib/rails-pg-extras.rb
+++ b/lib/rails-pg-extras.rb
@@ -152,7 +152,15 @@ def self.table_info(args: {}, in_format: :display_table)
 
   def self.connection
     if (db_url = ENV["RAILS_PG_EXTRAS_DATABASE_URL"])
-      ActiveRecord::Base.establish_connection(db_url).connection
+      connector = ActiveRecord::Base.establish_connection(db_url)
+
+      if connector.respond_to?(:connection)
+        connector.connection
+      elsif connector.respond_to?(:lease_connection)
+        connector.lease_connection
+      else
+        raise "Unsupported connector: #{connector.class}"
+      end
     else
       ActiveRecord::Base.connection
     end
diff --git a/lib/rails_pg_extras/version.rb b/lib/rails_pg_extras/version.rb
index e57f7b5..3384f3f 100644
--- a/lib/rails_pg_extras/version.rb
+++ b/lib/rails_pg_extras/version.rb
@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module RailsPgExtras
-  VERSION = "5.4.3"
+  VERSION = "5.4.4"
 end