# coding=utf-8

"""
Collect metrics from postgresql

#### Dependencies

 * psycopg2

"""

import diamond.collector

try:
    import psycopg2
    import psycopg2.extras
except ImportError:
    psycopg2 = None


class PostgresqlCollector(diamond.collector.Collector):
    """
    PostgreSQL collector class
    """

    def get_default_config_help(self):
        """
        Return help text for collector
        """
        config_help = super(PostgresqlCollector, self).get_default_config_help()
        config_help.update(
            {
                "host": "Hostname",
                "dbname": "DB to connect to in order to get list of DBs in PgSQL",
                "user": "Username",
                "password": "Password",
                "port": "Port number",
                "password_provider": "Whether to auth with supplied password or"
                " .pgpass file  <password|pgpass>",
                "sslmode": "Whether to use SSL - <disable|allow|require|...>",
                "underscore": "Convert _ to .",
                "extended": "Enable collection of extended database stats.",
                "metrics": "List of enabled metrics to collect",
                "pg_version": "The version of postgres that you'll be monitoring"
                " eg. in format 9.2",
                "has_admin": "Admin privileges are required to execute some"
                " queries.",
            }
        )
        return config_help

    def get_default_config(self):
        """
        Return default config.
        """
        config = super(PostgresqlCollector, self).get_default_config()
        config.update(
            {
                "path": "postgres",
                "host": "localhost",
                "dbname": "postgres",
                "user": "postgres",
                "password": "postgres",
                "port": 5432,
                "password_provider": "password",
                "sslmode": "disable",
                "underscore": False,
                "extended": False,
                "metrics": [],
                "pg_version": 9.2,
                "has_admin": True,
            }
        )
        return config

    def collect(self):
        """
        Do pre-flight checks, get list of db names, collect metrics, publish
        """
        if psycopg2 is None:
            self.log.error("Unable to import module psycopg2")
            return {}

        # Get list of databases
        dbs = self._get_db_names()

        if len(dbs) == 0:
            self.log.error("I have 0 databases!")

            return {}

        if self.config["metrics"]:
            metrics = self.config["metrics"]
        elif diamond.collector.str_to_bool(self.config["extended"]):
            metrics = registry["extended"]

            if (
                diamond.collector.str_to_bool(self.config["has_admin"])
                and "WalSegmentStats" not in metrics
            ):
                metrics.append("WalSegmentStats")

        else:
            metrics = registry["basic"]

        # Iterate every QueryStats class
        for metric_name in set(metrics):
            if metric_name not in metrics_registry:
                self.log.error(
                    "metric_name %s not found in metric registry" % metric_name
                )
                continue

            for dbase in dbs:
                conn = self._connect(database=dbase)
                try:
                    klass = metrics_registry[metric_name]
                    stat = klass(dbase, conn, underscore=self.config["underscore"])
                    stat.fetch(self.config["pg_version"])
                    for metric, value in stat:
                        if value is not None:
                            self.publish(metric, 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 stat.multi_db is False:
                        break
                finally:
                    conn.close()

    def _get_db_names(self):
        """
        Try to get a list of db names
        """
        query = """
            SELECT datname FROM pg_database
            WHERE datallowconn AND NOT datistemplate
            AND NOT datname='postgres' AND NOT datname='rdsadmin' ORDER BY 1
        """
        conn = self._connect(self.config["dbname"])
        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):
        """
        Connect to given database
        """
        conn_args = {
            "host": self.config["host"],
            "user": self.config["user"],
            "password": self.config["password"],
            "port": self.config["port"],
            "sslmode": self.config["sslmode"],
        }

        if database:
            conn_args["database"] = database
        else:
            conn_args["database"] = "postgres"

        # libpq will use ~/.pgpass only if no password supplied
        if self.config["password_provider"] == "pgpass":
            del conn_args["password"]

        try:
            conn = psycopg2.connect(**conn_args)
        except Exception as e:
            self.log.error(e)
            raise e

        # Avoid using transactions, set isolation level to autocommit
        conn.set_isolation_level(0)
        return conn


