text: # coding=utf-8
  
  """
  Collect metrics from postgresql
  
  #### Dependencies
  
   * psycopg2
  
  """
  
  import diamond.collector
  
  try:
      import psycopg2
      import psycopg2.extras
      psycopg2  # workaround for pyflakes issue #13
  except ImportError:
      psycopg2 = None
  
  
  class PostgresqlCollector(diamond.collector.Collector):
  
      def get_default_config_help(self):
          config_help = super(PostgresqlCollector, self).get_default_config_help()
          config_help.update({
              'host': 'Hostname',
              'user': 'Username',
              'password': 'Password',
              'port': 'Port number',
              'underscore': 'Convert _ to .',
              'extended': 'Enable collection of extended database stats',
          })
          return config_help
  
      def get_default_config(self):
          """
          Return default config.
          """
          config = super(PostgresqlCollector, self).get_default_config()
          config.update({
              'path': 'postgres',
              'host': 'localhost',
              'user': 'postgres',
              'password': 'postgres',
              'port': 5432,
              'underscore': False,
              'extended': False,
              'method': 'Threaded'
          })
          return config
  
      def collect(self):
          if psycopg2 is None:
              self.log.error('Unable to import module psycopg2')
              return {}
  
          # Create database-specific connections
          self.connections = {}
          for db in self._get_db_names():
              self.connections[db] = self._connect(database=db)
  
          if self.config['extended']:
              metrics = registry['extended']
          else:
              metrics = registry['basic']
  
          # Iterate every QueryStats class
          for klass in metrics:
              stat = klass(self.connections, underscore=self.config['underscore'])
              stat.fetch()
              [self.publish(metric, value) for metric, value in stat]
  
          # Cleanup
          [conn.close() for conn in self.connections.itervalues()]
  
      def _get_db_names(self):
          query = """
              SELECT datname FROM pg_database
              WHERE datallowconn AND NOT datistemplate
              AND NOT datname='postgres' ORDER BY 1
          """
          conn = self._connect()
          cursor = conn.cursor(cursor_factory=psycopg2.extras.DictCursor)
          cursor.execute(query)
          datnames = [d['datname'] for d in cursor.fetchall()]
          conn.close()
  
          # Exclude `postgres` database list, unless it is the
          # only database available (required for querying pg_stat_database)
          if not datnames:
              datnames = ['postgres']
          return datnames
  
      def _connect(self, database=None):
          conn_args = {
            'host':self.config['host'],
            'user':self.config['user'],
            'password':self.config['password'],
            'port':self.config['port']
          }
  
          if database:
            conn_args['database'] = database
  
          conn = psycopg2.connect(**conn_args)
  
          # Avoid using transactions, set isolation level to autocommit
          conn.set_isolation_level(0)
          return conn
  
  
  class QueryStats(object):
      def __init__(self, conns, underscore=False):
          self.connections = conns
          self.underscore = underscore
  
      def _translate_datname(self, db):
          if self.underscore:
              db = db.replace("_", ".")
          return db
  
      def fetch(self):
          self.data = list()
  
          for db, conn in self.connections.iteritems():
              cursor = conn.cursor(cursor_factory=psycopg2.extras.DictCursor)
              cursor.execute(self.query)
  
              for row in cursor.fetchall():
                  # If row is length 2, assume col1, col2 forms key: value
                  if len(row) == 2:
                      self.data.append({
                          'datname': self._translate_datname(db),
                          'metric': row[0],
                          'value': row[1],
                      })
  
                  # If row > length 2, assume each column name maps to
                  # key => value
                  else:
                      for key, value in row.iteritems():
                          if key in ('datname', 'schemaname', 'relname',
                                     'indexrelname',):
                              continue
  
                          self.data.append({
                              'datname': self._translate_datname(row.get(
                                  'datname', db)),
                              'schemaname': row.get('schemaname', None),
                              'relname': row.get('relname', None),
                              'indexrelname': row.get('indexrelname', None),
                              'metric': key,
                              'value': value,
                          })
  
              # Setting multi_db to True will run this query on all known
              # databases. This is bad for queries that hit views like
              # pg_database, which are shared across databases.
              #
              # If multi_db is False, bail early after the first query
              # iteration. Otherwise, continue to remaining databases.
              if not self.multi_db:
                  break
  
      def __iter__(self):
          for data_point in self.data:
              yield (self.path % data_point, data_point['value'])
  
  
  class DatabaseStats(QueryStats):
      """
      Database-level summary stats
      """
      path = "database.%(datname)s.%(metric)s"
      multi_db = False
      query = """
          SELECT pg_stat_database.datname as datname,
                 pg_stat_database.numbackends as numbackends,
                 pg_stat_database.xact_commit as xact_commit,
                 pg_stat_database.xact_rollback as xact_rollback,
                 pg_stat_database.blks_read as blks_read,
                 pg_stat_database.blks_hit as blks_hit,
                 pg_stat_database.tup_returned as tup_returned,
                 pg_stat_database.tup_fetched as tup_fetched,
                 pg_stat_database.tup_inserted as tup_inserted,
                 pg_stat_database.tup_updated as tup_updated,
                 pg_stat_database.tup_deleted as tup_deleted,
                 pg_stat_database.conflicts as conflicts,
                 pg_database_size(pg_database.datname) AS size
          FROM pg_database
          JOIN pg_stat_database
          ON pg_database.datname = pg_stat_database.datname
          WHERE pg_stat_database.datname
          NOT IN ('template0','template1','postgres')
      """
  
  
  class UserTableStats(QueryStats):
      path = "%(datname)s.tables.%(schemaname)s.%(relname)s.%(metric)s"
      multi_db = True
      query = """
          SELECT relname,
                 schemaname,
                 seq_scan,
                 seq_tup_read,
                 idx_scan,
                 idx_tup_fetch,
                 n_tup_ins,
                 n_tup_upd,
                 n_tup_del,
                 n_tup_hot_upd,
                 n_live_tup,
                 n_dead_tup
          FROM pg_stat_user_tables
      """
  
  
  class UserIndexStats(QueryStats):
      path = "%(datname)s.indexes.%(schemaname)s.%(relname)s." \
             "%(indexrelname)s.%(metric)s"
      multi_db = True
      query = """
          SELECT relname,
                 schemaname,
                 indexrelname,
                 idx_scan,
                 idx_tup_read,
                 idx_tup_fetch
          FROM pg_stat_user_indexes
      """
  
  
  class UserTableIOStats(QueryStats):
      path = "%(datname)s.tables.%(schemaname)s.%(relname)s.%(metric)s"
      multi_db = True
      query = """
          SELECT relname,
                 schemaname,
                 heap_blks_read,
                 heap_blks_hit,
                 idx_blks_read,
                 idx_blks_hit
          FROM pg_statio_user_tables
      """
  
  
  class UserIndexIOStats(QueryStats):
      path = "%(datname)s.indexes.%(schemaname)s.%(relname)s." \
             "%(indexrelname)s.%(metric)s"
      multi_db = True
      query = """
          SELECT relname,
                 schemaname,
                 indexrelname,
                 idx_blks_read,
                 idx_blks_hit
          FROM pg_statio_user_indexes
      """
  
  
  class ConnectionStateStats(QueryStats):
      path = "%(datname)s.connections.%(metric)s"
      multi_db = True
      query = """
          SELECT tmp.state AS key,COALESCE(count,0) FROM
                 (VALUES ('active'),
                         ('waiting'),
                         ('idle'),
                         ('idletransaction'),
                         ('unknown')
                  ) AS tmp(state)
          LEFT JOIN
               (SELECT CASE WHEN waiting THEN 'waiting'
                            WHEN current_query='<IDLE>' THEN 'idle'
                            WHEN current_query='<IDLE> in transaction'
                                THEN 'idletransaction'
                            WHEN current_query='<insufficient privilege>'
                                THEN 'unknown'
                            ELSE 'active' END AS state,
                       count(*) AS count
                 FROM pg_stat_activity
                 WHERE procpid != pg_backend_pid()
                 GROUP BY CASE WHEN waiting THEN 'waiting'
                               WHEN current_query='<IDLE>' THEN 'idle'
                               WHEN current_query='<IDLE> in transaction'
                                   THEN 'idletransaction'
                               WHEN current_query='<insufficient privilege>'
                                   THEN 'unknown' ELSE 'active' END
               ) AS tmp2
          ON tmp.state=tmp2.state ORDER BY 1
      """
  
  
  class LockStats(QueryStats):
      path = "%(datname)s.locks.%(metric)s"
      multi_db = False
      query = """
          SELECT lower(mode) AS key,
                 count(*) AS value
          FROM pg_locks
          WHERE database IS NOT NULL
          GROUP BY mode ORDER BY 1
      """
  
  
  class RelationSizeStats(QueryStats):
      path = "%(datname)s.sizes.%(schemaname)s.%(relname)s.%(metric)s"
      multi_db = True
      query = """
          SELECT pg_class.relname,
                 pg_namespace.nspname as schemaname,
                 pg_relation_size(pg_class.oid) as relsize
          FROM pg_class
          INNER JOIN
            pg_namespace
          ON pg_namespace.oid = pg_class.relnamespace
          WHERE reltype != 0
          AND relkind != 'S'
          AND nspname NOT IN ('pg_catalog', 'information_schema')
      """
  
  
  class BackgroundWriterStats(QueryStats):
      path = "bgwriter.%(metric)s"
      multi_db = False
      query = """
          SELECT checkpoints_timed,
                 checkpoints_req,
                 buffers_checkpoint,
                 buffers_clean,
                 maxwritten_clean,
                 buffers_backend,
                 buffers_backend_fsync,
                 buffers_alloc
          FROM pg_stat_bgwriter
      """
  
  
  class WalSegmentStats(QueryStats):
      path = "wals.%(metric)s"
      multi_db = False
      query = """
          SELECT count(*) AS segments
          FROM pg_ls_dir('pg_xlog') t(fn)
          WHERE fn ~ '^[0-9A-Z]{24}\$'
      """
  
  
  class TransactionCount(QueryStats):
      path = "transactions.%(metric)s"
      multi_db = False
      query = """
          SELECT 'commit' AS type,
                 sum(pg_stat_get_db_xact_commit(oid))
          FROM pg_database
          UNION ALL
          SELECT 'rollback',
                 sum(pg_stat_get_db_xact_rollback(oid))
          FROM pg_database
      """
  
  
  class IdleInTransactions(QueryStats):
      path = "%(datname)s.longest_running.%(metric)s"
      multi_db = True
      query = """
          SELECT 'idle_in_transaction',
                 max(COALESCE(ROUND(EXTRACT(epoch FROM now()-query_start)),0))
                     AS idle_in_transaction
          FROM pg_stat_activity
          WHERE current_query = '<IDLE> in transaction'
          GROUP BY 1
      """
  
  
  class LongestRunningQueries(QueryStats):
      path = "%(datname)s.longest_running.%(metric)s"
      multi_db = True
      query = """
          SELECT 'query',
              COALESCE(max(extract(epoch FROM CURRENT_TIMESTAMP-query_start)),0)
          FROM pg_stat_activity
          WHERE current_query NOT LIKE '<IDLE%'
          UNION ALL
          SELECT 'transaction',
              COALESCE(max(extract(epoch FROM CURRENT_TIMESTAMP-xact_start)),0)
          FROM pg_stat_activity
          WHERE 1=1
      """
  
  
  class UserConnectionCount(QueryStats):
      path = "%(datname)s.user_connections.%(metric)s"
      multi_db = True
      query = """
          SELECT usename,
                 count(*) as count
          FROM pg_stat_activity
          WHERE procpid != pg_backend_pid()
          GROUP BY usename
          ORDER BY 1
      """
  
  
  class DatabaseConnectionCount(QueryStats):
      path = "database.%(metric)s.connections"
      multi_db = False
      query = """
          SELECT datname,
                 count(datname) as connections
          FROM pg_stat_activity
          GROUP BY pg_stat_activity.datname
      """
  
  
  class TableScanStats(QueryStats):
      path = "%(datname)s.scans.%(metric)s"
      multi_db = True
      query = """
          SELECT 'relname' AS relname,
                 COALESCE(sum(seq_scan),0) AS sequential,
                 COALESCE(sum(idx_scan),0) AS index
          FROM pg_stat_user_tables
      """
  
  
  class TupleAccessStats(QueryStats):
      path = "%(datname)s.tuples.%(metric)s"
      multi_db = True
      query = """
          SELECT COALESCE(sum(seq_tup_read),0) AS seqread,
                 COALESCE(sum(idx_tup_fetch),0) AS idxfetch,
                 COALESCE(sum(n_tup_ins),0) AS inserted,
                 COALESCE(sum(n_tup_upd),0) AS updated,
                 COALESCE(sum(n_tup_del),0) AS deleted,
                 COALESCE(sum(n_tup_hot_upd),0) AS hotupdated
          FROM pg_stat_user_tables
      """
  
  
  registry = {
      'basic': (
          DatabaseStats,
          DatabaseConnectionCount,
      ),
      'extended': (
          DatabaseStats,
          DatabaseConnectionCount,
          UserTableStats,
          UserIndexStats,
          UserTableIOStats,
          UserIndexIOStats,
          ConnectionStateStats,
          LockStats,
          RelationSizeStats,
          BackgroundWriterStats,
          WalSegmentStats,
          TransactionCount,
          IdleInTransactions,
          LongestRunningQueries,
          UserConnectionCount,
          TableScanStats,
          TupleAccessStats,
      ),
  }
  