common_func.py 1.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445
  1. from sqlalchemy import create_engine
  2. from urllib import parse
  3. import json
  4. import numpy as np
  5. # 创建数据库连接引擎
  6. def get_db_engine(db_info):
  7. db_con_str = 'mysql+pymysql://%s:%s@%s:%d/%s' % \
  8. (db_info['username'],
  9. parse.quote_plus(db_info['password']),
  10. db_info['host'],
  11. db_info['port'],
  12. db_info['database'])
  13. engine = create_engine(db_con_str, connect_args={'charset': 'utf8'})
  14. return engine
  15. # insert into 替换为 replace into 存在重复唯一键时,先删除此行数据,然后插入新的数据
  16. def mysql_replace_into(table, conn, keys, data_iter):
  17. from sqlalchemy.dialects.mysql import insert
  18. from sqlalchemy.ext.compiler import compiles
  19. from sqlalchemy.sql.expression import Insert
  20. @compiles(Insert)
  21. def replace_string(insert, compiler, **kw):
  22. s = compiler.visit_insert(insert, **kw)
  23. s = s.replace("INSERT INTO", "REPLACE INTO")
  24. return s
  25. data = [dict(zip(keys, row)) for row in data_iter]
  26. conn.execute(table.table.insert(replace_string=""), data)
  27. class NpEncoder(json.JSONEncoder):
  28. def default(self, obj):
  29. if isinstance(obj, np.integer):
  30. return int(obj)
  31. elif isinstance(obj, np.floating):
  32. return float(obj)
  33. elif isinstance(obj, np.ndarray):
  34. return obj.tolist()
  35. else:
  36. return super(NpEncoder, self).default(obj)