Multiple improvements/features (pagination, etc)

Gemfile.lock

@@ -5,6 +5,7 @@ PATH
       coderay
       pg
       sinatra
+      sinatra-param
 
 GEM
   remote: https://rubygems.org/
@@ -14,13 +15,12 @@ GEM
     diff-lcs (1.3)
     eventmachine (1.2.5)
     method_source (0.9.0)
-    mustermann (1.0.2)
     pg (1.0.0)
     pry (0.11.3)
       coderay (~> 1.1.0)
       method_source (~> 0.9.0)
-    rack (2.0.4)
-    rack-protection (2.0.1)
+    rack (1.6.9)
+    rack-protection (1.5.5)
       rack
     rack-test (0.8.3)
       rack (>= 1.0, < 3)
@@ -38,11 +38,12 @@ GEM
       diff-lcs (>= 1.2.0, < 2.0)
       rspec-support (~> 3.7.0)
     rspec-support (3.7.1)
-    sinatra (2.0.1)
-      mustermann (~> 1.0)
-      rack (~> 2.0)
-      rack-protection (= 2.0.1)
-      tilt (~> 2.0)
+    sinatra (1.4.8)
+      rack (~> 1.5)
+      rack-protection (~> 1.4)
+      tilt (>= 1.3, < 3)
+    sinatra-param (1.4.0)
+      sinatra (~> 1.3)
     thin (1.7.2)
       daemons (~> 1.0, >= 1.0.9)
       eventmachine (~> 1.0, >= 1.0.4)

config.yml.example

@@ -1,5 +1,6 @@
-user: foo
-password: bar
-database: rails_development
-host: localhost
-port:
+default:
+  user: postgres
+  password:
+  database: postgres
+  host: postgres
+  port: 5432

lib/pg_web_stats.rb

@@ -3,46 +3,78 @@
 require 'yaml'
 
 class PgWebStats
-  attr_accessor :config, :connection
+  attr_accessor :default_server, :connections
 
   def initialize(config_path = 'config.yml')
     hash = config_path.is_a?(Hash) ? config_path : YAML.load_file(config_path)
-    self.config = Hash[hash.map{ |k, v| [k.to_s, v] }]
-    self.connection = PG.connect(
-      dbname: config['database'],
-      host: config['host'],
-      user: config['user'] || config['username'],
-      password: config['password'],
-      port: config['port']
-    )
+
+    self.default_server = nil
+    self.connections = Hash.new
+    hash.each do |name, config|
+      self.default_server = name unless self.default_server
+      self.connections[name] = PG.connect(
+        dbname: config['database'],
+        host: config['host'],
+        user: config['user'] || config['username'],
+        password: config['password'],
+        port: config['port']
+      )
+    end
+    if self.connections.length < 1
+      raise RuntimeError, "configuration must contain at least one server"
+    end
   end
 
-  def get_stats(params = { order: "total_time desc" })
-    query = build_stats_query(params)
+  def reset_stats(server)
+    connections[server].exec("SELECT pg_stat_statements_reset()")
+  end
+
+  def get_stats(server, params = { order: "total_time desc" })
+    connection = connections[server]
+    query = build_stats_query_base(connection, "COUNT(*) AS count", params)
+
+    count = 0
+    connection.exec(query) do |result|
+      count = result[0]["count"].to_i
+    end
+
+    query = build_stats_query(connection, "*", params)
 
     results = []
     connection.exec(query) do |result|
       result.each do |row|
-        results << Row.new(row, users, databases)
+        results << Row.new(row, users(server), databases(server))
       end
     end
 
-    results
+    {total: count, items: results}
   end
 
-  def users
-    @users ||= select_by_oid("select oid, rolname from pg_authid order by rolname;", 'rolname')
+  def users(server)
+    @users ||= Hash.new
+
+    if not @users.has_key? server
+      @users[server] = select_by_oid(server, "select oid, rolname from pg_authid order by rolname;", 'rolname')
+    end
+
+    @users[server]
   end
 
-  def databases
-    @databases ||= select_by_oid("select oid, datname from pg_database order by datname;", 'datname')
+  def databases(server)
+    @databases ||= Hash.new
+
+    if not @databases.has_key? server
+      @databases[server] = select_by_oid(server, "select oid, datname from pg_database order by datname;", 'datname')
+    end
+
+    @databases[server]
   end
 
   private
 
-  def select_by_oid(select_query, row_name)
+  def select_by_oid(server, select_query, row_name)
     @selection = {}
-    connection.exec(select_query) do |result|
+    connections[server].exec(select_query) do |result|
       result.each do |row|
         @selection[row['oid']] = row[row_name]
       end
@@ -51,32 +83,50 @@ def select_by_oid(select_query, row_name)
     @selection
   end
 
-  def build_stats_query(params)
-    order_by = params[:order]
-
-    query = "SELECT * FROM pg_stat_statements"
+  def build_stats_query_base(connection, what, params)
+    query = "SELECT #{what} FROM pg_stat_statements"
 
     where_conditions = []
 
     userid = params[:userid]
     if userid && !userid.empty?
-      where_conditions << "userid='#{userid.gsub("'", "''")}'"
+      where_conditions << "userid=#{userid}"
     end
 
     dbid = params[:dbid]
     if dbid && !dbid.empty?
-      where_conditions << "dbid='#{dbid.gsub("'", "''")}'"
+      where_conditions << "dbid=#{dbid}"
     end
 
     q = params[:q]
     if q && !q.empty?
