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"""

        schemes = dict()
        schemes[1] = (
            """
        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);
            """,
        )
        schemes[2] = (
            """DROP TABLE IF EXISTS links""",
            """
            CREATE TABLE links (
                id            SERIAL PRIMARY KEY,
                revision_id   INTEGER NOT NULL,
                project       VARCHAR(255) NOT NULL,
                package       VARCHAR(255) NOT NULL
            )
            """,
            "UPDATE scheme SET version=2",
        )
        schemes[3] = (
            """
            ALTER TABLE revisions ADD broken boolean NOT NULL DEFAULT(FALSE)
            """,
            "UPDATE scheme SET version=3",
        )
        schemes[4] = (
            """
            ALTER TABLE revisions ADD expanded_srcmd5 VARCHAR(255)
            """,
            "UPDATE scheme SET version=4",
        )
        schemes[5] = (
            """
            CREATE TABLE files (
                id              SERIAL PRIMARY KEY,
                revision_id     INTEGER NOT NULL,
                name            VARCHAR(255) NOT NULL,
                md5             VARCHAR(255) NOT NULL,
                size            INTEGER NOT NULL,
                mtime           INTEGER NOT NULL
            )
            """,
            "UPDATE scheme SET version=5",
        )
        schemes[6] = (
            """
            CREATE TABLE requests (
                id              SERIAL PRIMARY KEY,
                number          INTEGER NOT NULL,
                revision_id     INTEGER NOT NULL,
                creator         VARCHAR(255) NOT NULL,
                type            VARCHAR(10) NOT NULL,
                state           VARCHAR(10) NOT NULL,
                source_package  VARCHAR(255),
                source_project  VARCHAR(255)
            )
            """,
            "UPDATE scheme SET version=6",
        )
        schemes[7] = (
            "ALTER TABLE requests DROP COLUMN revision_id",
            "UPDATE scheme SET version=7",
        )
        schemes[8] = (
            "ALTER TABLE requests ADD COLUMN source_rev VARCHAR(40)",
            "UPDATE scheme SET version=8",
        )
        schemes[9] = (
            "ALTER TABLE revisions ADD COLUMN request_number INTEGER",
            "UPDATE revisions SET request_number=requestid",
            "ALTER TABLE revisions DROP COLUMN requestid",
            "UPDATE scheme SET version=9",
        )
        schemes[10] = (
            "ALTER TABLE revisions ADD COLUMN request_id INTEGER",
            """ALTER TABLE revisions 
               ADD CONSTRAINT request_id_foreign_key
               FOREIGN KEY (request_id) 
               REFERENCES requests (id)""",
            "UPDATE scheme SET version=10",
        )
        schemes[11] = (
            "create index request_number_idx on revisions (request_number);",
            "UPDATE scheme SET version=11",
        )
        schema_version = self.schema_version()
        if (schema_version + 1) not in schemes:
            return
        try:
            cur = self.conn.cursor()
            # create table one by one
            for version, commands in schemes.items():
                if version <= schema_version:
                    continue
                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()