testDB.py 4.25 KB
#!/usr/bin/python
# -*- coding: utf-8 -*-

import logging
import traceback
import MySQLdb, json
from common.testException import TestDBException

logger = logging.getLogger("db_logger")

class TestDB():
    db = None
    def __init__(self):
        try:
            self.db = MySQLdb.connect(host='localhost', user='root', passwd='PASSWORD', db='mydb', port=3306, charset='utf8')
        except Exception as e:
            logger.error(traceback.format_exc())
            raise TestDBException(-1, "TestDBException: cannot connect TestDB(%s)"%(e[1]))

    def commitDB(self):
        try:
            self.db.commit()
        except Exception as e:
            logger.error(traceback.format_exc())
            raise TestDBException(-1, "TestDBException: cannot commit db(%s)"%(e[1]))
    def rollbackDB(self):
        try:
            self.db.rollback()
        except Exception as e:
            logger.error(traceback.format_exc())
            raise TestDBException(-1, "TestDBException: cannot rollback db(%s)"%(e[1]))
    def closeDB(self):
        try:
            self.db.close()
        except Exception as e:
            logger.error(traceback.format_exc())
            raise TestDBException(-1, "TestDBException: cannot close db connection(%s)"%(e[1]))

    def getDictCursor(self):
        try:
            return self.db.cursor(MySQLdb.cursors.DictCursor)
        except Exception as e:
            logger.error(traceback.format_exc())
            raise TestDBException(-1, "TestDBException: cannot get mysql dict cursor(%s)"%(e[1]))

    def getCommonCursor(self):
        try:
            return self.db.cursor()
        except Exception as e:
            logger.error(traceback.format_exc())
            raise TestDBException(-1, "TestDBException: cannot get mysql common cursor(%s)"%(e[1]))

    def executeQuery(self, cursor, query):
        try:
            logger.info(query)
            row = cursor.execute(query)
            return row
        except Exception as e:
            logger.error(traceback.format_exc())
            logger.error(query)
            raise TestDBException(-1, "TestDBException: cannot execute query(%s)" %(query))

    def executeSelectQuery(self, query, one=False):
        try:
            cursor = self.getDictCursor()
            
            self.executeQuery(cursor, query)

            if one:
                res = cursor.fetchone()
            else:
                res = cursor.fetchall()

            self.closeDbCursor(cursor)

            return res

        except TestDBException as e:
            raise e
        except Exception as e:
            print(traceback.format_exc())
            raise TestDBException(-1, "TestDBException: executeSelectQuery()")

    def closeDbCursor(self, cursor):
        try:
            cursor.close()
        except Exception as e:
            logger.error(traceback.format_exc())
            raise TestDBException(-1, "TestDBException: cannot close db cursor(%s)"%(e[1]))

    # insert api call log function
    def insertIntoApiCallLog(self, url, ip, method, response_time, result):
        try:
            _id = None

            cursor = self.getDictCursor()
            query = ("""INSERT INTO api_call_log (url, ip, method, ret_code, ret_message, response_time) VALUES ('%s', '%s', '%s', '%s', '%s', '%s');"""
                % (url, ip, method, result.get('code'), result.get('message'), response_time))
            row = self.executeQuery(cursor, query)
            
            if row > 0:
                _id = cursor.lastrowid
                self.commitDB()

            self.closeDbCursor(cursor)
            return _id

        except Exception as e:
            print(traceback.format_exc())
            logger.error(traceback.format_exc())
            raise TestDBException(-1, "TestDBException: insertIntoApiCallLog()")

    def selectApiCallLog(self):
        try:
            _id = None

            cursor = self.getDictCursor()
            query = """SELECT * FROM api_call_log;"""

            self.executeQuery(cursor, query)

            logList = cursor.fetchall()

            self.closeDbCursor(cursor)

            return logList

        except Exception as e:
            print(traceback.format_exc())
            logger.error(traceback.format_exc())
            raise TestDBException(-1, "TestDBException: selectApiCallLog()")