Update app.py

app.py

@@ -12,6 +12,37 @@ from database import (
     get_table_schema
 )
 
+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:
+            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:
+                return pd.DataFrame()
+                
+            columns = result.keys()
+            df = pd.DataFrame(rows, columns=columns)
+            return df
+
+    except Exception as e:
+        return pd.DataFrame({"Error": [str(e)]})
+
 def get_table_info():
     """
     Gets the current table name and column information.
@@ -206,12 +237,6 @@ def query_sql(user_query: str) -> str:
     except ValueError:
         return result
 
-def update_schema():
-    schema = get_table_schema()
-    if schema:
-        return f"### Current Schema:\n```\n{schema}\n```"
-    return "No data loaded"
-
 # Create the Gradio interface
 with gr.Blocks() as demo:
     with gr.Group() as upload_group:
@@ -261,24 +286,32 @@ with gr.Blocks() as demo:
 
     def handle_upload(file_obj):
         if file_obj is None:
-            return "Please upload a file.", None, None, "No schema available", gr.update(visible=True), gr.update(visible=False)
+            return (
+                "Please upload a file.",
+                None,
+                "No schema available",
+                gr.update(visible=True),
+                gr.update(visible=False)
+            )
         
         success, message = process_uploaded_file(file_obj)
         if success:
             df = get_data_table()
-            schema = get_table_schema()
+            _, _, column_info = get_table_info()
+            schema = "\n".join([
+                f"- {col} ({info['type']}){' primary key' if info['is_primary'] else ''}"
+                for col, info in column_info.items()
+            ])
             return (
                 message,
                 df,
-                schema,
-                f"### Current Schema:\n```\n{schema}\n```" if schema else "No schema available",
+                f"### Current Schema:\n```\n{schema}\n```",
                 gr.update(visible=False),
                 gr.update(visible=True)
             )
         return (
             message,
             None,
-            None,
             "No schema available",
             gr.update(visible=True),
             gr.update(visible=False)
@@ -286,8 +319,12 @@ with gr.Blocks() as demo:
 
     def refresh_data():
         df = get_data_table()
-        schema = get_table_schema()
-        return df, f"### Current Schema:\n```\n{schema}\n```" if schema else "No schema available"
+        _, _, column_info = get_table_info()
+        schema = "\n".join([
+            f"- {col} ({info['type']}){' primary key' if info['is_primary'] else ''}"
+            for col, info in column_info.items()
+        ])
+        return df, f"### Current Schema:\n```\n{schema}\n```"
 
     # Event handlers
     file_input.upload(
@@ -297,7 +334,6 @@ with gr.Blocks() as demo:
             status,
             data_table,
             schema_display,
-            schema_display,
             upload_group,
             query_group
         ]