(db_id, table, create_table_statement, table_upper=False)
| 58 | |
| 59 | |
| 60 | def get_foreign_keys(db_id, table, create_table_statement, table_upper=False): |
| 61 | foreign_keys = [] |
| 62 | for row in create_table_statement.split('\n'): |
| 63 | if row.lower().startswith("foreign key") and row.lower().count("foreign key") == 1: |
| 64 | if " on " in row: |
| 65 | row = row.split(" on ")[0] |
| 66 | if " ON " in row: |
| 67 | row = row.split(" ON ")[0] |
| 68 | # row.replace(" ON DELETE CASCADE", "") |
| 69 | # row.replace(" on delete cascade", "") |
| 70 | row = row.replace(",", " ").replace("(", " ").replace(")", " ") |
| 71 | row = row.split() |
| 72 | if len(row) != 6: # multiple keys |
| 73 | # print(db_id, create_table_statement) |
| 74 | for i, tok in enumerate(row): |
| 75 | if tok.lower() == "references": |
| 76 | references_pos = i |
| 77 | num_keys = references_pos - 2 |
| 78 | # print(num_keys) |
| 79 | # print(row) |
| 80 | for i in range(num_keys): |
| 81 | if table_upper == True: |
| 82 | key = f"{table.upper()}.{row[2 + i].lower()} = {row[4 + num_keys - 1].upper()}.{row[4 + num_keys + i].lower()}" |
| 83 | else: |
| 84 | key = f"{table}.{row[2 + i]} = {row[4 + num_keys - 1]}.{row[4 + num_keys + i]}" |
| 85 | foreign_keys.append(key) |
| 86 | # print(foreign_keys) |
| 87 | continue |
| 88 | else: |
| 89 | # assert len(row) == 6 |
| 90 | if table_upper == True: |
| 91 | key = f"{table.upper()}.{row[2].lower()} = {row[4].upper()}.{row[5].lower()}" |
| 92 | else: |
| 93 | key = f"{table}.{row[2]} = {row[4]}.{row[5]}" |
| 94 | foreign_keys.append(key) |
| 95 | return foreign_keys |
| 96 | |
| 97 | |
| 98 | def extract_tablecolumn_prompt(prompt_db, db_id, db_path, limit_value=3, normalization=True): |
no test coverage detected