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, project VARCHAR(255) NOT NULL, package VARCHAR(255) NOT NULL ) """, """ UPDATE scheme SET version=2 """, ) 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()