class QueryStats(object):
    query = None
    path = None

    def __init__(self, dbname, conn, parameters=None, underscore=False):
        self.conn = conn
        self.dbname = dbname
        self.underscore = underscore
        self.parameters = parameters
        self.data = list()

    def _translate_datname(self, datname):
        """
        Replace '_' with '.'
        """
        if self.underscore:
            datname = datname.replace("_", ".")
        return datname

    def fetch(self, pg_version):
        if float(pg_version) >= 10.0 and hasattr(self, "post_100_query"):
            q = self.post_100_query
        elif float(pg_version) >= 9.6 and hasattr(self, "post_96_query"):
            q = self.post_96_query
        elif float(pg_version) >= 9.2 and hasattr(self, "post_92_query"):
            q = self.post_92_query
        else:
            q = self.query

        cursor = self.conn.cursor(cursor_factory=psycopg2.extras.DictCursor)
        try:
            cursor.execute(q, self.parameters)
            rows = cursor.fetchall()
            for row in rows:
                # If row is length 2, assume col1, col2 forms key: value
                if len(row) == 2:
                    self.data.append(
                        {
                            "datname": self._translate_datname(self.dbname),
                            "metric": row[0],
                            "value": row[1],
                        }
                    )

                # If row > length 2, assume each column name maps to
                # key => value
                else:
                    for key, value in iter(row.items()):
                        if key in (
                            "datname",
                            "schemaname",
                            "relname",
                            "indexrelname",
                            "funcname",
                        ):
                            continue

                        self.data.append(
                            {
                                "datname": self._translate_datname(
                                    row.get("datname", self.dbname)
                                ),
                                "schemaname": row.get("schemaname", None),
                                "relname": row.get("relname", None),
                                "indexrelname": row.get("indexrelname", None),
                                "funcname": row.get("funcname", None),
                                "metric": key,
                                "value": value,
                            }
                        )

        # Clean up
        finally:
            cursor.close()

    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
    post_92_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_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', 'rdsadmin')
    """
    query = post_92_query.replace(
        "pg_stat_database.temp_files as temp_files,", ""
    ).replace("pg_stat_database.temp_bytes as temp_bytes,", "")


class UserFunctionStats(QueryStats):
    # http://www.pateldenish.com/2010/11/postgresql-track-functions-to-tune.html
    path = "%(datname)s.functions.%(funcname)s.%(metric)s"
    multi_db = True
    query = """
        SELECT funcname,
               calls,
               total_time/calls as time_per_call
        FROM pg_stat_user_functions
        WHERE calls <> 0
    """


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
    """
    post_92_query = """
        SELECT tmp.mstate AS state,COALESCE(count,0) FROM
               (VALUES ('active'),
                       ('waiting'),
                       ('idle'),
                       ('idletransaction'),
                       ('unknown')
               ) AS tmp(mstate)
        LEFT JOIN
             (SELECT CASE WHEN waiting THEN 'waiting'
                          WHEN state = 'idle' THEN 'idle'
                          WHEN state LIKE 'idle in transaction%'
                              THEN 'idletransaction'
                          WHEN state = 'disabled'
                              THEN 'unknown'
                          WHEN query = '<insufficient privilege>'
                              THEN 'unknown'
                          ELSE 'active' END AS mstate,
                     count(*) AS count
               FROM pg_stat_activity
               WHERE pid != pg_backend_pid()
               GROUP BY CASE WHEN waiting THEN 'waiting'
                             WHEN state = 'idle' THEN 'idle'
                             WHEN state LIKE 'idle in transaction%'
                                 THEN 'idletransaction'
                             WHEN state = 'disabled'
                                 THEN 'unknown'
                             WHEN query = '<insufficient privilege>'
                                 THEN 'unknown'
                             ELSE 'active'
                        END
             ) AS tmp2
        ON tmp.mstate=tmp2.mstate ORDER BY 1
    """
    post_96_query = """
        SELECT tmp.state AS key,COALESCE(count,0) FROM
               (VALUES ('active'),
                       ('waiting'),
                       ('idle'),
                       ('idletransaction'),
                       ('unknown')
                ) AS tmp(state)
        LEFT JOIN
             (SELECT CASE WHEN wait_event IS NOT NULL THEN 'waiting'
                          WHEN state= 'idle' THEN 'idle'
                          WHEN state= 'idle in transaction'
                          THEN 'idletransaction'
                          WHEN state = 'active' THEN 'active'
                          ELSE 'unknown' END AS state,
                     count(*) AS count
               FROM pg_stat_activity
               WHERE pid != pg_backend_pid()
               GROUP BY CASE WHEN wait_event IS NOT NULL THEN 'waiting'
                          WHEN state= 'idle' THEN 'idle'
                          WHEN state= 'idle in transaction'
                          THEN 'idletransaction'
                          WHEN state = 'active' THEN 'active'
                          ELSE 'unknown' 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_alloc
        FROM pg_stat_bgwriter
    """


