pghero 2.2.1 → 2.7.0

data/lib/pghero.rb

@@ -5,6 +5,7 @@ require "forwardable"
 # methods
 require "pghero/methods/basic"
 require "pghero/methods/connections"
+require "pghero/methods/constraints"
 require "pghero/methods/explain"
 require "pghero/methods/indexes"
 require "pghero/methods/kill"
@@ -26,32 +27,37 @@ require "pghero/version"
 
 module PgHero
   autoload :Connection, "pghero/connection"
+  autoload :Stats, "pghero/stats"
   autoload :QueryStats, "pghero/query_stats"
+  autoload :SpaceStats, "pghero/space_stats"
 
   class Error < StandardError; end
   class NotEnabled < Error; end
 
+  MUTEX = Mutex.new
+
   # settings
   class << self
-    attr_accessor :long_running_query_sec, :slow_query_ms, :slow_query_calls, :explain_timeout_sec, :total_connections_threshold, :cache_hit_rate_threshold, :env, :show_migrations, :config_path
+    attr_accessor :long_running_query_sec, :slow_query_ms, :slow_query_calls, :explain_timeout_sec, :total_connections_threshold, :cache_hit_rate_threshold, :env, :show_migrations, :config_path, :filter_data
   end
   self.long_running_query_sec = (ENV["PGHERO_LONG_RUNNING_QUERY_SEC"] || 60).to_i
   self.slow_query_ms = (ENV["PGHERO_SLOW_QUERY_MS"] || 20).to_i
   self.slow_query_calls = (ENV["PGHERO_SLOW_QUERY_CALLS"] || 100).to_i
-  self.explain_timeout_sec = (ENV["PGHERO_EXPLAIN_TIMEOUT_SEC"] || 10).to_i
+  self.explain_timeout_sec = (ENV["PGHERO_EXPLAIN_TIMEOUT_SEC"] || 10).to_f
   self.total_connections_threshold = (ENV["PGHERO_TOTAL_CONNECTIONS_THRESHOLD"] || 500).to_i
   self.cache_hit_rate_threshold = 99
   self.env = ENV["RAILS_ENV"] || ENV["RACK_ENV"] || "development"
   self.show_migrations = true
   self.config_path = ENV["PGHERO_CONFIG_PATH"] || "config/pghero.yml"
+  self.filter_data = ENV["PGHERO_FILTER_DATA"].to_s.size > 0
 
   class << self
     extend Forwardable
     def_delegators :primary_database, :access_key_id, :analyze, :analyze_tables, :autoindex, :autovacuum_danger,
-      :best_index, :blocked_queries, :connection_sources, :connection_states, :connection_stats,
+      :best_index, :blocked_queries, :connections, :connection_sources, :connection_states, :connection_stats,
       :cpu_usage, :create_user, :database_size, :db_instance_identifier, :disable_query_stats, :drop_user,
       :duplicate_indexes, :enable_query_stats, :explain, :historical_query_stats_enabled?, :index_caching,
-      :index_hit_rate, :index_usage, :indexes, :invalid_indexes, :kill, :kill_all, :kill_long_running_queries,
+      :index_hit_rate, :index_usage, :indexes, :invalid_constraints, :invalid_indexes, :kill, :kill_all, :kill_long_running_queries,
       :last_stats_reset_time, :long_running_queries, :maintenance_info, :missing_indexes, :query_stats,
       :query_stats_available?, :query_stats_enabled?, :query_stats_extension_enabled?, :query_stats_readable?,
       :rds_stats, :read_iops_stats, :region, :relation_sizes, :replica?, :replication_lag, :replication_lag_stats,
@@ -68,6 +74,22 @@ module PgHero
       @time_zone || Time.zone
     end
 
+    # use method instead of attr_accessor to ensure
+    # this works if variable set after PgHero is loaded
+    def username
+      @username ||= config["username"] || ENV["PGHERO_USERNAME"]
+    end
+
+    # use method instead of attr_accessor to ensure
+    # this works if variable set after PgHero is loaded
+    def password
+      @password ||= config["password"] || ENV["PGHERO_PASSWORD"]
+    end
+
+    def stats_database_url
+      @stats_database_url ||= config["stats_database_url"] || ENV["PGHERO_STATS_DATABASE_URL"]
+    end
+
     def config
       @config ||= begin
         require "erb"
