"""
Datadog Postgres Query Translator

This script will translate the results of any Postgres into Datadog metrics
and submit them to Datadog.
"""
import psycopg2

try:
    from checks import AgentCheck
except ImportError:
    from datadog_checks.checks import AgentCheck

HOSTNAME = 'localhost'
USERNAME = 'postgres'
PASSWORD = 'postgres'
DATABASE = 'postgres'
QUERY = 'SELECT * FROM CITY'

def execute_postgres_query(query):
    connection = psycopg2.connect(host=HOSTNAME, database=DATABASE,
                                  user=USERNAME, password=PASSWORD)

    cursor = connection.cursor()
    cursor.execute(query)
    column_name_list = [description[0] for description in cursor.description]
    row_list = cursor.fetchall()
    cursor.close()
    connection.close()

    return row_list, column_name_list


class PostgresQueryCheck(AgentCheck):
    def check(self, instance):
        row_list, column_name_list = execute_postgres_query(QUERY)
        for row_entry in row_list:
            this_row_dict = {column_name: row_entry[index]
                             for index, column_name
                             in enumerate(column_name_list)}

            self.gauge(
                'postgres.custom.population',
                this_row_dict['population'],
                tags=['countrycode:{}'.format(this_row_dict['countrycode']),
                      'district:{}'.format(this_row_dict['district'])]
            )