CoolFace
Apppublic

LightRT/text2sql_backend

sourceHugging Faceupdated 1mo agoView on Hugging Face
0likes
scheme.py36 linesDownload Raw Back to src
1from sqlalchemy import create_engine , inspect2from langchain_core.documents import Document3 4def create_scheme(database_url : str) -> list[dict] :5 6    scheme_docs = []7 8    engine = create_engine(database_url)9    inspector = inspect(engine)10 11    tables = inspector.get_table_names()12 13    for table in tables :14        clean_table = table.replace("_" , " ")15 16        columns = inspector.get_columns(table)17 18        clean_columns = []19 20        for col in columns :21            clean_col = col['name'].replace("_" , " ")22            clean_columns.append(clean_col)23        24        doc = f"Table: {clean_table}.\nColumns: {', '.join(clean_columns)}"25 26        scheme_docs.append(27            Document(28                page_content=doc,29                metadata={30                    "table_name" : table31                }32            )33        )34    35    return scheme_docs36