Add multiple database support

app/controllers/rails_pg_extras/web/queries_controller.rb

@@ -4,6 +4,10 @@ class QueriesController < RailsPgExtras::Web::ApplicationController
     helper_method :unavailable_extensions
 
     def index
+      if params[:database].present?
+        RailsPgExtras.configuration.selected_database = params[:database].downcase.to_sym
+      end
+
       if params[:query_name].present?
         @query_name = params[:query_name].to_sym.presence_in(@all_queries.keys)
         return unless @query_name
@@ -35,10 +39,12 @@ def query_disabled?(query_name)
     end
 
     def unavailable_extensions
-      return @unavailable_extensions if defined?(@unavailable_extensions)
+      RailsPgExtras.with_selected_database do
+        return @unavailable_extensions if defined?(@unavailable_extensions)
 
-      enabled_extensions = ActiveRecord::Base.connection.extensions
-      @unavailable_extensions = REQUIRED_EXTENSIONS.delete_if { |ext| ext.to_s.in?(enabled_extensions) }
+        enabled_extensions = ExtrasDatabaseAdapter.connection.extensions
+        @unavailable_extensions = REQUIRED_EXTENSIONS.delete_if { |ext| ext.to_s.in?(enabled_extensions) }
+      end
     end
   end
 end

app/views/rails_pg_extras/web/shared/_queries_selector.html.erb

