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)