class WalSegmentStats(QueryStats):
    path = "wals.%(metric)s"
    multi_db = False
    post_100_query = """
        SELECT count(*) AS segments
        FROM pg_ls_dir('pg_wal') t(fn)
        WHERE fn ~ '^[0-9A-Z]{24}$'
    """
    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.idle_in_tranactions.%(metric)s"
    multi_db = True
    base_query = """
        SELECT 'idle_in_transactions',
               max(COALESCE(ROUND(EXTRACT(epoch FROM now()-query_start)),0))
                   AS idle_in_transaction
        FROM pg_stat_activity
        WHERE %s
        GROUP BY 1
    """
    query = base_query % ("current_query = '<IDLE> in transaction'",)
    post_92_query = base_query % ("state LIKE 'idle in transaction%'",)


class LongestRunningQueries(QueryStats):
    path = "%(datname)s.longest_running.%(metric)s"
    multi_db = True
    base_query = """
        SELECT 'query',
            COALESCE(max(extract(epoch FROM CURRENT_TIMESTAMP-query_start)),0)
        FROM pg_stat_activity
        WHERE %s
        UNION ALL
        SELECT 'transaction',
            COALESCE(max(extract(epoch FROM CURRENT_TIMESTAMP-xact_start)),0)
        FROM pg_stat_activity
        WHERE 1=1
    """
    query = base_query % ("current_query NOT LIKE '<IDLE%'",)
    post_92_query = base_query % ("state NOT LIKE 'idle%'",)


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
    """
    post_92_query = query.replace("procpid", "pid")


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
    """


class DatabaseReplicationStats(QueryStats):
    path = "database.replication.%(metric)s"
    multi_db = False
    query = """
        SELECT EXTRACT(epoch FROM
            current_timestamp - pg_last_xact_replay_timestamp()) as replay_lag
    """


class DatabaseXidAge(QueryStats):
    path = "%(datname)s.datfrozenxid.%(metric)s"
    multi_db = False
    query = """
        SELECT datname, age(datfrozenxid) AS age
        FROM pg_database WHERE datallowconn = TRUE
    """


metrics_registry = {
    "DatabaseStats": DatabaseStats,
    "DatabaseConnectionCount": DatabaseConnectionCount,
    "UserFunctionStats": UserFunctionStats,
    "UserTableStats": UserTableStats,
    "UserIndexStats": UserIndexStats,
    "UserTableIOStats": UserTableIOStats,
    "UserIndexIOStats": UserIndexIOStats,
    "ConnectionStateStats": ConnectionStateStats,
    "LockStats": LockStats,
    "RelationSizeStats": RelationSizeStats,
    "BackgroundWriterStats": BackgroundWriterStats,
    "WalSegmentStats": WalSegmentStats,
    "TransactionCount": TransactionCount,
    "IdleInTransactions": IdleInTransactions,
    "LongestRunningQueries": LongestRunningQueries,
    "UserConnectionCount": UserConnectionCount,
    "TableScanStats": TableScanStats,
    "TupleAccessStats": TupleAccessStats,
    "DatabaseReplicationStats": DatabaseReplicationStats,
    "DatabaseXidAge": DatabaseXidAge,
}

registry = {
    "basic": [
        "DatabaseStats",
        "DatabaseConnectionCount",
    ],
    "extended": [
        "DatabaseStats",
        "DatabaseConnectionCount",
        "DatabaseReplicationStats",
        "DatabaseXidAge",
        "UserFunctionStats",
        "UserTableStats",
        "UserIndexStats",
        "UserTableIOStats",
        "UserIndexIOStats",
        "ConnectionStateStats",
        "LockStats",
        "RelationSizeStats",
        "BackgroundWriterStats",
        "TransactionCount",
        "IdleInTransactions",
        "LongestRunningQueries",
        "UserConnectionCount",
        "TableScanStats",
        "TupleAccessStats",
    ],
}
