pghero 2.4.1 → 2.7.0

data/lib/pghero/database.rb

@@ -23,6 +23,11 @@ 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
@@ -50,15 +55,16 @@ 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
 
     def aws_access_key_id
@@ -73,8 +79,43 @@ module PgHero
       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
-      @db_instance_identifier ||= config["aws_db_instance_identifier"] || config["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
+
+        if @filter_data
+          begin
+            require "pg_query"
+          rescue LoadError
+            raise Error, "pg_query required for filter_data"
+          end
+        end
+      end
+
+      @filter_data
     end
 
     # TODO remove in next major version
@@ -85,27 +126,46 @@ module PgHero
 
     private
 
+    # check adapter lazily
     def connection_model
-      @connection_model ||= begin
-        url = config["url"]
-        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
+      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
-        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
+        @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

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,13 +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)))
-          result.map { |row| Hash[row.map { |col, val| [col.to_sym, result.column_types[col].send(:cast_value, 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
@@ -51,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)
@@ -63,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)

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

data/lib/pghero/methods/explain.rb

@@ -6,7 +6,7 @@ module PgHero
         explanation = nil
 
         # use transaction for safety
-        with_transaction(statement_timeout: (explain_timeout_sec * 1000), rollback: true) do
+        with_transaction(statement_timeout: (explain_timeout_sec * 1000).round, rollback: true) do
           if (sql.sub(/;\z/, "").include?(";") || sql.upcase.include?("COMMIT")) && !explain_safe?
             raise ActiveRecord::StatementInvalid, "Unsafe statement"
           end

data/lib/pghero/methods/maintenance.rb

@@ -57,7 +57,9 @@ module PgHero
             last_vacuum,
             last_autovacuum,
             last_analyze,
-            last_autoanalyze
+            last_autoanalyze,
+            n_dead_tup AS dead_rows,
+            n_live_tup AS live_rows
           FROM
             pg_stat_user_tables
           ORDER BY

data/lib/pghero/methods/queries.rb

@@ -2,7 +2,7 @@ module PgHero
   module Methods
     module Queries
       def running_queries(min_duration: nil, all: false)
-        select_all <<-SQL
+        query = <<-SQL
           SELECT
             pid,
             state,
@@ -24,6 +24,8 @@ module PgHero
           ORDER BY
             COALESCE(query_start, xact_start) DESC
         SQL
+
+        select_all(query, query_columns: [:query])
       end
 
       def long_running_queries
@@ -33,7 +35,7 @@ module PgHero
       # from https://wiki.postgresql.org/wiki/Lock_Monitoring
       # and https://big-elephants.com/2013-09/exploring-query-locks-in-postgres/
       def blocked_queries
-        select_all <<-SQL
+        query = <<-SQL
           SELECT
             COALESCE(blockingl.relation::regclass::text,blockingl.locktype) as locked_item,
             blockeda.pid AS blocked_pid,
@@ -65,6 +67,8 @@ module PgHero
           ORDER BY
             blocked_duration DESC
         SQL
+
+        select_all(query, query_columns: [:blocked_query, :current_or_recent_query_in_blocking_process])
       end
     end
   end

data/lib/pghero/methods/query_stats.rb

@@ -2,7 +2,7 @@ module PgHero
   module Methods
     module QueryStats
       def query_stats(historical: false, start_at: nil, end_at: nil, min_average_time: nil, min_calls: nil, **options)
-        current_query_stats = historical && end_at && end_at < Time.now ? [] : current_query_stats(options)
+        current_query_stats = historical && end_at && end_at < Time.now ? [] : current_query_stats(**options)
         historical_query_stats = historical && historical_query_stats_enabled? ? historical_query_stats(start_at: start_at, end_at: end_at, **options) : []
 
         query_stats = combine_query_stats((current_query_stats + historical_query_stats).group_by { |q| [q[:query_hash], q[:user]] })
@@ -56,8 +56,46 @@ module PgHero
         true
       end
 
-      def reset_query_stats(raise_errors: false)
-        execute("SELECT pg_stat_statements_reset()")
+      # TODO scope by database in PgHero 3.0
+      # (add database: database_name to options)
+      def reset_query_stats(**options)
+        reset_instance_query_stats(**options)
+      end
+
+      # resets query stats for the entire instance
+      # it's possible to reset stats for a specific
+      # database, user or query hash in Postgres 12+
+      def reset_instance_query_stats(database: nil, user: nil, query_hash: nil, raise_errors: false)
+        if database || user || query_hash
+          raise PgHero::Error, "Requires PostgreSQL 12+" if server_version_num < 120000
+
+          if database
+            database_id = execute("SELECT oid FROM pg_database WHERE datname = #{quote(database)}").first.try(:[], "oid")
+            raise PgHero::Error, "Database not found: #{database}" unless database_id
+          else
+            database_id = 0
+          end
+
+          if user
+            user_id = execute("SELECT usesysid FROM pg_user WHERE usename = #{quote(user)}").first.try(:[], "usesysid")
+            raise PgHero::Error, "User not found: #{user}" unless user_id
+          else
+            user_id = 0
+          end
+
+          if query_hash
+            query_id = query_hash.to_i
+            # may not be needed
+            # but not intuitive that all query hashes are reset with 0
+            raise PgHero::Error, "Invalid query hash: #{query_hash}" if query_id == 0
+          else
+            query_id = 0
+          end
+
+          execute("SELECT pg_stat_statements_reset(#{quote(user_id.to_i)}, #{quote(database_id.to_i)}, #{quote(query_id.to_i)})")
+        else
+          execute("SELECT pg_stat_statements_reset()")
+        end
         true
       rescue ActiveRecord::StatementInvalid => e
         raise e if raise_errors
@@ -104,31 +142,32 @@ module PgHero
           query_stats[database_id] = query_stats(limit: 1000000, database: database_name)
         end
 
-        supports_query_hash = supports_query_hash?
+        query_stats = query_stats.select { |_, v| v.any? }
+
+        # nothing to do
+        return if query_stats.empty?
 
-        if query_stats.any? { |_, v| v.any? } && reset_query_stats(raise_errors: raise_errors)
+        # use mapping, not query stats here
+        # TODO add option for this, and make default in PgHero 3.0
+        if false # mapping.size == 1 && server_version_num >= 120000
           query_stats.each do |db_id, db_query_stats|
-            if db_query_stats.any?
-              values =
-                db_query_stats.map do |qs|
-                  [
-                    db_id,
-                    qs[:query],
-                    qs[:total_minutes] * 60 * 1000,
-                    qs[:calls],
-                    now,
-                    supports_query_hash ? qs[:query_hash] : nil,
-                    qs[:user]
-                  ]
-                end
-
-              columns = %w[database query total_time calls captured_at query_hash user]
-              insert_stats("pghero_query_stats", columns, values)
+            if reset_query_stats(database: mapping[db_id], raise_errors: raise_errors)
+              insert_query_stats(db_id, db_query_stats, now)
+            end
+          end
+        else
+          if reset_query_stats(raise_errors: raise_errors)
+            query_stats.each do |db_id, db_query_stats|
+              insert_query_stats(db_id, db_query_stats, now)
             end
           end
         end
       end
 
+      def clean_query_stats
+        PgHero::QueryStats.where(database: id).where("captured_at < ?", 14.days.ago).delete_all
+      end
+
       def slow_queries(query_stats: nil, **options)
         query_stats ||= self.query_stats(options)
         query_stats.select { |q| q[:calls].to_i >= slow_query_calls.to_i && q[:average_time].to_f >= slow_query_ms.to_f }
@@ -166,14 +205,15 @@ module PgHero
         if query_stats_enabled?
           limit ||= 100
           sort ||= "total_minutes"
-          select_all <<-SQL
+          total_time = server_version_num >= 130000 ? "(total_plan_time + total_exec_time)" : "total_time"
+          query = <<-SQL
             WITH query_stats AS (
               SELECT
                 LEFT(query, 10000) AS query,
                 #{supports_query_hash? ? "queryid" : "md5(query)"} AS query_hash,
                 rolname AS user,
-                (total_time / 1000 / 60) AS total_minutes,
-                (total_time / calls) AS average_time,
+                (#{total_time} / 1000 / 60) AS total_minutes,
+                (#{total_time} / calls) AS average_time,
                 calls
               FROM
                 pg_stat_statements
@@ -182,6 +222,7 @@ module PgHero
               INNER JOIN
                 pg_roles ON pg_roles.oid = pg_stat_statements.userid
               WHERE
+                calls > 0 AND
                 pg_database.datname = #{database ? quote(database) : "current_database()"}
                 #{query_hash ? "AND queryid = #{quote(query_hash)}" : nil}
             )
@@ -200,6 +241,11 @@ module PgHero
               #{quote_table_name(sort)} DESC
             LIMIT #{limit.to_i}
           SQL
+
+          # we may be able to skip query_columns
+          # in more recent versions of Postgres
+          # as pg_stat_statements should be already normalized
+          select_all(query, query_columns: [:query])
         else
           raise NotEnabled, "Query stats not enabled"
         end
@@ -208,7 +254,7 @@ module PgHero
       def historical_query_stats(sort: nil, start_at: nil, end_at: nil, query_hash: nil)
         if historical_query_stats_enabled?
           sort ||= "total_minutes"
-          select_all_stats <<-SQL
+          query = <<-SQL
             WITH query_stats AS (
               SELECT
                 #{supports_query_hash? ? "query_hash" : "md5(query)"} AS query_hash,
@@ -244,6 +290,10 @@ module PgHero
               #{quote_table_name(sort)} DESC
             LIMIT 100
           SQL
+
+          # we can skip query_columns if all stored data is normalized
+          # for now, assume it's not
+          select_all_stats(query, query_columns: [:query, :explainable_query])
         else
           raise NotEnabled, "Historical query stats not enabled"
         end
@@ -276,6 +326,24 @@ module PgHero
       def normalize_query(query)
         squish(query.to_s.gsub(/\?(, ?\?)+/, "?").gsub(/\/\*.+?\*\//, ""))
       end
+
+      def insert_query_stats(db_id, db_query_stats, now)
+        values =
+          db_query_stats.map do |qs|
+            [
+              db_id,
+              qs[:query],
+              qs[:total_minutes] * 60 * 1000,
+              qs[:calls],
+              now,
+              supports_query_hash? ? qs[:query_hash] : nil,
+              qs[:user]
+            ]
+          end
+
+        columns = %w[database query total_time calls captured_at query_hash user]
+        insert_stats("pghero_query_stats", columns, values)
+      end
     end
   end
 end