pgbouncerhero 2.0.0 → 3.1.0

data/lib/pgbouncerhero/connection.rb

@@ -1,6 +1,5 @@
 module PgBouncerHero
   class Connection
-
     def initialize(host, port, user, password, dbname)
       @host = host
       @port = port
@@ -11,22 +10,60 @@ module PgBouncerHero
     end
 
     def connection
-      @connection ||= begin
-         begin
-          PG.connect(
-            host: @host,
-            port: @port,
-            user: @user,
-            password: @password,
-            dbname: @dbname,
-            connect_timeout: @timeout
-          )
-        rescue Exception => e
-          Rails.logger.error("[PGBouncerHero] Host:#{@host} | Database Name:#{@dbname} | Timeout: #{@timeout}s => #{e}")
-          nil
-        end
-      end
+      disconnect! if @connection && connection_invalid?(@connection)
+      @connection ||= connect
+    rescue StandardError => e
+      Rails.logger.error("[PGBouncerHero] Host:#{@host} | Database Name:#{@dbname} | Timeout: #{@timeout}s => #{e}")
+      nil
+    end
+
+    def connected?
+      !connection.nil?
+    end
+
+    def with_connection
+      raw_connection = connection
+      return unless raw_connection
+
+      yield raw_connection
+    rescue PG::Error
+      disconnect!
+      raise
+    end
+
+    def disconnect!
+      @connection&.finish unless @connection&.finished?
+    rescue PG::Error
+      nil
+    ensure
+      @connection = nil
     end
 
+    def method_missing(method_name, *, **, &)
+      with_connection { |raw_connection| raw_connection.public_send(method_name, *, **, &) }
+    end
+
+    def respond_to_missing?(method_name, include_private = false)
+      PG::Connection.public_instance_methods(include_private).include?(method_name) || super
+    end
+
+    private
+
+    def connect
+      PG.connect(
+        host: @host,
+        port: @port,
+        user: @user,
+        password: @password,
+        dbname: @dbname,
+        connect_timeout: @timeout
+      )
+    end
+
+    def connection_invalid?(connection)
+      connection.finished? || connection.status != PG::CONNECTION_OK
+    rescue PG::Error
+      true
+    end
   end
 end

data/lib/pgbouncerhero/database.rb

@@ -1,15 +1,23 @@
+require "connection_pool"
+require "monitor"
+
 module PgBouncerHero
   class Database
-
     include Methods::Basics
 
-    attr_reader :id, :config, :group
+    DEFAULT_POOL_SIZE = 5
+    DEFAULT_POOL_TIMEOUT = 5
+
+    attr_reader :id, :config, :group, :pool_size, :pool_timeout
 
     def initialize(group, id, config)
       @id = id
       @config = config || {}
-      @url = URI.parse(config["url"])
+      @url = URI.parse(@config["url"].to_s)
       @group = group
+      @pool_size = positive_integer_setting("pool_size", "PGBOUNCERHERO_POOL_SIZE", DEFAULT_POOL_SIZE)
+      @pool_timeout = positive_float_setting("pool_timeout", "PGBOUNCERHERO_POOL_TIMEOUT", DEFAULT_POOL_TIMEOUT)
+      @pool_monitor = Monitor.new
     end
 
     def name
@@ -17,7 +25,28 @@ module PgBouncerHero
     end
 
     def connection
-      @connection ||= connection_model.new(host, port, user, password, dbname).connection
+      proxy = connection_proxy
+      proxy if with_connection { true }
+    end
+
+    def with_connection
+      pool = connection_pool
+      pool.with do |managed_connection|
+        managed_connection.with_connection { |raw_connection| yield raw_connection }
+      end
+    rescue ConnectionPool::TimeoutError => e
+      Rails.logger.error("[PGBouncerHero] #{name} connection pool timed out after #{pool_timeout}s: #{e.message}")
+      nil
+    end
+
+    def disconnect!
+      pool = @pool_monitor.synchronize do
+        current_pool = @connection_pool
+        @connection_pool = nil
+        @connection_proxy = nil
+        current_pool
+      end
+      pool&.shutdown(&:disconnect!)
     end
 
     def host
