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() def import_rev(self, revision): cur = self.conn.cursor() cur.execute( """INSERT INTO revisions (project, package, rev, unexpanded_srcmd5, commit_time, userid, comment, requestid) VALUES(%s, %s, %s, %s, %s, %s, %s, %s)""", ( revision.project, revision.package, revision.rev, revision.unexpanded_srcmd5, revision.time, revision.userid, revision.comment, revision.requestid, ), ) if __name__ == "__main__": db = DB() db.create_tables()