@@ -1,7 +1,12 @@
-<%= form_tag queries_path, id: "queries", method: :get do |f| %>
-  <%= select_tag :query_name, options_for_select(@all_queries, params[:query_name]),
-                 { prompt: "--- select query ---", class: "border p-2 font-bold", autofocus: true } %>
-<% end %>
+<%= form_tag queries_path, id: "queries", method: :get do |f| %>
+  <% if ENV["RAILS_PG_EXTRAS_DATABASE_URL"].blank? && RailsPgExtras.configuration.show_all_active_record_databases %>
+    <%= select_tag :database, options_for_select(ExtrasDatabaseAdapter.database_list, params[:database] || ExtrasDatabaseAdapter.database_list[0]),
+                   { prompt: "--- select database ---", class: "border p-2 font-bold", autofocus: true } %>
+  <% end %>
+  <%= select_tag :query_name, options_for_select(@all_queries, params[:query_name]),
+                 { prompt: "--- select query ---", class: "border p-2 font-bold", autofocus: true } %>
+
+<% end %>
 
 <%= javascript_tag nonce: true do -%>
   document.getElementById('queries').addEventListener('change', (e) => {

lib/extras_database_adapter.rb

@@ -0,0 +1,20 @@
+class ExtrasDatabaseAdapter < ActiveRecord::Base
+  self.abstract_class = true
+
+  def self.database_list
+    base_configurations = ActiveRecord::Base.configurations.configurations.filter{|t| t.env_name == ActiveRecord::ConnectionHandling::DEFAULT_ENV.call.to_s}
+    base_configurations.map{|config| config.name}
+  end
+
+  def self.connect_databases
+    database_config = {}
+    # create a connects_to database hash with the roles being the name of the database config
+    # e.g. what will be called: connects_to database {primary: primary, primary_replica: primary_replica}
+    self.database_list.each do |database_name|
+      database_config[database_name.to_sym] = database_name.to_sym
+    end
+
+    ExtrasDatabaseAdapter.connects_to database: database_config
+    self
+  end
+end

lib/rails-pg-extras.rb

@@ -8,6 +8,7 @@
 require "rails_pg_extras/index_info_print"
 require "rails_pg_extras/table_info"
 require "rails_pg_extras/table_info_print"
+require "extras_database_adapter"
 
 module RailsPgExtras
   QUERIES = RubyPgExtras::QUERIES
@@ -26,44 +27,48 @@ module RailsPgExtras
   end
 
   def self.run_query(query_name:, in_format:, args: {})
-    if %i(calls outliers).include?(query_name)
-      pg_stat_statements_version_sql = "SELECT installed_version
-                                        FROM pg_available_extensions
-                                        WHERE name = 'pg_stat_statements'"
-      if (version = RailsPgExtras.connection.execute(pg_stat_statements_version_sql)
-        .to_a[0].fetch("installed_version", nil))
-        if Gem::Version.new(version) < Gem::Version.new(NEW_PG_STAT_STATEMENTS)
-          query_name = "#{query_name}_legacy".to_sym
-        elsif Gem::Version.new(version) >= Gem::Version.new(PG_STAT_STATEMENTS_17)
-          query_name = "#{query_name}_17".to_sym
+    with_selected_database do
+      if %i(calls outliers).include?(query_name)
+        pg_stat_statements_version_sql = "SELECT installed_version
+                                          FROM pg_available_extensions
+                                          WHERE name = 'pg_stat_statements'"
+        if (version = RailsPgExtras.connection.execute(pg_stat_statements_version_sql)
+          .to_a[0].fetch("installed_version", nil))
+          if Gem::Version.new(version) < Gem::Version.new(NEW_PG_STAT_STATEMENTS)
+            query_name = "#{query_name}_legacy".to_sym
+          elsif Gem::Version.new(version) >= Gem::Version.new(PG_STAT_STATEMENTS_17)
+            query_name = "#{query_name}_17".to_sym
+          end
         end
       end
-    end
 
-    sql = if (custom_args = DEFAULT_ARGS[query_name].merge(args)) != {}
-        RubyPgExtras.sql_for(query_name: query_name) % custom_args
-      else
-        RubyPgExtras.sql_for(query_name: query_name)
-      end
+      sql = if (custom_args = DEFAULT_ARGS[query_name].merge(args)) != {}
+          RubyPgExtras.sql_for(query_name: query_name) % custom_args
+        else
+          RubyPgExtras.sql_for(query_name: query_name)
+        end
 
-    result = connection.execute(sql)
+      result = connection.execute(sql)
 
-    RubyPgExtras.display_result(
-      result,
-      title: RubyPgExtras.description_for(query_name: query_name),
-      in_format: in_format,
-    )
+      RubyPgExtras.display_result(
+        result,
+        title: RubyPgExtras.description_for(query_name: query_name),
+        in_format: in_format,
+      )
+    end
   end
 
   def self.diagnose(in_format: :display_table)
-    data = RailsPgExtras::DiagnoseData.call
+    with_selected_database do
+      data = RailsPgExtras::DiagnoseData.call
 
-    if in_format == :display_table
-      RailsPgExtras::DiagnosePrint.call(data)
-    elsif in_format == :hash
-      data
-    else
-      raise "Invalid 'in_format' argument!"
+      if in_format == :display_table
+        RailsPgExtras::DiagnosePrint.call(data)
+      elsif in_format == :hash
+        data
+      else
+        raise "Invalid 'in_format' argument!"
+      end
     end
   end
 
@@ -150,19 +155,30 @@ def self.table_info(args: {}, in_format: :display_table)
     end
   end
 
-  def self.connection
-    if (db_url = ENV["RAILS_PG_EXTRAS_DATABASE_URL"])
-      connector = ActiveRecord::Base.establish_connection(db_url)
+  def self.with_selected_database
+    ExtrasDatabaseAdapter.connect_databases
 
-      if connector.respond_to?(:connection)
-        connector.connection
-      elsif connector.respond_to?(:lease_connection)
-        connector.lease_connection
-      else
-        raise "Unsupported connector: #{connector.class}"
+    if (db_url = ENV["RAILS_PG_EXTRAS_DATABASE_URL"] || ENV["DATABASE_URL"])
+      ExtrasDatabaseAdapter.establish_connection(db_url)
+      yield
+    elsif RailsPgExtras.configuration.selected_database.present?
+      ExtrasDatabaseAdapter.connected_to(role: RailsPgExtras.configuration.selected_database) do
+        yield
+      end
+    else
+      ExtrasDatabaseAdapter.connected_to(role: ExtrasDatabaseAdapter.database_list.first.to_sym) do
+        yield
       end
+    end
+  end
+
+  def self.connection
+    if ExtrasDatabaseAdapter.respond_to?(:connection)
+      ExtrasDatabaseAdapter.connection
+    elsif ExtrasDatabaseAdapter.respond_to?(:lease_connection)
+      ExtrasDatabaseAdapter.lease_connection
     else
-      ActiveRecord::Base.connection
+      raise "Unsupported connector: #{ExtrasDatabaseAdapter.class}"
     end
   end
 end

lib/rails_pg_extras/configuration.rb

@@ -8,6 +8,8 @@ class Configuration
 
     attr_reader :enabled_web_actions
     attr_accessor :public_dashboard
+    attr_accessor :selected_database
+    attr_accessor :show_all_active_record_databases
 
     def initialize(attrs)
       self.enabled_web_actions = attrs[:enabled_web_actions]

spec/spec_helper.rb

@@ -3,6 +3,7 @@
 require "rubygems"
 require "bundler/setup"
 require "active_record"
+require_relative "../lib/extras_database_adapter"
 require_relative "../lib/rails-pg-extras"
 
 pg_version = ENV["PG_VERSION"]
@@ -30,9 +31,11 @@
     ActiveRecord::Base.establish_connection(
       ENV.fetch("DATABASE_URL")
     )
-    RailsPgExtras.connection.execute("CREATE EXTENSION IF NOT EXISTS pg_stat_statements;")
-    RailsPgExtras.connection.execute("CREATE EXTENSION IF NOT EXISTS pg_buffercache;")
-    RailsPgExtras.connection.execute("CREATE EXTENSION IF NOT EXISTS sslinfo;")
+    RailsPgExtras.with_selected_database do
+      RailsPgExtras.connection.execute("CREATE EXTENSION IF NOT EXISTS pg_stat_statements;")
+      RailsPgExtras.connection.execute("CREATE EXTENSION IF NOT EXISTS pg_buffercache;")
+      RailsPgExtras.connection.execute("CREATE EXTENSION IF NOT EXISTS sslinfo;")
+    end
   end
 
   config.after :suite do