@@ -42,6 +71,26 @@ module PgBouncerHero
 
     private
 
+    def execute(command)
+      with_connection { |raw_connection| raw_connection.exec(command) }
+    end
+
+    def connection_pool
+      @pool_monitor.synchronize do
+        @connection_pool ||= ConnectionPool.new(size: pool_size, timeout: pool_timeout) { build_connection }
+      end
+    end
+
+    def connection_proxy
+      @pool_monitor.synchronize do
+        @connection_proxy ||= ConnectionPool::Wrapper.new(pool: connection_pool)
+      end
+    end
+
+    def build_connection
+      connection_model.new(host, port, user, password, dbname)
+    end
+
     def connection_model
       @connection_model ||= begin
         Class.new(PgBouncerHero::Connection) do
@@ -51,5 +100,21 @@ module PgBouncerHero
         end
       end
     end
+
+    def positive_integer_setting(config_key, environment_key, default)
+      value = config.fetch(config_key, ENV.fetch(environment_key, default))
+      parsed = Integer(value)
+      raise ArgumentError, "#{config_key} must be greater than zero" unless parsed.positive?
+
+      parsed
+    end
+
+    def positive_float_setting(config_key, environment_key, default)
+      value = config.fetch(config_key, ENV.fetch(environment_key, default))
+      parsed = Float(value)
+      raise ArgumentError, "#{config_key} must be a finite number greater than zero" unless parsed.positive? && parsed.finite?
+
+      parsed
+    end
   end
 end

data/lib/pgbouncerhero/engine.rb

@@ -2,18 +2,21 @@ module PgBouncerHero
   class Engine < ::Rails::Engine
     isolate_namespace PgBouncerHero
 
-    initializer "pgbouncerhero", group: :all do |app|
-      if defined?(Sprockets) && Sprockets::VERSION >= "4"
-        app.config.assets.precompile << "pgbouncerhero/application.js"
-        app.config.assets.precompile << "pgbouncerhero/application.css"
-        app.config.assets.precompile << "pgbouncerhero/short-paragraph.png"
-      else
-        # use a proc instead of a string
-        app.config.assets.precompile << proc { |path| path == "pgbouncerhero/application.js" }
-        app.config.assets.precompile << proc { |path| path == "pgbouncerhero/application.css" }
-        app.config.assets.precompile << proc { |path| path == "pgbouncerhero/short-paragraph.png" }
+    initializer "pgbouncerhero.assets" do |app|
+      app.config.assets.paths << root.join("app/assets/builds")
+      app.config.assets.paths << root.join("app/javascript")
+    end
+
+    initializer "pgbouncerhero.importmap", before: "importmap" do |app|
+      PgBouncerHero.importmap.draw root.join("config/importmap.rb")
+      PgBouncerHero.importmap.cache_sweeper watches: root.join("app/javascript")
+
+      ActiveSupport.on_load(:action_controller_base) do
+        before_action { PgBouncerHero.importmap.cache_sweeper.execute_if_updated }
       end
+    end
 
+    initializer "pgbouncerhero.configuration" do
       PgBouncerHero.time_zone = PgBouncerHero.config["time_zone"] if PgBouncerHero.config["time_zone"]
     end
   end

data/lib/pgbouncerhero/group.rb

@@ -2,7 +2,6 @@ require "pgbouncerhero/database"
 
 module PgBouncerHero
   class Group
-
     attr_reader :id, :config, :databases
 
     def initialize(id, config)
@@ -16,6 +15,5 @@ module PgBouncerHero
     def name
       @name ||= id.to_s
     end
-
   end
 end

data/lib/pgbouncerhero/methods/basics.rb

