Spaces:
Running
Running
| import sqlite3 | |
| import sqliteschema | |
| from typing import Dict, Any | |
| # ------------------------- | |
| # 1) Business SQLite Schema | |
| # ------------------------- | |
| def get_sqlite_schema() -> Dict[str, Any]: | |
| """ | |
| Extracts the schema of the hardcoded business SQLite database | |
| and returns it as a dictionary. | |
| """ | |
| db_path = "business_database.db" | |
| extractor = sqliteschema.SQLiteSchemaExtractor(db_path) | |
| schema_dict = extractor.fetch_database_schema_as_dict() | |
| return schema_dict | |
| # ------------------------- | |
| # 2) Uploaded SQLite Schema | |
| # ------------------------- | |
| def get_uploaded_sqlite_schema(db_path: str) -> Dict[str, Any]: | |
| """ | |
| Extracts the schema from a user uploaded SQLite file. | |
| """ | |
| extractor = sqliteschema.SQLiteSchemaExtractor(db_path) | |
| schema_dict = extractor.fetch_database_schema_as_dict() | |
| return schema_dict | |
| # ------------------------- | |
| # 3) MySQL Schema | |
| # ------------------------- | |
| def get_mysql_schema(host: str, port: str, user: str, password: str, database: str) -> Dict[str, Any]: | |
| """ | |
| Connects to a MySQL database and extracts its schema. | |
| """ | |
| import mysql.connector | |
| conn = mysql.connector.connect( | |
| host=host, | |
| port=int(port), | |
| user=user, | |
| password=password, | |
| database=database | |
| ) | |
| cursor = conn.cursor() | |
| cursor.execute("SHOW TABLES") | |
| tables = [t[0] for t in cursor.fetchall()] | |
| schema: Dict[str, Any] = {} | |
| for table in tables: | |
| cursor.execute(f"DESCRIBE {table}") | |
| cols = cursor.fetchall() | |
| schema[table] = [ | |
| {"Field": col[0], "Type": col[1], "Null": col[2], "Key": col[3], "Default": col[4], "Extra": col[5]} | |
| for col in cols | |
| ] | |
| conn.close() | |
| return schema | |
| # ------------------------- | |
| # 4) PostgreSQL Schema | |
| # ------------------------- | |
| def get_postgresql_schema(host: str, port: str, user: str, password: str, database: str) -> Dict[str, Any]: | |
| """ | |
| Connects to a PostgreSQL database and extracts its schema. | |
| """ | |
| import psycopg2 | |
| conn = psycopg2.connect( | |
| host=host, | |
| port=int(port), | |
| user=user, | |
| password=password, | |
| dbname=database | |
| ) | |
| cursor = conn.cursor() | |
| # list all public tables | |
| cursor.execute(""" | |
| SELECT table_name | |
| FROM information_schema.tables | |
| WHERE table_schema = 'public' | |
| AND table_type = 'BASE TABLE'; | |
| """) | |
| tables = [t[0] for t in cursor.fetchall()] | |
| schema: Dict[str, Any] = {} | |
| for table in tables: | |
| cursor.execute(f""" | |
| SELECT column_name, data_type | |
| FROM information_schema.columns | |
| WHERE table_name = '{table}'; | |
| """) | |
| cols = cursor.fetchall() | |
| schema[table] = [{"column_name": col[0], "data_type": col[1]} for col in cols] | |
| conn.close() | |
| return schema | |