Fix custom URL connection

pawurb · Nov 21, 2024 · 9df59fe · 9df59fe
1 parent f37b271
commit 9df59fe
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 2 deletions.
diff --git 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
@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module RailsPgExtras
-  VERSION = "5.4.3"
+  VERSION = "5.4.4"
 end