sql_command.py

'''

'''
##TODO:

import sqlite3
import pandas as pd

## run the following function to set up a dedicated database using SQLite.
def construct_db(data_file=None):
    excel_file = "/Users/yunshi/Downloads/360Data/Data Center/Working-On Task/演讲与培训/2023ChatGPT/Coding/Text2SQL/模拟数据.csv"  # Replace with your actual file path
    df = pd.read_csv(excel_file)
    print('df:', df.head())
    
    conn = sqlite3.connect('myexcelDB.db')  # Replace 'mydatabase.db' with your desired name
    # Create a cursor object to execute SQL commands
    cursor = conn.cursor()

    ##NOTE: Insert data from DataFrame into the table. 注意这里的if_exists选项，考虑是否要覆盖原始内容。这里如果需要指定index，那么就需要指定index_label，且index=False，否则会选择默认的index，这样就会产生duplicate错误。
    df.to_sql('table01', conn, if_exists='replace', index=False, index_label="产品ID")
    
    return None


def llm_query(sql_command):
    # Connect to the database (or create it if it doesn't exist)
    conn = sqlite3.connect('./myexcelDB.db')  # Replace 'myexcelDB.db' with your desired name

    # Create a cursor object to execute SQL commands
    cursor = conn.cursor()
    ## SQL command
    cursor.execute(sql_command)
    query_result = cursor.fetchall()  # Fetch all rows as a list of tuples
    # print('query_result:', query_result)
    
    ## 将列名也取出来。
    column_names = [description[0] for description in cursor.description]
    
    query_df = pd.DataFrame(query_result, columns=column_names)
    # query_df.set_index("产品ID", inplace=True)
    print('query_df:', query_df)
    
    return query_df

# llm_query("SELECT * FROM table01 WHERE 宽度 > 300") ## sample function call. 

# construct_db()
# sql = "SELECT 产品ID FROM table01 WHERE 长度 > 50"
# res = llm_query(sql)
# res