(prompt_db, db_id, db_path, limit_value=3, normalization=True)
| 96 | |
| 97 | |
| 98 | def extract_tablecolumn_prompt(prompt_db, db_id, db_path, limit_value=3, normalization=True): |
| 99 | table_query = "SELECT * FROM sqlite_master WHERE type='table';" |
| 100 | tables = sqlite3.connect(db_path).cursor().execute(table_query).fetchall() |
| 101 | prompt = "" |
| 102 | foreign_keys = [] |
| 103 | table_names = [] |
| 104 | for table in tables: |
| 105 | table_name = table[1] |
| 106 | if table_name == 'sqlite_sequence': |
| 107 | continue |
| 108 | table_names.append(table_name) |
| 109 | if normalization: |
| 110 | table_name = table_name.lower() |
| 111 | create_table_statement = table[-1] |
| 112 | create_table_statement = normalize_create_table(table_name, create_table_statement) |
| 113 | foreign_keys_one_table = get_foreign_keys(db_id, table_name, create_table_statement) |
| 114 | table_info_query = f"PRAGMA table_info({table_name});" |
| 115 | headers = [x[1] for x in sqlite3.connect(db_path).cursor().execute(table_info_query).fetchall()] |
| 116 | if normalization: |
| 117 | foreign_keys_one_table = [x.lower() for x in foreign_keys_one_table] |
| 118 | headers = [x.lower() for x in headers] |
| 119 | foreign_keys.extend(foreign_keys_one_table) |
| 120 | table_statement = "" |
| 121 | if prompt_db.startswith("Table(Columns)"): |
| 122 | table_statement += f"{table_name}({', '.join(headers)});\n" |
| 123 | if prompt_db.startswith("Columns=[]"): |
| 124 | table_statement += f"Table {table_name}, Columns = [{', '.join(headers)}];\n" |
| 125 | prompt += table_statement |
| 126 | if "+FK" in prompt_db: |
| 127 | prompt += "Foreign_keys = [" + ', '.join(foreign_keys) + "];\n" |
| 128 | prompt += '\n' |
| 129 | return prompt, table_names |
| 130 | |
| 131 | |
| 132 | def extract_create_table_prompt(prompt_db, db_id, db_path, limit_value=3, normalization=True): |
no test coverage detected