import logging import psycopg2 from psycopg2.extras import LoggingConnection 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(connection_factory=LoggingConnection, **params) logger = logging.getLogger(__name__) self.conn.initialize(logger) except (Exception, psycopg2.DatabaseError) as error: print(error) raise error def schema_version(self): # create a cursor with self.conn.cursor() as cur: # 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() 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] = ( """DROP TABLE IF EXISTS files""", """ 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] = ( """DROP TABLE IF EXISTS requests""", """ 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", ) schemes[12] = ( "create index request_number_idx2 on requests(number)", "UPDATE scheme SET version=12", ) schemes[13] = ( "DROP TABLE IF EXISTS users", """CREATE TABLE users ( id SERIAL PRIMARY KEY, userid VARCHAR(255) NOT NULL, email VARCHAR(255) NOT NULL, realname VARCHAR(255) NOT NULL ) """, "UPDATE scheme SET version=13", ) schemes[14] = ( "DROP TABLE IF EXISTS linked_revs", """ CREATE TABLE linked_revs ( id SERIAL PRIMARY KEY, revision_id INTEGER NOT NULL, linked_id INTEGER NOT NULL ) """, "UPDATE scheme SET version=14", ) schemes[15] = ( "ALTER TABLE revisions ALTER COLUMN rev TYPE real USING rev::real", "UPDATE scheme SET version=15", ) schemes[16] = ( """DROP TABLE IF EXISTS fake_revs""", """ CREATE TABLE fake_revs ( id SERIAL PRIMARY KEY, revision_id INTEGER NOT NULL, linked_id INTEGER NOT NULL ) """, "create index revs_linked on fake_revs (revision_id,linked_id)", "UPDATE scheme SET version=16", ) schemes[17] = ( "ALTER TABLE revisions ADD COLUMN files_hash VARCHAR(40)", "UPDATE scheme SET version=17", ) schemes[18] = ( "ALTER TABLE linked_revs ADD COLUMN considered BOOLEAN DEFAULT FALSE", "UPDATE scheme SET version=18", ) schemes[19] = ( "CREATE INDEX ON files(revision_id)", "UPDATE scheme SET version=19", ) schemes[20] = ( "CREATE INDEX ON linked_revs(linked_id)", "CREATE INDEX ON linked_revs(considered)", "UPDATE scheme SET version=20", ) schema_version = self.schema_version() if (schema_version + 1) not in schemes: return try: with self.conn.cursor() as cur: # create table one by one for version, commands in schemes.items(): if version <= schema_version: continue for command in commands: cur.execute(command) # 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()