-      where_conditions << "query LIKE '#{q.gsub("'", "''")}%'"
+      where_conditions << "query LIKE '#{connection.escape_string(q)}%'"
     end
 
     query += " WHERE #{where_conditions.join(" AND ")}" if where_conditions.size > 0
 
+    query
+  end
+
+  def build_stats_query(connection, what, params)
+    order_by = params[:order]
+
+    query = build_stats_query_base(connection, "*", params)
+
+    order_by = if params[:order_by] && params[:direction]
+      "#{params[:order_by]} #{params[:direction]}"
+    else
+      "total_time desc"
+    end
     query += " ORDER BY #{order_by}"
 
+    count = params[:count] ? params[:count].clamp(1, 1000) : 25
+    query += " LIMIT #{count}"
+
+    if params[:offset] > 0
+      query += " OFFSET #{params[:offset]}"
+    end
+
     query
   end
 end

lib/pg_web_stats_app.rb

@@ -1,44 +1,100 @@
 require 'sinatra'
+require 'sinatra/param'
 require 'pg_web_stats'
 
 class PgWebStatsApp < Sinatra::Base
+  helpers Sinatra::Param
+
   set :root,  File.expand_path(File.join(File.dirname(__FILE__), '../'))
   set :views, File.join(settings.root, 'views')
 
   helpers do
-    def sort_link(title, key, alt_title = nil)
-      direction = if params[:order_by] == key && params[:direction] == "desc"
+    def link(title, update, alt_title = nil)
+      update = Hash[update.map{ |k, v| [k.to_s, v] }]
+      url = "?" + URI.encode_www_form(params.merge(update))
+
+      "<a href='#{url}' title='#{alt_title}'>#{title}</a>"
+    end
+
+    def page_links
+      offset = params[:offset]
+      count = params[:count]
+
+      pages = @stats[:total].fdiv(count).floor  # Pages are 0-indexed, so floor
+      this_page = offset.fdiv(count).floor
+
+      Enumerator.new do |enum|
+        if offset > 0
+          enum.yield text: 'prev', offset: [0, offset - count].max
+        end
+
+        [0, this_page - 4].max.upto([this_page + 4, pages].min) do |page|
+          classname = page == this_page ? "active" : ""
+          enum.yield text: (page + 1).to_s, offset: page * count, class: classname
+        end
+
+        if @stats[:items].length >= count
+          enum.yield text: 'next', offset: offset + count
+        end
+      end
+    end
+
+    def sort_link(title, update, alt_title = nil)
+      direction = if params[:order_by] == update && params[:direction] == "desc"
         "asc"
       else
         "desc"
       end
+      update = {
+        order_by: update,
+        direction: direction,
+        offset: 0   # Changing sorting resets pagination
+      }
+      link title, update, alt_title
+    end
 
-      url = "?order_by=#{key}&direction=#{direction}"
-      url += "&userid=#{params[:userid]}" if params[:userid]
-      url += "&dbid=#{params[:dbid]}" if params[:dbid]
-      url += "&q=#{params[:q]}" if params[:q]
-
-      "<a href='#{url}' title='#{alt_title}'>#{title}</a>"
+    def page_link(info)
+      text = info.delete(:text)
+      classname = info.delete(:class)
+      attrs = classname && !classname.empty? ? " class=\"#{classname}\"" : ""
+      "<li#{attrs}>" + link(text, info) + "</li>"
     end
   end
 
   get '/' do
-    order_by = if params[:order_by] && params[:direction]
-      "#{params[:order_by]} #{params[:direction]}"
+    @servers = PG_WEB_STATS.connections.keys
+    if @servers.length == 1
+      redirect '/' + PG_WEB_STATS.default_server + '/', 307
     else
-      "total_time desc"
+      erb :servers, layout: :application
     end
+  end
 
-    @stats = PG_WEB_STATS.get_stats(
-      order: order_by,
-      userid: params[:userid],
-      dbid: params[:dbid],
-      q: params[:q]
-    )
+  post '/:server/reset-stats/' do |server|
+    PG_WEB_STATS.reset_stats(server)
+    redirect "/#{server}/", 303
+  end
+
+  get '/:server/' do |server|
+    param :q,          String
+    param :userid,     String, format: /^\d*$/
+    param :dbid,       String, format: /^\d*$/
+    param :count,      Integer, default: 25
+    param :offset,     Integer, default: 0
+    param :order_by,   String, default: "total_time"
+    param :direction,  String, in: ["asc", "desc"], default: "desc"
+
+    all_keys = %w{q userid dbid count offset order_by direction}
+    params.select {|key| all_keys.include? key}
+
+    if not PG_WEB_STATS.connections.has_key? server
+      halt(404)
+    end
 
-    @databases = PG_WEB_STATS.databases
-    @users = PG_WEB_STATS.users
+    @stats = PG_WEB_STATS.get_stats(server, params)
+    @databases = PG_WEB_STATS.databases(server)
+    @users = PG_WEB_STATS.users(server)
 
-    erb :queries, layout: :application
+    erb :queries, layout: :application, locals: {server: server}
   end
 end

pg_web_stats.gemspec

@@ -19,6 +19,7 @@ Gem::Specification.new do |spec|
 
   spec.add_dependency "pg"
   spec.add_dependency "sinatra"
+  spec.add_dependency "sinatra-param"
   spec.add_dependency "coderay"
   spec.add_development_dependency "bundler", "~> 1.5"
   spec.add_development_dependency "rake"