@@ -87,26 +109,49 @@ module PgHero
         elsif config_file_exists
           raise "Invalid config file"
         else
-          {
-            "databases" => {
-              "primary" => {
-                "url" => ENV["PGHERO_DATABASE_URL"] || ActiveRecord::Base.connection_config,
-                "db_instance_identifier" => ENV["PGHERO_DB_INSTANCE_IDENTIFIER"]
-              }
+          databases = {}
+
+          if !ENV["PGHERO_DATABASE_URL"] && spec_supported?
+            ActiveRecord::Base.configurations.configs_for(env_name: env, include_replicas: true).each do |db|
+              databases[db.spec_name] = {"spec" => db.spec_name}
+            end
+          end
+
+          if databases.empty?
+            databases["primary"] = {
+              "url" => ENV["PGHERO_DATABASE_URL"] || ActiveRecord::Base.connection_config
             }
+          end
+
+          if databases.size == 1
+            databases.values.first.merge!(
+              "db_instance_identifier" => ENV["PGHERO_DB_INSTANCE_IDENTIFIER"],
+              "gcp_database_id" => ENV["PGHERO_GCP_DATABASE_ID"],
+              "azure_resource_id" => ENV["PGHERO_AZURE_RESOURCE_ID"]
+            )
+          end
+
+          {
+            "databases" => databases
           }
         end
       end
     end
 
+    # ensure we only have one copy of databases
+    # so there's only one connection pool per database
     def databases
-      @databases ||= begin
-        Hash[
-          config["databases"].map do |id, c|
-            [id.to_sym, PgHero::Database.new(id, c)]
-          end
-        ]
+      unless defined?(@databases)
+        # only use mutex on initialization
+        MUTEX.synchronize do
+          # return if another process initialized while we were waiting
+          return @databases if defined?(@databases)
+
+          @databases = config["databases"].map { |id, c| [id.to_sym, Database.new(id, c)] }.to_h
+        end
       end
+
+      @databases
     end
 
     def primary_database
@@ -137,7 +182,7 @@ module PgHero
 
     def autoindex_all(create: false, verbose: true)
       each_database do |database|
-        puts "Autoindexing #{database}..." if verbose
+        puts "Autoindexing #{database.id}..." if verbose
         database.autoindex(create: create)
       end
     end
@@ -146,6 +191,26 @@ module PgHero
       ActiveSupport::NumberHelper.number_to_human_size(value, precision: 3)
     end
 
+    # delete previous stats
+    # go database by database to use an index
+    # stats for old databases are not cleaned up since we can't use an index
+    def clean_query_stats
+      each_database do |database|
+        database.clean_query_stats
+      end
+    end
+
+    def clean_space_stats
+      each_database do |database|
+        database.clean_space_stats
+      end
+    end
+
+    # private
+    def spec_supported?
+      ActiveRecord::VERSION::MAJOR >= 6
+    end
+
     private
 
     def each_database

data/lib/pghero/database.rb

@@ -2,6 +2,7 @@ module PgHero
   class Database
     include Methods::Basic
     include Methods::Connections
+    include Methods::Constraints
     include Methods::Explain
     include Methods::Indexes
     include Methods::Kill
@@ -22,16 +23,17 @@ module PgHero
     def initialize(id, config)
       @id = id
       @config = config || {}
+
+      # preload model to ensure only one connection pool
+      # this doesn't actually start any connections
+      @adapter_checked = false
+      @connection_model = build_connection_model
     end
 
     def name
       @name ||= @config["name"] || id.titleize
     end
 
-    def db_instance_identifier
-      @db_instance_identifier ||= @config["db_instance_identifier"]
-    end
-
     def capture_query_stats?
       config["capture_query_stats"] != false
     end
@@ -53,35 +55,118 @@ module PgHero
     end
 
     def explain_timeout_sec
-      (config["explain_timeout_sec"] || PgHero.config["explain_timeout_sec"] || PgHero.explain_timeout_sec).to_i
+      (config["explain_timeout_sec"] || PgHero.config["explain_timeout_sec"] || PgHero.explain_timeout_sec).to_f
     end
 
     def long_running_query_sec
       (config["long_running_query_sec"] || PgHero.config["long_running_query_sec"] || PgHero.long_running_query_sec).to_i
     end
 
+    # defaults to 100 megabytes
     def index_bloat_bytes
-      (config["index_bloat_bytes"] || PgHero.config["index_bloat_bytes"] || 100.megabytes).to_i
+      (config["index_bloat_bytes"] || PgHero.config["index_bloat_bytes"] || 104857600).to_i
     end
 
-    private
+    def aws_access_key_id
+      config["aws_access_key_id"] || PgHero.config["aws_access_key_id"] || ENV["PGHERO_ACCESS_KEY_ID"] || ENV["AWS_ACCESS_KEY_ID"]
+    end
 
-    def connection_model
-      @connection_model ||= begin
-        url = config["url"]
-        Class.new(PgHero::Connection) do
-          def self.name
-            "PgHero::Connection::Database#{object_id}"
+    def aws_secret_access_key
+      config["aws_secret_access_key"] || PgHero.config["aws_secret_access_key"] || ENV["PGHERO_SECRET_ACCESS_KEY"] || ENV["AWS_SECRET_ACCESS_KEY"]
+    end
+
+    def aws_region
+      config["aws_region"] || PgHero.config["aws_region"] || ENV["PGHERO_REGION"] || ENV["AWS_REGION"] || (defined?(Aws) && Aws.config[:region]) || "us-east-1"
+    end
+
+    # environment variable is only used if no config file
+    def aws_db_instance_identifier
+      @aws_db_instance_identifier ||= config["aws_db_instance_identifier"] || config["db_instance_identifier"]
+    end
+
+    # environment variable is only used if no config file
+    def gcp_database_id
+      @gcp_database_id ||= config["gcp_database_id"]
+    end
+
+    # environment variable is only used if no config file
+    def azure_resource_id
+      @azure_resource_id ||= config["azure_resource_id"]
+    end
+
+    # must check keys for booleans
+    def filter_data
+      unless defined?(@filter_data)
+        @filter_data =
+          if config.key?("filter_data")
+            config["filter_data"]
+          elsif PgHero.config.key?("filter_data")
+            PgHero.config.key?("filter_data")
+          else
+            PgHero.filter_data
           end
-          case url
-          when String
-            url = "#{url}#{url.include?("?") ? "&" : "?"}connect_timeout=5" unless url.include?("connect_timeout=")
-          when Hash
-            url[:connect_timeout] ||= 5
+
+        if @filter_data
+          begin
+            require "pg_query"
+          rescue LoadError
+            raise Error, "pg_query required for filter_data"
           end
-          establish_connection url if url
         end
       end
+
+      @filter_data
+    end
+
+    # TODO remove in next major version
+    alias_method :access_key_id, :aws_access_key_id
+    alias_method :secret_access_key, :aws_secret_access_key
+    alias_method :region, :aws_region
+    alias_method :db_instance_identifier, :aws_db_instance_identifier
+
+    private
+
+    # check adapter lazily
+    def connection_model
+      unless @adapter_checked
+        # 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 =~ /postg/i
+          raise Error, "Invalid connection URL"
+        end
+        @adapter_checked = true
+      end
+
+      @connection_model
+    end
+
+    # just return the model
+    # do not start a connection
+    def build_connection_model
+      url = config["url"]
+
+      # resolve spec
+      if !url && config["spec"]
+        raise Error, "Spec requires Rails 6+" unless PgHero.spec_supported?
+        resolved = ActiveRecord::Base.configurations.configs_for(env_name: PgHero.env, spec_name: config["spec"], include_replicas: true)
+        raise Error, "Spec not found: #{config["spec"]}" unless resolved
+        url = resolved.config
+      end
+
+      Class.new(PgHero::Connection) do
+        def self.name
+          "PgHero::Connection::Database#{object_id}"
+        end
+
+        case url
+        when String
+          url = "#{url}#{url.include?("?") ? "&" : "?"}connect_timeout=5" unless url.include?("connect_timeout=")
+        when Hash
+          url[:connect_timeout] ||= 5
+        end
+        establish_connection url if url
+      end
     end
   end
 end

data/lib/pghero/methods/basic.rb

@@ -18,6 +18,10 @@ module PgHero
         select_one("SELECT current_database()")
       end
 
+      def current_user
+        select_one("SELECT current_user")
+      end
+
       def server_version
         @server_version ||= select_one("SHOW server_version")
       end
@@ -32,14 +36,34 @@ module PgHero
 
       private
 
-      def select_all(sql, conn = nil)
+      def select_all(sql, conn: nil, query_columns: [])
         conn ||= connection
         # squish for logs
         retries = 0
         begin
           result = conn.select_all(add_source(squish(sql)))
-          cast_method = ActiveRecord::VERSION::MAJOR < 5 ? :type_cast : :cast_value
-          result.map { |row| Hash[row.map { |col, val| [col.to_sym, result.column_types[col].send(cast_method, val)] }] }
+          result = result.map { |row| Hash[row.map { |col, val| [col.to_sym, result.column_types[col].send(:cast_value, val)] }] }
+          if filter_data
+            query_columns.each do |column|
+              result.each do |row|
+                begin
+                  row[column] = PgQuery.normalize(row[column])
+                rescue PgQuery::ParseError
+                  # try replacing "interval $1" with "$1::interval"
+                  # see https://github.com/lfittl/pg_query/issues/169 for more info
+                  # this is not ideal since it changes the query slightly
+                  # we could skip normalization
+                  # but this has a very small chance of data leakage
+                  begin
+                    row[column] = PgQuery.normalize(row[column].gsub(/\binterval\s+(\$\d+)\b/i, "\\1::interval"))
+                  rescue PgQuery::ParseError
+                    row[column] = "<unable to filter data>"
+                  end
+                end
+              end
+            end
+          end
+          result
         rescue ActiveRecord::StatementInvalid => e
           # fix for random internal errors
           if e.message.include?("PG::InternalError") && retries < 2
@@ -52,8 +76,8 @@ module PgHero
         end
       end
 
-      def select_all_stats(sql)
-        select_all(sql, stats_connection)
+      def select_all_stats(sql, **options)
+        select_all(sql, **options, conn: stats_connection)
       end
 
       def select_all_size(sql)
@@ -64,12 +88,12 @@ module PgHero
         result
       end
 
-      def select_one(sql, conn = nil)
-        select_all(sql, conn).first.values.first
+      def select_one(sql, conn: nil)
+        select_all(sql, conn: conn).first.values.first
       end
 
       def select_one_stats(sql)
-        select_one(sql, stats_connection)
+        select_one(sql, conn: stats_connection)
       end
 
       def execute(sql)
@@ -81,7 +105,7 @@ module PgHero
       end
 
       def stats_connection
-        ::PgHero::QueryStats.connection
+        ::PgHero::Stats.connection
       end
 
       def insert_stats(table, columns, values)
@@ -125,22 +149,7 @@ module PgHero
       end
 
       def table_exists?(table)
-        ["PostgreSQL", "PostGIS"].include?(stats_connection.adapter_name) &&
-        select_one_stats(<<-SQL
-          SELECT EXISTS (
-            SELECT
-              1
-            FROM
-              pg_catalog.pg_class c
-            INNER JOIN
-              pg_catalog.pg_namespace n ON n.oid = c.relnamespace
-            WHERE
-              n.nspname = 'public'
-              AND c.relname = #{quote(table)}
-              AND c.relkind = 'r'
-          )
-        SQL
-        )
+        stats_connection.table_exists?(table)
       end
     end
   end

data/lib/pghero/methods/connections.rb

@@ -1,6 +1,41 @@
 module PgHero
   module Methods
     module Connections
+      def connections
+        if server_version_num >= 90500
+          select_all <<-SQL
+            SELECT
+              pg_stat_activity.pid,
+              datname AS database,
+              usename AS user,
+              application_name AS source,
+              client_addr AS ip,
+              state,
+              ssl
+            FROM
+              pg_stat_activity
+            LEFT JOIN
+              pg_stat_ssl ON pg_stat_activity.pid = pg_stat_ssl.pid
+            ORDER BY
+              pg_stat_activity.pid
+          SQL
+        else
+          select_all <<-SQL
+            SELECT
+              pid,
+              datname AS database,
+              usename AS user,
+              application_name AS source,
+              client_addr AS ip,
+              state
+            FROM
+              pg_stat_activity
+            ORDER BY
+              pid
+          SQL
+        end
+      end
+
       def total_connections
         select_one("SELECT COUNT(*) FROM pg_stat_activity")
       end