chromaDB/text_analysis/tools/mysql_helper.py

# coding:utf8
import os, sys
cur_dir = os.path.dirname(os.path.abspath(__file__)) or os.getcwd()
par_dir = os.path.abspath(os.path.join(cur_dir, os.path.pardir))
sys.path.append(cur_dir)
sys.path.append(par_dir)
import json
import re
# from log_util.set_logger import set_logger
# logging = set_logger('logs/error.log')
import pymysql.cursors
import traceback

def mysqlConn(data,logging):
    res={"successCode":"1","errorLog":"","results":""}
    p_host=data["Host"]
    p_port=int(data["Port"])
    p_db=data["Database"]
    p_user=data["User"]
    p_password=data["Password"]
    try:
        db = pymysql.connect(host=p_host, user=p_user, passwd=p_password, db=p_db, port=p_port,
                         charset='utf8', cursorclass=pymysql.cursors.DictCursor)
        db.ping(reconnect=True)
        cursor = db.cursor()
        sql = "SHOW TABLES"
        cursor.execute(sql)
        tables = cursor.fetchall()
        if tables:
            table_names = list(map(lambda x: list(x.values())[0], tables))
            res["results"] = table_names
        else:
            res["successCode"] = "0"
        cursor.close()
        db.close()
        return res
    except:
        res["successCode"] = "0"
        res["errorLog"]=traceback.format_exc()
        logging.error(traceback.format_exc())
        return res

def getTableColumnNames(data,logging):
    res={"successCode":"1","errorLog":"","results":""}
    p_host=data["Host"]
    p_port=int(data["Port"])
    p_db=data["Database"]
    p_user=data["User"]
    p_password=data["Password"]
    p_table=data["Table"]
    try:
        db = pymysql.connect(host=p_host, user=p_user, passwd=p_password, db=p_db, port=p_port,
                         charset='utf8', cursorclass=pymysql.cursors.DictCursor)
        db.ping(reconnect=True)
        cursor = db.cursor()
        sql = "DESCRIBE "+p_table
        cursor.execute(sql)
        tables = cursor.fetchall()
        if tables:
            table_names = list(map(lambda x: x['Field'], tables))
            res["results"] = table_names
        else:
            res["successCode"] = "0"
        cursor.close()
        db.close()
        return res
    except:
        res["successCode"] = "0"
        res["errorLog"]=traceback.format_exc()
        logging.error(traceback.format_exc())
        return res

def mysqlInsert(input,logging):
    res={"successCode":"1","errorLog":"","results":""}
    data=input["metadata"]["admin"]
    p_host=data["Host"]
    p_port=int(data["Port"])
    p_db=data["Database"]
    p_user=data["User"]
    p_password=data["Password"]
    p_table=data["Table"]
    p_columnName=data["columnName"]
    cN='('+','.join(p_columnName)+') '
    p_values=data["values"]
    val=tuple(p_values)
    try:
        db = pymysql.connect(host=p_host, user=p_user, passwd=p_password, db=p_db, port=p_port,
                         charset='utf8', cursorclass=pymysql.cursors.DictCursor)
        db.ping(reconnect=True)
        cursor = db.cursor()
        sql = "insert into " + p_table + cN + "values ("+ ','.join(['%s'] * len(val)) + ")"
        cursor.execute(sql,val)
        db.commit()
        cursor.close()
        db.close()
        return res
    except:
        res["successCode"] = "0"
        res["errorLog"]=traceback.format_exc()
        logging.error(traceback.format_exc())
        return res

def mysqlUpdate(input,logging):
    res={"successCode":"1","errorLog":"","results":""}
    data=input["metadata"]["admin"]
    p_host=data["Host"]
    p_port=int(data["Port"])
    p_db=data["Database"]
    p_user=data["User"]
    p_password=data["Password"]
    p_table=data["Table"]
    # p_set=data["Set"]
    p_set=get_updateSet(input)
    # where=process_where(data["Filter"])
    where=get_filter(data["Filter"])
    try:
        db = pymysql.connect(host=p_host, user=p_user, passwd=p_password, db=p_db, port=p_port,
                         charset='utf8', cursorclass=pymysql.cursors.DictCursor)
        db.ping(reconnect=True)
        cursor = db.cursor()
        sql = "UPDATE " + p_table + p_set + where
        print(sql)
        cursor.execute(sql)
        db.commit()
        cursor.close()
        db.close()
        return res
    except:
        res["successCode"] = "0"
        res["errorLog"]=traceback.format_exc()
        logging.error(traceback.format_exc())
        return res

def mysqlExecute(input,logging):
    res={"successCode":"1","errorLog":"","results":""}
    data=input["metadata"]["admin"]
    p_host=data["Host"]
    p_port=int(data["Port"])
    p_db=data["Database"]
    p_user=data["User"]
    p_password=data["Password"]
    execute=data["Execute"]
    try:
        db = pymysql.connect(host=p_host, user=p_user, passwd=p_password, db=p_db, port=p_port,
                         charset='utf8', cursorclass=pymysql.cursors.DictCursor)
        db.ping(reconnect=True)
        cursor = db.cursor()
        cursor.execute(execute)
        if 'select' in execute.lower():
            result = cursor.fetchall()
            res["results"]=json.dumps(result,ensure_ascii=False)
        else:
            db.commit()
        cursor.close()
        db.close()
        return res
    except:
        res["successCode"] = "0"
        res["errorLog"]=traceback.format_exc()
        logging.error(traceback.format_exc())
        return res

