def write_backup_file()

in hugegraph-llm/src/hugegraph_llm/utils/hugegraph_utils.py [0:0]

• 19 lines of code
• 10 McCabe index (conditional complexity)

def write_backup_file(client, backup_subdir, filename, query, all_pk_flag):
    with open(os.path.join(backup_subdir, filename), "w", encoding="utf-8") as f:
        if filename == "edges.json":
            data = client.gremlin().exec(query)["data"][0]["edges"]
            json.dump(data, f, ensure_ascii=False)
        elif filename == "vertices.json":
            data_full = client.gremlin().exec(query)["data"][0]["vertices"]
            data = [{key: value for key, value in vertex.items() if key != "id"}
                    for vertex in data_full] if all_pk_flag else data_full
            json.dump(data, f, ensure_ascii=False)
        elif filename == "schema.json":
            data_full = query
            if isinstance(data_full, dict) and "schema" in data_full:
                groovy_filename = filename.replace(".json", ".groovy")
                with open(os.path.join(backup_subdir, groovy_filename), "w", encoding="utf-8") as groovy_file:
                    groovy_file.write(str(data_full["schema"]))
            else:
                data = data_full
                json.dump(data, f, ensure_ascii=False)