database.py

database.py

import psycopg2 as pg
from psycopg2.extensions import ISOLATION_LEVEL_AUTOCOMMIT


class Database:
    def __init__(self, host='127.0.0.1', port='5432', username='postgres', password='', database=''):
        self.host = host
        self.port = port
        self.username = username
        self.password = password
        self.database = database

    def connect(self, database=True):
        try:
            if database:
                con = pg.connect(user=self.username, password=self.password, host=self.host, port=self.port,
                                 database=self.database)
            else:
                con = pg.connect(user=self.username, password=self.password, host=self.host, port=self.port)
            print('connected to database')
            return con
        except pg.OperationalError as e:
            if 'does not exist' in str(e):
                print('Database "{}" dos not exist'.format(self.database))
                print('Do you want to create this database? "y/n')
                choice = input()
                if choice in ('y', 'Y'):
                    self.create(con=self.connect(False), database_name=self.database)
                else:
                    exit('unable connect to database. {}'.format(e))

    @staticmethod
    def create(con, database_name):
        cur = con.cursor()
        con.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
        try:
            cur.execute("""CREATE DATABASE {}""".format(database_name))
            print('database {} successfully created'.format(database_name))
        except pg.ProgrammingError as e:
            print('error - {}'.format(e))
        finally:
            cur.close()

    def create_table(self,):