ExportToDatabase.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Wed Jul 14 13:43:45 2021

export the stationlist to database

@author: geopeter
"""
import psycopg2
import psycopg2.extras
from psycopg2 import sql
from psycopg2.extensions import ISOLATION_LEVEL_AUTOCOMMIT
import configparser

cfg = configparser.ConfigParser()
cfg.read('config.ini')

assert "POSTGRES" in cfg, "missing POSTGRES in config.ini"

param_postgres = cfg["POSTGRES"]

stationGPD = None

# Use existing connection to DB "postgres" to create DB "temperatures_berteld_morstein"
def create_db(dbName):
    print("Create DB: ",dbName)
    connection = psycopg2.connect(dbname='postgres', user=param_postgres["user"], password=param_postgres["password"], host=param_postgres["host"], port=param_postgres["port"])
    connection.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)  # Needs to be in AUTOCOMMIT mode for creating database
    with connection.cursor() as cursor:
        create_db_query = sql.SQL("""CREATE DATABASE {};""".format(dbName))
        cursor.execute(create_db_query)
        connection.close()
        
def drop_db(dbName):
    if dbexists(dbName):
        print("Drop DB: ",dbName)
        try:
            connection = psycopg2.connect(dbname='postgres', user=param_postgres["user"], password=param_postgres["password"], host=param_postgres["host"], port=param_postgres["port"])
            connection.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT) # Needs to be in AUTOCOMMIT mode for creating database
            with connection.cursor() as cursor:
                cursor.execute("DROP DATABASE {};".format(dbName));
                connection.close()
        except(psycopg2.DatabaseError) as error:
            # do nothing, because test db is clean
            print(error.message)
            return

def dbexists(dbName):
        try:
            connection = psycopg2.connect(dbname='postgres', user=param_postgres["user"], password=param_postgres["password"], host=param_postgres["host"], port=param_postgres["port"])
            cursor = connection.cursor()
            cursor.execute("SELECT datname FROM pg_database WHERE datname LIKE '{}';".format(dbName))
            
            db_exists = cursor.fetchall()[0]
            connection.close()
            return True
        except(Exception, psycopg2.DatabaseError) as error:
            # do nothing, because test db is clean
            return False


# Connect to DB "postgres" to check for database "temperatures_berteld_morstein"
def check_for_db_existence(stationList, dbName):
    print("Checking for database existence")
    if dbexists(dbName):
        print('DB existing exists')
    else: 
         create_db(dbName)
         create_table(stationList, dbName)

# Connect to DB "temperatures_berteld_morstein" to create table "temperatures"
def create_table(stationList, dbName):
    df_columns = list(stationList)
    columns = ['id INTEGER', 'lon NUMERIC', 'lat NUMERIC', 'country TEXT', 'file TEXT']
    for column in df_columns:
        if str(column).startswith('19') or str(column).startswith('20'):
            columns.append('"{}" NUMERIC'.format(column))
    columns_clean = str(columns).strip('[]').replace("'", "")

    with psycopg2.connect(database=dbName, user=param_postgres["user"], password=param_postgres["password"], host=param_postgres["host"], port=param_postgres["port"]) as connection:
        with connection.cursor() as cursor:
            query = sql.SQL("""CREATE TABLE stations ({});""".format(columns_clean))
            #print(query)
            cursor.execute(query)


def insert_data(stationList, dbName):
    with psycopg2.connect(database=dbName, user=param_postgres["user"], password=param_postgres["password"], host=param_postgres["host"], port=param_postgres["port"]) as connection:
        with connection.cursor() as cursor:

            if len(stationList) > 0:
                # print(stationList)
                cursor.execute("DELETE FROM stations;")

                df_columns = list(stationList)
                # create (col1,col2,...)

                # As integers like 2018, 2017, etc. are not possible as column names, double quotes have to be added. This requires some tricks and cleanups
                columns = []
                for column in df_columns:
                    columns.append('"' + column + '"')
                columns = str(columns).replace('[', '').replace(']', '').replace("'", "").replace('\n', '').replace(' ', '')
                stationList = stationList.round(decimals=3)

                # create VALUES('%s', '%s",...) one '%s' per column
                values = "VALUES({})".format(",".join(["%s" for _ in df_columns]))

                # create INSERT INTO table (columns) VALUES('%s',...)
                insert_stmt = """INSERT INTO {} ({}) {}""".format('stations', columns, values)
                psycopg2.extras.execute_batch(cursor, insert_stmt, stationList.values)

def export(stationList):
    check_for_db_existence(stationList, param_postgres['dbName'])
    insert_data(stationList, param_postgres['dbName'])