maikheb/nl2sql

1from sqlalchemy import create_engine, inspect
2import sqlite3
3
4def extract_schema_sqlite(db_path):
5    conn = sqlite3.connect(db_path)
6    cursor = conn.cursor()
7    cursor.execute("SELECT name FROM sqlite_master WHERE type='table';")
8    tables = [row[0] for row in cursor.fetchall()]
9    schema = {}
10
11    for table in tables:
12        cursor.execute(f"PRAGMA table_info({table});")
13        schema[table] = [col[1] for col in cursor.fetchall()]
14    conn.close()
15    return schema
16
17def extract_schema_rdbms(db_uri):
18    engine = create_engine(db_uri)
19    inspector = inspect(engine)
20    schema = {}
21    for table in inspector.get_table_names():
22        columns = [col["name"] for col in inspector.get_columns(table)]
23        schema[table] = columns
24    return schema
25