git-importer/db.py

import psycopg2
from config import config


class DB:
    def __init__(self):
        self.connect()

    def connect(self):
        try:
            # read the connection parameters
            params = config()
            # connect to the PostgreSQL server
            self.conn = psycopg2.connect(**params)

        except (Exception, psycopg2.DatabaseError) as error:
            print(error)
            raise error

    def schema_version(self):
        # create a cursor
        cur = self.conn.cursor()

        # execute a statement
        try:
            cur.execute("SELECT MAX(version) from scheme")
        except psycopg2.errors.UndefinedTable as error:
            cur.close()
            self.close()
            self.connect()
            return 0

        db_version = cur.fetchone()

        cur.close()
        return db_version[0]

    def close(self):
        if self.conn is not None:
            self.conn.close()
        self.conn = None

    def create_tables(self):
        """Create the tables if not existing - assumes connected"""

        commands = (
            """
        CREATE TABLE scheme (
            id SERIAL PRIMARY KEY,
            version SMALLINT NOT NULL
        )
        """,
            "INSERT INTO scheme (version) VALUES(1)",
            """ CREATE TABLE revisions (
                id SERIAL PRIMARY KEY,
                package VARCHAR(255) NOT NULL
                )
        """,
        )

        if self.schema_version() > 0:
            return
        try:
            cur = self.conn.cursor()
            # create table one by one
            for command in commands:
                cur.execute(command)
            # close communication with the PostgreSQL database server
            cur.close()
            # commit the changes
            self.conn.commit()
        except (Exception, psycopg2.DatabaseError) as error:
            print(error)
            self.close()
            raise error


if __name__ == "__main__":
    db = DB()
    db.create_tables()