@@ -1,42 +1,58 @@
 module PgBouncerHero
   module Methods
     module Basics
+      SHUTDOWN_MODES = {
+        nil => "",
+        immediate: "",
+        wait_for_clients: " WAIT_FOR_CLIENTS",
+        wait_for_servers: " WAIT_FOR_SERVERS"
+      }.freeze
+
       def summary
         if connection
           l = lists
           d = databases
           l = l.as_json
-          d = d.as_json.reject { |a| a['name'] == 'pgbouncer' }
-          l.push({databases_details: d})
+          d = d.as_json.reject { |a| a["name"] == "pgbouncer" }
+          l.push({ databases_details: d })
           l
         end
       end
       def databases
-        connection.exec("SHOW databases")
+        execute("SHOW databases")
       end
       def stats
-        connection.exec("SHOW stats")
+        execute("SHOW stats")
       end
       def lists
-        connection.exec("SHOW lists")
+        execute("SHOW lists")
       end
       def pools
-        connection.exec("SHOW pools")
+        execute("SHOW pools")
       end
       def clients
-        connection.exec("SHOW clients")
+        execute("SHOW clients")
       end
       def conf
-        connection.exec("SHOW config")
+        execute("SHOW config")
+      end
+      def state
+        execute("SHOW state")
       end
       def reload
-        connection.exec("RELOAD")
+        execute("RELOAD")
       end
       def suspend
-        connection.exec("SUSPEND")
+        execute("SUSPEND")
+      end
+      def resume
+        execute("RESUME")
       end
-      def shutdown
-        connection.exec("SHUTDOWN")
+      def shutdown(mode = nil)
+        suffix = SHUTDOWN_MODES.fetch(mode)
+        execute("SHUTDOWN#{suffix}")
+      rescue KeyError
+        raise ArgumentError, "unsupported shutdown mode: #{mode.inspect}"
       end
     end
   end

data/lib/pgbouncerhero/version.rb

@@ -1,3 +1,3 @@
 module PgBouncerHero
-  VERSION = "2.0.0"
+  VERSION = "3.1.0"
 end

data/lib/pgbouncerhero.rb

@@ -1,27 +1,43 @@
-require "pgbouncerhero/version"
+require "erb"
+require "monitor"
+require "pathname"
+require "yaml"
+require "rails"
+require "pg"
+require "importmap-rails"
+require "turbo-rails"
+require "stimulus-rails"
 
-# methods
+require "pgbouncerhero/version"
 require "pgbouncerhero/methods/basics"
-
+require "pgbouncerhero/connection"
 require "pgbouncerhero/group"
-require "pgbouncerhero/engine" if defined?(Rails)
+require "pgbouncerhero/engine"
 
-# models
-require "pgbouncerhero/connection"
+module PgBouncerHero
+  class ConfigurationError < StandardError; end
 
-# others
-require "jquery-rails"
-require "semantic-ui-sass"
+  BOOLEAN_SETTINGS = {
+    true => true,
+    false => false,
+    "1" => true,
+    "0" => false,
+    "true" => true,
+    "false" => false,
+    "yes" => true,
+    "no" => false,
+    "on" => true,
+    "off" => false
+  }.freeze
+  private_constant :BOOLEAN_SETTINGS
 
-module PgBouncerHero
-  # settings
-  class << self
-    attr_accessor :env
-  end
-  self.env = ENV["RAILS_ENV"] || ENV["RACK_ENV"] || "development"
+  STATE_MONITOR = Monitor.new
+  private_constant :STATE_MONITOR
+
+  mattr_accessor :importmap, default: Importmap::Map.new
 
   class << self
-    extend Forwardable
+    attr_reader :env
 
     def time_zone=(time_zone)
       @time_zone = time_zone.is_a?(ActiveSupport::TimeZone) ? time_zone : ActiveSupport::TimeZone[time_zone.to_s]
