TeeA's picture
Upload script/convert.py with huggingface_hub
4174c27 verified
raw
history blame
1.81 kB
def convert_obj_to_schema(obj:dict) -> str:
schema = f"CREATE DATABASE {obj['db_id']}; USE {obj['db_id']};"
tables = {}
for table_name in obj['table_names']:
table_name = "_".join(table_name.split(" "))
tables[table_name] = []
for col_idx, column in enumerate(obj['column_names']):
table_id, column_name = column
# _ combination
column_name = "_".join(column_name.split(" "))
if table_id == -1:
continue
table_name = obj['table_names'][table_id]
column_type = obj['column_types'][col_idx]
# _ combination
table_name = "_".join(table_name.split(" "))
column_attributes = [column_name, column_type]
if col_idx in obj['primary_keys']:
column_attributes.append("PRIMARY KEY")
tables[table_name].append(column_attributes)
for fk in obj['foreign_keys']:
if col_idx == fk[1]:
table_id_other, column_name_other = obj['column_names'][fk[0]]
table_name_other = obj['table_names'][table_id_other]
column_name_other = "_".join(column_name_other.split(" "))
table_name_other = "_".join(table_name_other.split(" "))
tables[table_name].append([f"FOREIGN KEY {column_name} REFERENCES {table_name_other}({column_name_other})"])
for table_name in tables:
column_attributes = ""
for idx, column_attribute in enumerate(tables[table_name]):
column_attributes += " ".join(column_attribute) + ","
column_attributes = column_attributes[:-1] # remove endless ','
schema += f" CREATE TABLE {table_name}({column_attributes});"
return {
"db_id": obj['db_id'],
"schema": schema
}