Update app.py

app.py

@@ -12,16 +12,40 @@ from database import (
     get_table_schema
 )
 
-def process_uploaded_file(file):
+def process_sql_file(file_path):
     """
-    Process the uploaded CSV file and load it into the database.
+    Process an SQL file and execute its contents.
     """
     try:
-        if file is None:
-            return False, "Please upload a file."
+        # Read the SQL file
+        with open(file_path, 'r') as file:
+            sql_content = file.read()
             
+        # Split into individual statements
+        # This assumes statements are separated by semicolons
+        statements = sql_content.split(';')
+        
+        # Clear existing database
+        clear_database()
+        
+        # Execute each statement
+        with engine.begin() as conn:
+            for statement in statements:
+                if statement.strip():  # Skip empty statements
+                    conn.execute(text(statement))
+                    
+        return True, "SQL file successfully executed! Proceeding to query interface..."
+        
+    except Exception as e:
+        return False, f"Error processing SQL file: {str(e)}"
+
+def process_csv_file(file_path):
+    """
+    Process a CSV file and load it into the database.
+    """
+    try:
         # Read the CSV file
-        df = pd.read_csv(file)
+        df = pd.read_csv(file_path)
         
         if len(df.columns) == 0:
             return False, "Error: File contains no columns"
@@ -34,8 +58,29 @@ def process_uploaded_file(file):
         records = df.to_dict('records')
         insert_rows_into_table(records, table)
         
-        return True, "File successfully loaded! Proceeding to query interface..."
+        return True, "CSV file successfully loaded! Proceeding to query interface..."
         
+    except Exception as e:
+        return False, f"Error processing CSV file: {str(e)}"
+
+def process_uploaded_file(file):
+    """
+    Process the uploaded file (either SQL or CSV).
+    """
+    try:
+        if file is None:
+            return False, "Please upload a file."
+            
+        # Get file extension
+        file_ext = os.path.splitext(file)[1].lower()
+        
+        if file_ext == '.sql':
+            return process_sql_file(file)
+        elif file_ext == '.csv':
+            return process_csv_file(file)
+        else:
+            return False, "Error: Unsupported file type. Please upload either a .sql or .csv file."
+            
     except Exception as e:
         return False, f"Error processing file: {str(e)}"
 
@@ -44,8 +89,20 @@ def get_data_table():
     Fetches all data from the current table and returns it as a Pandas DataFrame.
     """
     try:
+        # Get list of tables
         with engine.connect() as con:
-            result = con.execute(text("SELECT * FROM data_table"))
+            tables = con.execute(text(
+                "SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%'"
+            )).fetchall()
+            
+        if not tables:
+            return pd.DataFrame()
+            
+        # Use the first table found
+        table_name = tables[0][0]
+            
+        with engine.connect() as con:
+            result = con.execute(text(f"SELECT * FROM {table_name}"))
             rows = result.fetchall()
             
             if not rows:
@@ -98,7 +155,7 @@ def query_sql(user_query: str) -> str:
         return "Error: No data table exists. Please upload a file first."
         
     schema_info = (
-        "The database has a table named 'data_table' with the following schema:\n"
+        f"The database has the following schema:\n"
         f"{schema}\n"
         "Generate a valid SQL SELECT query using ONLY these column names.\n"
         "DO NOT explain your reasoning, and DO NOT return anything other than the SQL query itself."
@@ -184,17 +241,26 @@ with gr.Blocks() as demo:
         gr.Markdown("""
         # Data Query Interface
 
-        Upload your CSV file to begin.
+        Upload your data file to begin.
+        
+        ### Supported File Types:
+        - SQL (.sql): SQL file containing CREATE TABLE and INSERT statements
+        - CSV (.csv): CSV file with headers that will be automatically converted to a table
         
-        ### Requirements:
-        - File must be in CSV format
+        ### CSV Requirements:
+        - Must include headers
         - First column will be used as the primary key
-        - All columns will be automatically typed based on their content
+        - Column types will be automatically detected
+        
+        ### SQL Requirements:
+        - Must contain valid SQL statements
+        - Statements must be separated by semicolons
+        - Should include CREATE TABLE and data insertion statements
         """)
         
         file_input = gr.File(
-            label="Upload CSV File",
-            file_types=[".csv"],
+            label="Upload Data File",
+            file_types=[".csv", ".sql"],
             type="filepath"
         )
         status = gr.Textbox(label="Status", interactive=False)