Added Database Schemas

app.py

@@ -40,19 +40,22 @@ print('Model Loaded...')
 print(f'Model Device: {model.device}')
 
 # Get Databases
-def get_databases():
-    databases = conn.execute("PRAGMA show_databases").fetchall()
-    return [item[0] for item in databases] 
+def get_schemas():
+    schemas = conn.execute("""
+    SELECT DISTINCT schema_name 
+    FROM information_schema.schemata 
+    WHERE schema_name NOT IN ('information_schema', 'pg_catalog')
+    """).fetchall()
+    return [item[0] for item in schemas] 
 
 # Get Tables
-def get_tables(database):
-    conn.execute(f"USE {database}")
-    tables = conn.execute("SHOW TABLES").fetchall()
+def get_tables(schema_name):
+    tables = conn.execute(f"SELECT table_name FROM information_schema.tables WHERE table_schema = '{schema_name}'").fetchall()
     return [table[0] for table in tables]
 
 # Update Tables
-def update_tables(selected_db):
-    tables = get_tables(selected_db)
+def update_tables(schema_name):
+    tables = get_tables(schema_name)
     return gr.update(choices=tables)
 
 # Get Schema
@@ -130,7 +133,7 @@ def text2sql(table, query_input):
     }
 
 # Load Databases Names
-databases = get_databases()
+
 
 # Custom CSS styling
 custom_css = """
@@ -164,7 +167,7 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="purple", secondary_hue="indigo"
     with gr.Row():
 
         with gr.Column(scale=1, variant='panel'):
-            database_dropdown = gr.Dropdown(choices=databases, label="Select Database", interactive=True)
+            schema_dropdown = gr.Dropdown(choices=get_schemas(), label="Select Schema", interactive=True)
             tables_dropdown = gr.Dropdown(choices=[], label="Available Tables", value=None)
 
         with gr.Column(scale=2):
@@ -185,7 +188,7 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="purple", secondary_hue="indigo"
                 with gr.Tab("Schema"):
                     table_schema = gr.Textbox(lines=TAB_LINES, label="Table Schema", value="", interactive=False)
  
-        database_dropdown.change(update_tables, inputs=database_dropdown, outputs=tables_dropdown)
+        schema_dropdown.change(update_tables, inputs=schema_dropdown, outputs=tables_dropdown)
         generate_query_button.click(text2sql, inputs=[tables_dropdown, query_input], outputs=[table_schema, input_prompt, generated_query, result_output])
 
 if __name__ == "__main__":