ampli/py/util.py

from argparse import ArgumentTypeError
import psycopg2 as pg
from configparser import ConfigParser
import pandas.io.sql as psql
import pandas as p
import datetime as dt
import numpy as np
from pprint import pprint

def datevalid(d):
    try:
        return dt.datetime.strptime(d, "%Y-%m-%d")
    except ValueError:
        raise ArgumentTypeError("Invaid date: {}".format(d))


def config(filename='database.ini', section='postgresql'):
    """Config parser from
    http://www.postgresqltutorial.com/postgresql-python/connect/"""
    # create a parser
    parser = ConfigParser()
    # read config file
    parser.read(filename)

    # get section, default to postgresql
    db = {}
    if parser.has_section(section):
        params = parser.items(section)
        for param in params:
            db[param[0]] = param[1]
    else:
        raise Exception(
            'Section {0} not found in the {1} file'.format(section, filename))

    return db


def getQuery(query, qparams=[], verbose = True):
    """
    Get single query
    """
    conn = False

    try:
        params = config()

        if verbose:
            print("Connecting to database")

        conn = pg.connect(**params)
        cur = conn.cursor()

        # Get table
        if verbose:
            print("Retrieving table")

        dataframe = psql.read_sql(query, conn, params=qparams)

        cur.close()

        if verbose:
            print("Table recieved")

        return dataframe

    except (Exception, pg.DatabaseError) as error:
        raise error

    finally:
        if conn is not False:
            conn.close()
            if verbose:
                print('Database connection closed')


def pickleQuery(query, path, qparams=[]):
    dq = getQuery(query, qparams)
    dq.to_pickle(path)
    print("Table pickled")


def gettemp(datestart, dateend, station, verbose = True):
    query = """
    SELECT record_no, station, temp_date, temp_date + temp_time AS temp_timestamp, tmax_c, tmin_c,
        tgmin, tmean, rhmean
    FROM weather.temperature_fact
    WHERE station = %(station)s AND
        temp_date >= to_date(%(datestart)s, 'yyyy-mm-dd') AND
        temp_date <  to_date(%(dateend)s,   'yyyy-mm-dd')
    ORDER BY temp_date, temp_time;
    """
    pdict = {
        'datestart': datestart,
        'dateend':   dateend,
        'station':   station
    }
    if verbose:
        print("Getting data with parameters:")
        pprint(pdict)
    qdf = getQuery(query, pdict, verbose)
    if verbose:
        print("converting")
    qdf.temp_date = p.to_datetime(qdf.temp_date)
    if verbose:
        print('Done')
    return qdf


if __name__ == "__main__":
    dv = getQuery('SELECT version()').version[0]
    print('PostgreSQL database version:')
    print(dv)