csv2sql / sql.py
mobinln's picture
add csv ingestion and prompt templates
5875608
import pandas as pd
import sqlite3
def csv_to_sqlite(csv_file, db_name, table_name):
# Read the CSV file into a pandas DataFrame
df = pd.read_csv(csv_file)
# Connect to the SQLite database (it will create the database file if it doesn't exist)
conn = sqlite3.connect(db_name)
cursor = conn.cursor()
# Infer the schema based on the DataFrame columns and data types
def create_table_from_df(df, table_name):
# Get column names and types
col_types = []
for col in df.columns:
dtype = df[col].dtype
if dtype == "int64":
col_type = "INTEGER"
elif dtype == "float64":
col_type = "REAL"
else:
col_type = "TEXT"
col_types.append(f'"{col}" {col_type}')
# Create the table schema
col_definitions = ", ".join(col_types)
create_table_query = (
f"CREATE TABLE IF NOT EXISTS {table_name} ({col_definitions});"
)
# print(create_table_query)
# Execute the table creation query
cursor.execute(create_table_query)
print(f"Table '{table_name}' created with schema: {col_definitions}")
# Create table schema
create_table_from_df(df, table_name)
# Insert CSV data into the SQLite table
df.to_sql(table_name, conn, if_exists="replace", index=False)
# Commit and close the connection
conn.commit()
conn.close()
print(f"Data loaded into '{table_name}' table in '{db_name}' SQLite database.")
def run_sql_query(db_name, query):
"""
Executes a SQL query on a SQLite database and returns the results.
Args:
db_name (str): The name of the SQLite database file.
query (str): The SQL query to run.
Returns:
list: Query result as a list of tuples, or an empty list if no results or error occurred.
"""
try:
# Connect to the SQLite database
conn = sqlite3.connect(db_name)
cursor = conn.cursor()
# Execute the SQL query
cursor.execute(query)
# Fetch all results
results = cursor.fetchall()
# Close the connection
conn.close()
# Return results or an empty list if no results were found
return results if results else []
except sqlite3.Error as e:
print(f"An error occurred while executing the query: {e}")
return []
def get_table_schema(db_name, table_name):
"""
Retrieves the schema (columns and data types) for a given table in the SQLite database.
Args:
db_name (str): The name of the SQLite database file.
table_name (str): The name of the table.
Returns:
list: A list of tuples with column name, data type, and other info.
"""
conn = sqlite3.connect(db_name)
cursor = conn.cursor()
# Use PRAGMA to get the table schema
cursor.execute(f"PRAGMA table_info({table_name});")
schema = cursor.fetchall()
conn.close()
return schema