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 pandas as pd
# import numpy as np

stationGPD = None


# Connect to DB "postgres" to check for database "temperatures_berteld_morstein"
def check_for_db_existence(stationList):
    print("Checking for database existence")
    with psycopg2.connect(database='postgres', user='postgres', password='postgres', host='localhost', port=5432) as connection:
        with connection.cursor() as cursor:
            cursor.execute("SELECT datname FROM pg_database WHERE datname LIKE 'temperatures_berteld_morstein';")

            try:
                db_exists = cursor.fetchall()[0]
                print('DB existing')

            except IndexError:  # DB temperatures_berteld_morstein doesn't exist
                print('DB not existing')
                try:
                    create_db(connection, cursor)
                    create_table(stationList)
                    print('Successfully created database and table')
                except (Exception, psycopg2.DatabaseError) as error:
                    print(error)


# Use existing connection to DB "postgres" to create DB "temperatures_berteld_morstein"
def create_db(connection, cursor):
    connection.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)  # Needs to be in AUTOCOMMIT mode for creating database
    create_db_query = sql.SQL("CREATE DATABASE temperatures_berteld_morstein;")
    print(create_db_query)
    cursor.execute(create_db_query)


# Connect to DB "temperatures_berteld_morstein" to create table "temperatures"
def create_table(stationList):
    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='temperatures_berteld_morstein', user='postgres', password='postgres', host='localhost', port=5432) 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):
    with psycopg2.connect(database='temperatures_berteld_morstein', user='postgres', password='postgres', host='localhost', port=5432) as connection:
        with connection.cursor() as cursor:

            if len(stationList) > 0:
                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)
    insert_data(stationList)

    # return stationList