db.py 1.5 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061
  1. import json
  2. import pymysql
  3. def read(file: "str") -> "dict":
  4. with open(file, "r", encoding="utf-8") as fp:
  5. data = json.load(fp)
  6. fp.close()
  7. return data
  8. def insert(db, cursor, data: "list[tuple]"):
  9. # create sql
  10. values = []
  11. for it in data:
  12. if it[4] is not None: # village
  13. values.append(f"('{it[0]}', '{it[1]}', '{it[2]}', null, '{it[4]}')")
  14. else: # other
  15. values.append(f"('{it[0]}', '{it[1]}', '{it[2]}', '{it[3]}', null)")
  16. values_str = ", ".join(values)
  17. sql = f"INSERT INTO area_all(id, name, level, child, type) VALUES {values_str};"
  18. # do insert
  19. try:
  20. res = cursor.execute(sql)
  21. db.commit()
  22. print("success:", len(data), res)
  23. except Exception as e:
  24. print(e)
  25. print(sql)
  26. db.rollback()
  27. exit(-1)
  28. def split(arr: "list", batch: "int") -> "list[list]":
  29. size = len(arr)
  30. if size <= batch:
  31. return [arr]
  32. step = size // batch
  33. res = [arr[i * batch: (i + 1) * batch] for i in range(step)]
  34. if size % batch:
  35. res.append(arr[step * batch:])
  36. return res
  37. def main():
  38. db = pymysql.connect(host="47.96.4.54", port=10052, user="root", password="hm123456789", database="dam_subsidy")
  39. cursor = db.cursor()
  40. data = read("db_data-1.json")
  41. batch = 1024
  42. for k, ll in data.items():
  43. steps = split(ll, batch)
  44. for step in steps:
  45. insert(db, cursor, step)
  46. db.close()
  47. if __name__ == "__main__":
  48. main()