query_stats.rb in pghero-2.7.0

- old
+ new

@@ -54,12 +54,50 @@
       def disable_query_stats
         execute("DROP EXTENSION IF EXISTS pg_stat_statements")
         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
         false
       end
@@ -102,35 +140,36 @@
         query_stats = {}
         mapping.each do |database_id, database_name|
           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? }
 
-        if query_stats.any? { |_, v| v.any? } && reset_query_stats(raise_errors: raise_errors)
-          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
+        # nothing to do
+        return if query_stats.empty?
 
-              columns = %w[database query total_time calls captured_at query_hash user]
-              insert_stats("pghero_query_stats", columns, values)
+        # 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 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 }
       end
 
@@ -284,9 +323,27 @@
 
       # removes comments
       # combines ?, ?, ? => ?
       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