Get the database schema from the database URI
(root_dir, dataset, db_id, limit_value=3)
| 396 | |
| 397 | |
| 398 | def generate_db_prompt_bird(root_dir, dataset, db_id, limit_value=3) -> str: |
| 399 | """Get the database schema from the database URI |
| 400 | """ |
| 401 | DATA_PATH = root_dir |
| 402 | dev_db_path = f"{DATA_PATH}/{dataset}/dev_databases" |
| 403 | db_uri = dev_db_path + "/" + db_id + "/" + db_id + ".sqlite" |
| 404 | # table_path = f"{DATA_PATH}/{dataset}/dev_tables.json" |
| 405 | # with open(table_path, "r") as f: |
| 406 | # tables = json.load(f) |
| 407 | # table_names = tables[db_id]["table_names"] |
| 408 | conn = sqlite3.connect(db_uri) |
| 409 | |
| 410 | # Create a cursor object |
| 411 | cursor = conn.cursor() |
| 412 | cursor.execute("SELECT name FROM sqlite_master WHERE type='table'") |
| 413 | tables = cursor.fetchall() |
| 414 | table_names = [] |
| 415 | for table in tables: |
| 416 | if table == 'sqlite_sequence': |
| 417 | continue |
| 418 | if isinstance(table, tuple): |
| 419 | table_name = table[0] |
| 420 | table_names.append(table_name) |
| 421 | db = SQLDatabase.from_uri("sqlite:///" + db_uri) |
| 422 | db._sample_rows_in_table_info = limit_value |
| 423 | return db.get_table_info_no_throw(), table_names |
| 424 | |
| 425 | |
| 426 | def get_db_schemas(bench_root: str, db_name: str): |