import sqlite3 import os import glob import json import logging import csv from hardlevel import eval_hardness from sqlalchemy import text from utils.database import create_engine_with_timeout def execute_sql(database_uri, sql: str): row_count, col_count, msg, results = 0, 0, "", [] ddl_keywords = ['CREATE', 'ALTER', 'DROP', 'TRUNCATE', 'UPDATE'] for keyword in ddl_keywords: if sql.strip().upper().startswith(keyword): return None, None, "ddl is forbidden", results try: engine = create_engine_with_timeout(database_uri) with engine.connect() as conn: for i in conn.execute(text(sql)): row_count += 1 col_count = len(i) results.append(i) except Exception as e: return None, None, str(e), results return row_count, col_count, msg, results def compare_sql(database_uri, sql, sql_generated): logging.info("compare_sql %s, %s, database_uri: %s", sql, sql_generated, database_uri) row_count_sql, col_count_sql, msg_sql, results_sql = execute_sql(database_uri, sql) row_count_gen, col_count_gen, msg_gen, results_gen = execute_sql(database_uri, sql_generated) logging.debug( "sql, row_count_sql, col_count_sql, msg_sql, results_sql: %s, %s, %s, %s, %s", sql, row_count_sql, col_count_sql, msg_sql, results_sql, ) logging.debug( "sql_generated, row_count_gen, col_count_gen, msg_gen, results_gen: %s, %s, %s, %s, %s", sql_generated, row_count_gen, col_count_gen, msg_gen, results_gen, ) if row_count_sql is None or col_count_sql is None: return "raw query failed" if row_count_gen is None or col_count_gen is None: return "run sql failed" if row_count_sql > row_count_gen or col_count_sql > col_count_gen or not sql_results_issubset(results_sql, results_gen): # 生成结果 < 标准答案，任何一条记录，在标准答案里，但是不在生成答案里，判定为错误 # 判断标准答案是否是生成答案的子集，如果不是，则符合条件 logging.info("results_gen < results_sql") return False elif row_count_sql == row_count_gen and col_count_sql <= col_count_gen: # 生成答案和标准答案，行数相等，生成答案列包含标准答案，判定为✅ # 判断标准答案是否是生成答案的子集，如果不是，则符合条件 result = sql_results_issubset(results_sql, results_gen) logging.debug("results_sql is subset of results_gen: %s", result) return result else: # human check return "human check" def sql_not_valid(sql): if sql in ("job failed"): return False return True def compare_row(dbname, sql, generated_sql, sqlite=False, sqlite_local=False): if not (sql_not_valid(sql) and sql_not_valid(generated_sql)): logging.debug("sql %s, %s not valid, ignore", sql, generated_sql) return None if generated_sql == "sql not found": return "sql not found" database_uri = get_database_uri(dbname, sqlite, sqlite_local) same = compare_sql(database_uri, sql, generated_sql) logging.debug("sql, generated_sql, same: %s, %s, %s", sql, generated_sql, same) return same def get_database_uri(dbname, sqlite=False, spider_local=False): if os.getenv("BIRD_DATASET"): return "sqlite:///spider/bird/dev_databases/{}/{}.sqlite".format(dbname, dbname) if sqlite: dbname = dbname.removeprefix("new_") if spider_local: return "sqlite:///database/{}/{}.sqlite?mode=ro".format(dbname, dbname) if os.getenv("RUN_IN_CLOUD"): return "sqlite:///spider/database/{}/{}.sqlite?mode=ro".format(dbname, dbname) else: return "spider://{}".format(dbname) else: sql_dialect = "mysql+pymysql://2rrudfkjr8MQq3Q.root:9pS7AkerMB1Py3GE@gateway01.us-west-2.prod.aws.tidbcloud.com:4000/{}?ssl_ca=%2Fetc%2Fssl%2Fcerts%2Fca-certificates.crt&ssl_verify_identity=true".format(dbname) # noqa return sql_dialect def compare_csv(dbname, filename): with open(filename) as f: reader = csv.DictReader(f) for i in reader: sql = i["sql"] generated_sql = i["generated_sql"] print(compare_row(dbname, sql, generated_sql)) def sql_results_issubset(results_a, results_b): # 比较两个SQL执行结果，如果返回True，代表results_a是results_b的子集 # SQL 的执行结果，可能行和列的顺序都不同 results_list_b = [] all_match = True for rowb in results_b: results_list_b.append(set(rowb)) for rowa in results_a: row_match = False row_set_a = set(rowa) for row_set_b in results_list_b: if row_set_a.issubset(row_set_b): row_match = True break all_match = all_match and row_match if not all_match: break return all_match def compare_generated_sql(dbname, csv_filename): results = [] with open(csv_filename) as f: for i in csv.DictReader(f): results.append(i) with open(csv_filename, "w") as f: writer = csv.writer(f) csv_header = ["database", "question", "sql", "generated_sql", "difficulty", "equal", "note"] writer.writerow(csv_header) for i in results: sql = i["sql"] generated_sql = i["generated_sql"] hardness = i["difficulty"] database = i.get("\ufeffdatabase") or i.get("database") equal = compare_sql(get_database_uri(dbname, sqlite=True), sql, generated_sql) row = [database, i["question"], sql, generated_sql, hardness, equal] if os.getenv("DEBUG"): row += [i["description"], i["clarified_task"]] writer.writerow(row) f.flush() def compare_din(dbname, csv_filename, json_filename, output_file): json_result_mapper = {} with open(json_filename) as f: for i in json.loads(f.read()): json_result_mapper[i["question"]] = i with open(output_file, "w") as fo: writer = csv.writer(fo) writer.writerow(["database", "question", "sql", "generated_sql", "difficulty", "equal"]) with open(csv_filename) as f: for i in csv.DictReader(f): question = i["question"] sql = i["sql"] din_sql = json_result_mapper[question]["sql"] hardness = i["difficulty"] equal = compare_sql(get_database_uri(dbname, sqlite=True), sql, din_sql) writer.writerow([i["database"], i["question"], sql, din_sql, hardness, equal]) fo.flush() def filter_all_cases_has_data(): for i in glob.glob("./cases_new*_st_sqlite.csv"): yield i def run_all_report_by_csv(): for i in filter_all_cases_has_data(): dbname = i.removeprefix("./cases_new_").removesuffix("_st_sqlite.csv") compare_generated_sql(dbname, i) if __name__ == "__main__": logging.basicConfig(level=logging.DEBUG) # run_all_report_by_csv() # compare_generated_sql("music_1", "./cases_new_music_1_st_sqlite.csv") sql = 'SELECT count(DISTINCT allergy) FROM Allergy_type;' gen_sql = 'SELECT COUNT(DISTINCT `allergy_type`.`allergy`) FROM `allergy_type`;' dburl = get_database_uri("allergy_1", sqlite=True, spider_local=True) print(compare_sql(dburl, sql, gen_sql))