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", ) 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()