# def process_where(data):
#     '''
#     组装where
#     :param data: data["Filter"],{"key":"age","value":"20","operator":">"},{"logicalSymbol":"and"},{"key":"weight","value":"50","operator":"<"}
#     :return: WHERE age>20 and weight<50
#     '''
#     if data=="" or data==[]:
#         return ""
#     where = " WHERE "
#     for line in data:
#         if "key" in line.keys():
#             val = line["value"]
#             if isinstance(val, str):
#                 val = "\'" + val + "\'"
#             tmp = str(line["key"]) + " " + line["operator"] + " " + str(val)
#             where += tmp
#         else:
#             where += " " + line["logicalSymbol"] + " "
#     return where
#
# def process_filter(data):
#     '''
#     组装key,value,operator
#     :param data: data["Filter"],{"key":"age",value:"20","operator":"="}
#     :return: age=20
#     '''
#     if data=="" or data==[]:
#         return ""
#     res=data["key"]+" "+data["operator"]+" "+data["value"]
#     return res

def get_updateSet(input):
    metadata=input["metadata"]
    user=metadata["user"]
    sets=metadata["admin"]["Set"]
    res=[]
    for line in sets:
        part=line.split("=")
        tmp = []
        for p in part:
            user_match=re.findall('##(.*?)##', p)
            if user_match!=[]:
                tmp.append(user[user_match[0]])
        res.append(str(tmp[0])+"="+str(tmp[1]))
    result=" SET "+",".join(res)
    return result

def get_filter(data):
    if "OR" not in data.keys():
        return ""
    op_or=data["OR"]
    res = ""
    if len(op_or) == 1:
        tmp = []
        line = op_or[0]["AND"]
        for single_line in line:
            val = single_line["value"]
            if isinstance(val, str):
                val = "\'" + val + "\'"
            tmp.append(str(single_line["key"]) + single_line["operator"] + str(val))
            if single_line != line[-1]:
                tmp.append("and")
        res = " WHERE "+" ".join(tmp)
    elif len(op_or) > 1:
        tmp = []
        for single_and in op_or:
            line = single_and["AND"]
            for sigle_line in line:
                val = sigle_line["value"]
                if isinstance(val, str):
                    val = "\'" + val + "\'"
                tmp.append(str(sigle_line["key"]) + sigle_line["operator"] + str(val))
                if sigle_line != line[-1]:
                    tmp.append("and")
            if single_and != op_or[-1]:
                tmp.append("or")
        res = " WHERE "+" ".join(tmp)
    return res


def mysqlQuery(input,logging):
    res={"successCode":"1","errorLog":"","results":""}
    data=input["metadata"]["admin"]
    p_host=data["Host"]
    p_port=int(data["Port"])
    p_db=data["Database"]
    p_user=data["User"]
    p_password=data["Password"]
    p_table=data["Table"]
    p_columnNames=data["columnNames"]
    # p_filter=data["Filter"]
    column='*'
    if len(p_columnNames)==1:
        column=p_columnNames[0]
    elif len(p_columnNames)>1:
        column=','.join(p_columnNames)
    where=get_filter(data["Filter"])
    try:
        db = pymysql.connect(host=p_host, user=p_user, passwd=p_password, db=p_db, port=p_port,
                         charset='utf8', cursorclass=pymysql.cursors.DictCursor)
        db.ping(reconnect=True)
        cursor = db.cursor()
        sql = "SELECT " + column +" From "+ p_table + where
        # print(sql)
        cursor.execute(sql)
        result = cursor.fetchall()
        res["results"]=json.dumps(result,ensure_ascii=False)
        cursor.close()
        db.close()
        return res
    except:
        res["successCode"] = "0"
        res["errorLog"]=traceback.format_exc()
        logging.error(traceback.format_exc())
        return res

def mysqlDelete(input,logging):
    res={"successCode":"1","errorLog":"","results":""}
    data=input["metadata"]["admin"]
    p_host=data["Host"]
    p_port=int(data["Port"])
    p_db=data["Database"]
    p_user=data["User"]
    p_password=data["Password"]
    p_table=data["Table"]
    # where=process_where(data["Filter"])
    where=get_filter(data["Filter"])
    try:
        db = pymysql.connect(host=p_host, user=p_user, passwd=p_password, db=p_db, port=p_port,
                         charset='utf8', cursorclass=pymysql.cursors.DictCursor)
        db.ping(reconnect=True)
        cursor = db.cursor()
        sql = "DELETE From "+ p_table + where
        cursor.execute(sql)
        db.commit()
        cursor.close()
        db.close()
        return res
    except:
        res["successCode"] = "0"
        res["errorLog"]=traceback.format_exc()
        logging.error(traceback.format_exc())
        return res


if __name__=="__main__":
    input={"metadata":{"admin":{
			"type":"query",
			"Table":"student",
			"columnNames":["name","age"],
			"Set":["##tag1##=##value1##","##tag2##=##value2##"],
			"Filter":{
				"OR":[
						{
							"AND":[{"key":"age","value":20,"operator":">"},{"key":"weight","value":50,"operator":"<"}]
						},
						{
							"AND":[{"key":"name","value":"ff","operator":"="}]
						}
					]
			},
			"Host":"172.26.28.30",
			"Port":"3306",
			"Database":"test",
			"User":"crawl",
			"Password":"crawl123"
        }},
        "user": {
            "tag1": "age",
            "tag2": "weight",
            "value1": 2,
            "value2": 100
        }
    }
    res=mysqlUpdate(input,"")
    print(res)