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

stationGPD = None


# Connect to DB "postgres" to check for database "temperatures_berteld_morstein"
def check_for_db_existence():
    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()
                    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():
    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 temperatures (id INTEGER, lon NUMERIC, lat NUMERIC, country TEXT);")
            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,...)
                columns = ",".join(df_columns)

                # 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('temperatures', columns, values)

                psycopg2.extras.execute_batch(cursor, insert_stmt, stationList.values)


def export(stationList):
    check_for_db_existence()
    insert_data(stationList)

    # return stationList