averrie/TEXT2SQL

1BQ_GET_TABLES_TEMPLATE = """2import os3import pandas as pd4from google.cloud import bigquery5 6os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = "/workspace/bigquery_credential.json"7client = bigquery.Client()8 9query = f\"\"\"10SELECT11    table_name, ddl12FROM13    `{database_name}.{dataset_name}.INFORMATION_SCHEMA.TABLES`14WHERE15    table_type != 'VIEW'16\"\"\"17 18query_job = client.query(query)19 20try:21    results = query_job.result().to_dataframe() 22    if results.empty:23        print("No data found for the specified query.")24    else:25        results.to_csv("{save_path}", index=False)26        print(f"DB metadata is saved to {save_path}")27except Exception as e:28    print("Error occurred while fetching data: ", e)29"""30 31BQ_GET_TABLE_INFO_TEMPLATE = """32import os33import pandas as pd34from google.cloud import bigquery35 36os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = "/workspace/bigquery_credential.json"37client = bigquery.Client()38 39query = f\"\"\"40    SELECT field_path, data_type, description41    FROM `{database_name}.{dataset_name}.INFORMATION_SCHEMA.COLUMN_FIELD_PATHS`42    WHERE table_name = '{table}';43\"\"\"44 45query_job = client.query(query)46 47try:48    output = query_job.result().to_dataframe() 49    if output.empty:50        print("No data found for the specified query.")51    else:52        output.to_csv("{save_path}", index=False)53        print(f"Results saved to {save_path}")54except Exception as e:55    print("Error occurred while fetching data: ", e)56"""57 58BQ_SAMPLE_ROWS_TEMPLATE = """59import os60import pandas as pd61from google.cloud import bigquery62import json63os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = "/workspace/bigquery_credential.json"64client = bigquery.Client()65 66query = f\"\"\"67    SELECT68    *69    FROM70    `{database_name}.{dataset_name}.{table}`71    TABLESAMPLE SYSTEM (0.0001 PERCENT)72    LIMIT {row_number};73\"\"\"74 75query_job = client.query(query)76 77try:78    output = query_job.result().to_dataframe() 79    if output.empty:80        print("No data found for the specified query.")81    else:82        sample_rows = output.to_dict(orient='records')83        json_data = json.dumps(sample_rows, indent=4, default=str)84        with open("{save_path}", 'w') as json_file: 85            json_file.write(json_data)86        print(f"Sample rows saved to {save_path}")87except Exception as e:88    print("Error occurred while fetching data: ", e)89"""90 91BQ_EXEC_SQL_QUERY_TEMPLATE = """92import os93import pandas as pd94from google.cloud import bigquery95 96os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = "/workspace/bigquery_credential.json"97client = bigquery.Client()98 99sql_query = f\"\"\"{sql_query}\"\"\"100 101query_job = client.query(sql_query)102 103try:104    results = query_job.result().to_dataframe()105    if results.empty:106        print("No data found for the specified query.")107    else:108        if {is_save}:109            results.to_csv("{save_path}", index=False)110            print(f"Results saved to {save_path}")111        else:112            print(results)113except Exception as e:114    print("Error occurred while fetching data: ", e)115"""116 117 118SF_EXEC_SQL_QUERY_TEMPLATE = """119import os120import json121import pandas as pd122import snowflake.connector123 124# Load Snowflake credentials125snowflake_credential = json.load(open("/workspace/snowflake_credential.json"))126 127# Connect to Snowflake128conn = snowflake.connector.connect(129    **snowflake_credential130)131cursor = conn.cursor()132 133# Define the SQL query134sql_query = f\"\"\"135 136 137{sql_query}138 139 140\"\"\"141 142# Execute the SQL query143cursor.execute(sql_query)144 145try:146    # Fetch the results147    results = cursor.fetchall()148    columns = [desc[0] for desc in cursor.description]149    df = pd.DataFrame(results, columns=columns)150 151    # Check if the result is empty152    if df.empty:153        print("No data found for the specified query.")154    else:155        # Save or print the results based on the is_save flag156        if {is_save}:157            df.to_csv("{save_path}", index=False)158            print(f"Results saved to {save_path}")159        else:160            print(df)161except Exception as e:162    print("Error occurred while fetching data: ", e)163finally:164    cursor.close()165    conn.close()166"""167 168 169 170LOCAL_SQL_TEMPLATE = """171import pandas as pd172import os173import sqlite3174import duckdb175 176def detect_db_type(file_path):177    if file_path.endswith('.db') or file_path.endswith('.sqlite') :178        return 'sqlite'179    elif file_path.endswith('.duckdb'):180        return 'duckdb'181    else:182        try:183            conn = duckdb.connect(database=file_path, read_only=True)184            conn.execute('SELECT 1')185            conn.close()186            return 'duckdb'187        except:188            return 'sqlite'189 190def execute_sql(file_path, command, output_path_path):191    db_type = detect_db_type(file_path)192    193    # Make sure the file path is correct194    if not os.path.exists(file_path) and db_type == 'sqlite':195        print(f"ERROR: Database file not found: {{file_path}}")196        return197 198    # Connect to the database199    if db_type == 'sqlite':200        conn = sqlite3.connect(file_path)201    elif db_type == 'duckdb':202        conn = duckdb.connect(database=file_path, read_only=True)203    else:204        print(f"ERROR: Unsupported database type {{db_type}}")205        return206    207    try:208        # Execute the SQL command and fetch the results209        df = pd.read_sql_query(command, conn)210        211        # Check if the output should be saved to a CSV file or printed directly212        if output_path_path.lower().endswith(".csv"):213            df.to_csv(output_path_path, index=False)214            print(f"Output saved to: {{output_path_path}}")215        else:216            print(df)217    except Exception as e:218        print(f"ERROR: {{e}}")219    finally:220        # Close the connection to the database221        conn.close()222 223# Example usage224file_path = "{file_path}"    # Path to your database file225command = f\"\"\"{sql_command}\"\"\"    # SQL command to be executed226output_path = "{output_path}"# Path to save the output as a CSV or "directly"227 228execute_sql(file_path, command, output_path)229"""230