@@ -31,38 +47,105 @@ module PgBouncerHero
       @time_zone || Time.zone
     end
 
+    def env=(env)
+      STATE_MONITOR.synchronize do
+        @env = env.to_s
+        reset!
+      end
+    end
+
+    def config_path
+      STATE_MONITOR.synchronize { @config_path || default_config_path }
+    end
+
+    def config_path=(path)
+      STATE_MONITOR.synchronize do
+        @config_path = path && Pathname(path)
+        reset!
+      end
+    end
+
     def config
-      Thread.current[:PgBouncerHero_config] ||= begin
-        path = "config/pgbouncerhero.yml"
-
-        config = YAML.load(ERB.new(File.read(path)).result) if File.exist?(path)
-        config ||= {}
-
-        if config[env]
-          config[env]
-        elsif config["pgbouncers"] # preferred format
-          config
-        else
-          {
-            "pgbouncers" => {
-              "default" => {
-                "primary" => {
-                  "url" => ENV["PGBOUNCERHERO_DATABASE_URL"]
-                }
-              }
-            }
-          }
+      STATE_MONITOR.synchronize { @config ||= selected_config(load_config) }
+    end
+
+    def groups
+      STATE_MONITOR.synchronize do
+        @groups ||= config.fetch("pgbouncers").to_h do |group_id, _databases|
+          [ group_id.parameterize, PgBouncerHero::Group.new(group_id, config.fetch("pgbouncers")) ]
         end
       end
     end
 
-    def groups
-      @groups ||= begin
-        mapped = config['pgbouncers'].map do |group_id, hash|
-          [group_id.parameterize, PgBouncerHero::Group.new(group_id, config['pgbouncers'])]
+    def read_only?
+      value = ENV.fetch("PGBOUNCERHERO_READ_ONLY") { config.fetch("read_only", false) }
+      BOOLEAN_SETTINGS.fetch(value.is_a?(String) ? value.downcase : value)
+    rescue KeyError
+      raise ConfigurationError, "read_only must be a boolean"
+    end
+
+    def disconnect!
+      STATE_MONITOR.synchronize do
+        @groups&.each_value do |group|
+          group.databases.each(&:disconnect!)
         end
-        Hash[mapped]
       end
     end
+
+    def reset!
+      STATE_MONITOR.synchronize do
+        disconnect!
+        @config = nil
+        @groups = nil
+      end
+    end
+
+    private
+
+    def default_config_path
+      (Rails.root || Pathname.pwd).join("config/pgbouncerhero.yml")
+    end
+
+    def load_config
+      return {} unless config_path.file?
+
+      loaded = YAML.safe_load(ERB.new(config_path.read).result, aliases: true)
+      return {} if loaded.nil?
+      return loaded if loaded.is_a?(Hash)
+
+      raise ConfigurationError, "#{config_path} must contain a YAML mapping"
+    end
+
+    def selected_config(loaded)
+      selected = loaded.fetch(env, loaded)
+      unless selected.is_a?(Hash)
+        raise ConfigurationError, "configuration for #{env.inspect} must be a YAML mapping"
+      end
+
+      return default_config unless selected.key?("pgbouncers")
+
+      pgbouncers = selected["pgbouncers"]
+      unless pgbouncers.is_a?(Hash) && pgbouncers.any? && pgbouncers.values.all? { |databases| databases.is_a?(Hash) && databases.any? }
+        raise ConfigurationError, '"pgbouncers" must contain at least one group with one database'
+      end
+
+      selected
+    end
+
+    def default_config
+      {
+        "pgbouncers" => {
+          "default" => {
+            "primary" => {
+              "url" => ENV["PGBOUNCERHERO_DATABASE_URL"]
+            }
+          }
+        }
+      }
+    end
   end
+
+  self.env = ENV["RAILS_ENV"] || ENV["RACK_ENV"] || "development"
 end
+
+at_exit { PgBouncerHero.disconnect! }