nl2sql_library/utils/utility_functions.py [18:145]: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - load_dotenv() # bigquery_connection_string = # "bigquery://sl-test-project-363109/nl2sql_spider" PROJ_CONFIG_FILE = "utils/proj_config.json" SQL_LOG_FILE = "utils/sqlgen_log.json" proj_config_dict = { "default": { "proj_name": "sl-test-project-363109", "dataset": "nl2sql_spider", "metadata_file": "spider_md_cache.json", }, "config": {"proj_name": "", "dataset": "", "metadata_file": ""}, } def initialize_db(proj="sl-test-project-363109", dataset="nl2sql_spider"): """ Initialize the BQ connection string """ print("Init init DB ", proj, dataset) bigquery_connection_string = f"bigquery://{proj}/{dataset}" print(bigquery_connection_string) logger.info(f"BQ connection string: {bigquery_connection_string}") return bigquery_connection_string def config_project( proj_name="sl-test-project-363109", dataset="nl2sql_spider", metadata_file="spider_md_cache.json", ): """ Save the project configuration details """ try: with open(PROJ_CONFIG_FILE, "r", encoding="utf-8") as infile: proj_config = json.load(infile) # except: except FileNotFoundError: logger.error("File not found, using default configuration") proj_config = proj_config_dict logger.info( f"Saving following data : {proj_name}, {dataset}, {metadata_file}" ) proj_config["config"]["proj_name"] = proj_name proj_config["config"]["dataset"] = dataset proj_config["config"]["metadata_file"] = metadata_file json_obj = json.dumps(proj_config, indent=4) with open(PROJ_CONFIG_FILE, "w", encoding="utf-8") as outfile: outfile.write(json_obj) # json.dump(proj_config, outfile) with open(PROJ_CONFIG_FILE, "r", encoding="utf-8") as infile: proj_config = json.load(infile) logger.info( f"New file name : {proj_config['config']['metadata_file']}" ) initialize_db(proj=proj_name, dataset=dataset) def get_project_config(): """ Return the Project configuration details """ try: with open(PROJ_CONFIG_FILE, "r", encoding="utf-8") as infile: proj_config = json.load(infile) # except: except FileNotFoundError: logger.error("File not found, using default configuration") proj_config = proj_config_dict proj_config["config"] = proj_config["default"] json_obj = json.dumps(proj_config, indent=4) with open(PROJ_CONFIG_FILE, "w", encoding="utf-8") as outfile: outfile.write(json_obj) # json.dump(proj_config, outfile) return proj_config def execute_bq_query(sql_query=""): """ Execute the given query on BigQuery """ project = get_project_config()["config"]["proj_name"] client = bigquery.Client(project=project) logger.info(f"Execute bq query : {sql_query}") # query_job = client.query( # "select * from `q_and_a_db.questions_and_gensqls` \ # where created_by = 'CoT executor' " ) try: query_job = client.query(sql_query) results = query_job.result() results = query_job.to_dataframe() logger.info("Execution result = ", results) return results except Exception: return "Execution failed" # Below function for future upgrades enabling logging within BQ # def bq_insert_data(result_id="dummy", # question="dummy", # sql="dummy", executor="dummy", # feedback="False"): # """ # Adds a new row to BigQuery table # """ # project = get_project_config()['config']['proj_name'] # dataset = "q_and_a_db" # client = bigquery.Client(project=project) # table_ref = f"{project}.{dataset}.questions_and_gensqls" # insert_query = f'INSERT INTO {table_ref} VALUES\ # ("{result_id}", "{question}", "{sql}", "{executor}", {feedback})' # logger.info("Query to insert=", insert_query) # query_job = client.query(insert_query) # query_job.result() def log_sql(result_id="dummy", question="dummy", sql="dummy", executor="dummy", feedback="False" ): - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - nl2sql_src/utils/utility_functions.py [18:147]: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - load_dotenv() # bigquery_connection_string = # "bigquery://sl-test-project-363109/nl2sql_spider" PROJ_CONFIG_FILE = "utils/proj_config.json" SQL_LOG_FILE = "utils/sqlgen_log.json" proj_config_dict = { "default": { "proj_name": "sl-test-project-363109", "dataset": "nl2sql_spider", "metadata_file": "spider_md_cache.json", }, "config": {"proj_name": "", "dataset": "", "metadata_file": ""}, } def initialize_db(proj="sl-test-project-363109", dataset="nl2sql_spider"): """ Initialize the BQ connection string """ print("Init init DB ", proj, dataset) bigquery_connection_string = f"bigquery://{proj}/{dataset}" print(bigquery_connection_string) logger.info(f"BQ connection string: {bigquery_connection_string}") return bigquery_connection_string def config_project( proj_name="sl-test-project-363109", dataset="nl2sql_spider", metadata_file="spider_md_cache.json", ): """ Save the project configuration details """ try: with open(PROJ_CONFIG_FILE, "r", encoding="utf-8") as infile: proj_config = json.load(infile) # except: except FileNotFoundError: logger.error("File not found, using default configuration") proj_config = proj_config_dict logger.info( f"Saving following data : {proj_name}, {dataset}, {metadata_file}" ) proj_config["config"]["proj_name"] = proj_name proj_config["config"]["dataset"] = dataset proj_config["config"]["metadata_file"] = metadata_file json_obj = json.dumps(proj_config, indent=4) with open(PROJ_CONFIG_FILE, "w", encoding="utf-8") as outfile: outfile.write(json_obj) # json.dump(proj_config, outfile) with open(PROJ_CONFIG_FILE, "r", encoding="utf-8") as infile: proj_config = json.load(infile) logger.info( f"New file name : {proj_config['config']['metadata_file']}" ) initialize_db(proj=proj_name, dataset=dataset) def get_project_config(): """ Return the Project configuration details """ try: with open(PROJ_CONFIG_FILE, "r", encoding="utf-8") as infile: proj_config = json.load(infile) # except: except FileNotFoundError: logger.error("File not found, using default configuration") proj_config = proj_config_dict proj_config["config"] = proj_config["default"] json_obj = json.dumps(proj_config, indent=4) with open(PROJ_CONFIG_FILE, "w", encoding="utf-8") as outfile: outfile.write(json_obj) # json.dump(proj_config, outfile) return proj_config def execute_bq_query(sql_query=""): """ Execute the given query on BigQuery """ project = get_project_config()["config"]["proj_name"] client = bigquery.Client(project=project) logger.info(f"Execute bq query : {sql_query}") # query_job = client.query( # "select * from `q_and_a_db.questions_and_gensqls` \ # where created_by = 'CoT executor' " ) try: query_job = client.query(sql_query) results = query_job.result() results = query_job.to_dataframe() logger.info("Execution result = ", results) return results except Exception: return "Execution failed" # Below function for future upgrades enabling logging within BQ # def bq_insert_data(result_id="dummy", # question="dummy", # sql="dummy", executor="dummy", # feedback="False"): # """ # Adds a new row to BigQuery table # """ # project = get_project_config()['config']['proj_name'] # dataset = "q_and_a_db" # client = bigquery.Client(project=project) # table_ref = f"{project}.{dataset}.questions_and_gensqls" # insert_query = f'INSERT INTO {table_ref} VALUES\ # ("{result_id}", "{question}", "{sql}", "{executor}", {feedback})' # logger.info("Query to insert=", insert_query) # query_job = client.query(insert_query) # query_job.result() def log_sql(result_id="dummy", question="dummy", sql="dummy", executor="dummy", feedback="False" ): - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -