Minimal DB abstraction layer

database.py

# coding: utf-8
# [email protected]
# Mon Jul 27 20:04:26 CEST 2015
# pylint: disable-msg=C0103
# pylint: disable-msg=C0301
# pylint: disable-msg=W0611
# pylint: disable-msg=W0612
# pylint: disable-msg=W0702
# pylint: disable-msg=W0703
# pylint: disable-msg=W0621
# pylint: disable-msg=R0913
# vim:ts=4 sts=4 tw=100:
"""
Implements the database layer via sqlitedb and mysqldb objects
"""

import mysql.connector as mysql
import sqlite3
import sys

class sqlitedb(object):

    def __init__(self, db='vuln.db'):
        self.db = db
        self.con = None
        self.online = False

    def connect(self):
        if (self.online):
            print('Already connected')
        else:
            try:
                self.con = sqlite3.connect(self.db)
            except sqlite3.Error as e:
                print ("Error {}".format(e.msg))
                raise(e)
            else:
                self.online = True

    def execute(self, query=None):
        if not (self.online):
            print('Not connected')
        else:
            try:
                c = self.con.cursor()
                c.execute(query)
                self.con.commit()
            except sqlite3.Error as e:
                print ("Error {}".format(e.msg))
                raise(e)

    def close(self):
        if not (self.online):
            print('Not connected')
        else:
            if self.con:
                self.con.close()
                self.online = False

    def __test__(self):
        self.db = sqlitedb('test.db')
        self.db.connect()
        self.db.execute("CREATE TABLE person (id INTEGER PRIMARY KEY ASC, name varchar(250) NOT NULL)")
        self.db.execute("""
                  CREATE TABLE address
                  (id INTEGER PRIMARY KEY ASC, street_name varchar(250), street_number varchar(250),
                   post_code varchar(250) NOT NULL, person_id INTEGER NOT NULL,
                   FOREIGN KEY(person_id) REFERENCES person(id))
                  """)
        self.db.execute("INSERT INTO person VALUES(1, 'pythoncentral')")
        self.db.execute("INSERT INTO address VALUES(1, 'python road', '1', '00000', 1)")
        self.db.close()



class mysqldb(object):

    def __init__(self, host=None, user=None, dbname=None, password=None):
        self.con = None
        self.host = host
        self.user = user
        self.dbname = dbname
        self.password = password

    def connect(self):
        try:
            self.con = mysql.connect(host=self.host, user=self.user,
                password=self.password, db=self.dbname)
        except mysql.Error as e:
            print ("Error {}".format(e.msg))
            raise(e)

    def execute(self, query):
        try:
            cur = self.con.cursor()
            #cur.query(query)
            cur.execute(query)
            rows = cur.fetchall()
            n_rows = int(cur.rowcount())
            return rows
        except mysql.Error as e:
            print ("Error %d: %s" % (e.args[0],e.args[1]))
            raise (e)

    def close(self):
        if self.con:
            self.con.close()

    def __test__(self):
        host='localhost'
        user='testuser'
        password='password'
        dbname='testdb'
        self.connect()
        rows = self.execute('SELECT VERSION()')
        for row in rows: print (row)
        self.close()