git-importer/lib/db.py

import psycopg2

from lib.config import config


class DB:
    def __init__(self, section="production"):
        self.config_section = section
        self.connect()
        self.create_tables()

    def connect(self):
        try:
            # read the connection parameters
            params = config(section=self.config_section)
            # 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)",
            """ DROP TABLE IF EXISTS revisions""",
            """ CREATE TABLE revisions (
                id                  SERIAL PRIMARY KEY,
                project             VARCHAR(255) NOT NULL,
                package             VARCHAR(255) NOT NULL,
                rev                 INTEGER      NOT NULL,
                unexpanded_srcmd5   VARCHAR(255) NOT NULL,
                commit_time         timestamp    NOT NULL,
                userid              VARCHAR(255) NOT NULL,
                comment             TEXT,
                requestid           INTEGER
                )
            """,
            """
            CREATE UNIQUE INDEX ppr ON revisions (project, package, rev);
            """,
        )

        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

    def cursor(self):
        return self.conn.cursor()


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