CSVAgent / database.py
Quazim0t0's picture
add default uuid feature (#1)
ef08977 verified
from sqlalchemy import (
create_engine,
MetaData,
Table,
Column,
String,
Integer,
Float,
insert,
text,
)
# Use a persistent SQLite database file
engine = create_engine("sqlite:///database.db")
metadata_obj = MetaData()
def create_dynamic_table(df):
"""
Creates a table dynamically based on DataFrame schema.
Args:
df: pandas DataFrame containing the data
Returns:
SQLAlchemy Table object
"""
# Drop existing table if it exists
if 'data_table' in metadata_obj.tables:
metadata_obj.remove(metadata_obj.tables['data_table'])
# Define columns with a new auto-incrementing primary key
columns = [Column("uuid", Integer, primary_key=True, autoincrement=True)] # New primary key
# Add columns based on DataFrame dtypes
for col_name, dtype in df.dtypes.items():
if 'int' in str(dtype):
col_type = Integer
elif 'float' in str(dtype):
col_type = Float
else:
col_type = String(255) # Using a generous length for string columns
columns.append(Column(col_name, col_type))
# Create new table
table = Table('data_table', metadata_obj, *columns)
# Create table in database
metadata_obj.create_all(engine, tables=[table])
return table
def clear_database():
"""
Removes all tables from the database.
"""
metadata_obj.drop_all(engine)
metadata_obj.clear()
def insert_rows_into_table(rows, table):
"""
Inserts rows into the specified table.
Args:
rows: List of dictionaries containing the row data
table: SQLAlchemy Table object
"""
with engine.begin() as connection:
connection.execute(insert(table), rows)
def get_table_schema():
"""
Returns the current table schema as a string.
"""
try:
# Get list of tables
with engine.connect() as con:
tables = con.execute(text(
"SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%'"
)).fetchall()
if not tables:
return None
# Use the first table found
table_name = tables[0][0]
# Get column information
with engine.connect() as con:
columns = con.execute(text(f"PRAGMA table_info({table_name})")).fetchall()
schema = []
for col in columns:
col_name = col[1]
col_type = col[2]
is_primary = "primary key" if col[5] == 1 else ""
schema.append(f"- {col_name} ({col_type}) {is_primary}".strip())
return "\n".join(schema)
except Exception as e:
return None