import os import sys import json import time import shutil import requests REPORT_DIR = sys.path[0] + "/report" CONFIG_FILE = sys.path[0] + "/api_config.json" DATA_FILE = sys.path[0] + "/../data/table_info_with_answer.json" def get_input_file(report_dir): answer_file = report_dir + "/input.json" shutil.copyfile(DATA_FILE, answer_file) return answer_file def create_dir_if_not_exists(dir_path): if not os.path.exists(dir_path): os.makedirs(dir_path) pass return dir_path def get_cur_datetime(date_format="%Y%m%d%H%m%S"): timestamp = time.time() time_array = time.localtime(timestamp) return time.strftime(date_format, time_array) def read_json_file(file_path): f = open(file_path, encoding = "utf-8") content = f.read() f.close() return json.loads(content) def write_json_file(file_path, obj, indent = False): f = open(file_path, "w") if indent: content = f.write(json.dumps(obj, indent=2)) pass else: content = f.write(json.dumps(obj)) f.close() return file_path def run_cmd(cmd): return os.system(cmd) def get_sqlite3_table(file_path, table): mydb = sqlite3.connect(file_path) cur = mydb.cursor() cur.execute("select name from sqlite_master where type='table'") return def gen_import_cmd(file_path, database): ret_list = [] mydb = sqlite3.connect(file_path) cur = mydb.cursor() cur.execute("select name from sqlite_master where type='table'") tables = cur.fetchall() table_list = [item[0] for item in tables] for table in table_list: sql_cmd = "select sql from sqlite_master where name='" + table + "'" cur.execute(sql_cmd) for row in cur: sql = row[0].replace("NUMERIC", "INT(11)").replace("INTEGER", "INT(11)") item = {"database": database, "sql": sql, "table": table} ret_list.append(item) pass pass return ret_list def http_post_with_auth(url, token, data=None): headers = {"authorization": token, "content-type": "application/json"} resp = requests.post(url, headers=headers, data=json.dumps(data)) return json.loads(resp.text) # https://us-west-2.prod.aws.tidbcloud.com/api/v1/dataapps/sqleditor/orgs/1/projects/1372813089195381283/clusters/1379661944639654096/session def gen_sql_session(config_obj): url = config_obj["url_prefix"] + "/orgs/" + config_obj["org_id"] + "/projects/" + config_obj["project_id"] + "/clusters/" + config_obj["cluster_id"] + "/session" resp_obj = http_post_with_auth(url, config_obj["token"], {"database":""}) if "data" not in resp_obj: return "" return resp_obj["data"] def query_data_api(config_obj, sql): url = config_obj["url_prefix"] + "/orgs/" + config_obj["org_id"] + "/projects/" + config_obj["project_id"] + "/clusters/" + config_obj["cluster_id"] + "/statements" resp_obj = http_post_with_auth(url, config_obj["token"], {"sql":sql, "sessionid": config_obj["session_id"]}) if "data" not in resp_obj: return resp_obj return resp_obj["data"] def evaluate(config_obj, data_obj): report_obj = {"summary":{}, "data":[]} nogen_list = [] fail_list = [] total_num = 0 succ_num = 0 no_gen_num = 0 for table_item in data_obj: # use database query_data_api(config_obj, "use " + table_item["database"]) for index in range(len(table_item["answers"])): answer = table_item["answers"][index] total_num +=1 if answer.replace("\n", "").rstrip() == "No SQL Generated": no_gen_num +=1 nogen_list.append({"database": table_item["database"], "question": table_item["questions"][index], "answer": answer}) continue resp = query_data_api(config_obj, answer) report_item = {"database": table_item["database"], "sql": answer, "resp": resp} if "query" in resp and resp['query'] == 'Query OK!': succ_num +=1 else: print(report_item) fail_list.append(report_item) report_obj["data"].append(report_item) pass pass report_obj["summary"] = {"total": total_num, "succ": succ_num, "no_gen": no_gen_num} return report_obj, nogen_list, fail_list report_dir = create_dir_if_not_exists(REPORT_DIR + "/" + get_cur_datetime("%Y%m%d")) data_file = get_input_file(report_dir) if data_file == "": print("no datafile generated.") exit(0) data_obj = read_json_file(data_file) config_obj = read_json_file(CONFIG_FILE) if config_obj["session_id"] == "": session_id = gen_sql_session(config_obj) config_obj["session_id"] = session_id write_json_file(CONFIG_FILE, config_obj) report_obj, nogen_list, fail_list = evaluate(config_obj, data_obj) write_json_file(report_dir + "/output.json", report_obj) write_json_file(report_dir + "/nogen_list.json", nogen_list, True) write_json_file(report_dir + "/fail_list.json", fail_list, True)