db_helper.py 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342
  1. import traceback
  2. import time
  3. import redis
  4. import psycopg2
  5. import pymysql
  6. from psycopg2 import pool as pgpool
  7. from config import set_config
  8. from log import Log
  9. config_ = set_config()
  10. log_ = Log()
  11. conn_redis = None
  12. class RedisHelper(object):
  13. def __init__(self):
  14. """
  15. 初始化redis连接信息
  16. redis_info: redis连接信息, 格式:dict, {'host': '', 'port': '', 'password': ''}
  17. """
  18. redis_info = config_.REDIS_INFO
  19. self.host = redis_info['host']
  20. self.port = redis_info['port']
  21. self.password = redis_info['password']
  22. def connect(self):
  23. """
  24. 连接redis
  25. :return: conn
  26. """
  27. global conn_redis
  28. if conn_redis is None:
  29. pool = redis.ConnectionPool(host=self.host,
  30. port=self.port,
  31. password=self.password,
  32. decode_responses=True)
  33. conn = redis.Redis(connection_pool=pool)
  34. conn_redis = conn
  35. return conn_redis
  36. def key_exists(self, key_name):
  37. """
  38. 判断key是否存在
  39. :param key_name: key
  40. :return: 存在-True, 不存在-False
  41. """
  42. conn = self.connect()
  43. return conn.exists(key_name)
  44. def del_keys(self, key_name):
  45. """
  46. 删除key
  47. :param key_name: key
  48. :return: None
  49. """
  50. conn = self.connect()
  51. conn.delete(key_name)
  52. def get_data_from_redis(self, key_name):
  53. """
  54. 读取redis中的数据
  55. :param key_name: key
  56. :return: data
  57. """
  58. conn = self.connect()
  59. if not conn.exists(key_name):
  60. # key不存在
  61. return None
  62. data = conn.get(key_name)
  63. return data
  64. def set_data_to_redis(self, key_name, value, expire_time=24*3600):
  65. """
  66. 新增数据
  67. :param key_name: key
  68. :param value: 元素的值 videoId
  69. :param expire_time: 过期时间,单位:s,默认1天
  70. :return: None
  71. """
  72. conn = self.connect()
  73. conn.set(key_name, value, ex=int(expire_time))
  74. def add_data_with_zset(self, key_name, data, expire_time=7*24*3600):
  75. """
  76. 新增数据,有序set
  77. :param key_name: key
  78. :param data: 元素的值及对应分数 type-dict {value: score}
  79. :param expire_time: 过期时间,单位:s,默认7天,type-int
  80. :return: None
  81. """
  82. conn = self.connect()
  83. conn.zadd(key_name, data)
  84. # 设置过期时间
  85. conn.expire(key_name, int(expire_time))
  86. def get_data_zset_with_index(self, key_name, start, end, desc=True, with_scores=False):
  87. """
  88. 根据索引位置获取元素的值
  89. :param key_name: key
  90. :param start: 索引起始点 闭区间,包含start
  91. :param end: 索引结束点 闭区间,包含end
  92. :param desc: 分数排序方式,默认从大到小
  93. :param with_scores: 是否获取元素的分数,默认 False,只获取元素的值
  94. :return: data 元素值列表(不包含分数),value(videoId)类型转换为int, 包含分数时不进行类型转换
  95. """
  96. st_get_zset = time.time()
  97. conn = self.connect()
  98. if not conn.exists(key_name):
  99. return None
  100. data = conn.zrange(key_name, start, end, desc, with_scores)
  101. et_get_zset = time.time()
  102. log_.info(f'get zest with index: key_name = {key_name}, executeTime = {(et_get_zset - st_get_zset) * 1000}')
  103. if with_scores:
  104. return data
  105. else:
  106. return [eval(value) for value in data]
  107. def get_score_with_value(self, key_name, value):
  108. """
  109. 在zset中,根据元素的value获取对应的score
  110. :param key_name: key
  111. :param value: 元素的值
  112. :return: score value对应的score
  113. """
  114. conn = self.connect()
  115. if not conn.exists(key_name):
  116. return None
  117. return conn.zscore(key_name, value)
  118. def update_score_with_value(self, key_name, value, score, expire_time=24*3600):
  119. """
  120. 在zset中,修改元素value对应的score
  121. :param key_name: key
  122. :param value: 元素的值
  123. :param score: value对应的score更新值
  124. :param expire_time: 过期时间,单位:s,默认1天,type-int
  125. """
  126. conn = self.connect()
  127. if conn.exists(key_name):
  128. conn.zadd(key_name, {value: score})
  129. else:
  130. # key不存在时,需设置过期时间
  131. conn.zadd(key_name, {value: score})
  132. conn.expire(key_name, int(expire_time))
  133. def remove_value_from_zset(self, key_name, value):
  134. """
  135. 删除zset中的指定元素
  136. :param key_name: key
  137. :param value: 元素的值
  138. :return: None
  139. """
  140. conn = self.connect()
  141. conn.zrem(key_name, value)
  142. def get_index_with_data(self, key_name, value):
  143. """
  144. 根据元素的值获取在有序set中的位置,按照分数倒序(从大到小)
  145. :param key_name: key
  146. :param value: 元素的值
  147. :return: idx 位置索引
  148. """
  149. conn = self.connect()
  150. return conn.zrevrank(key_name, value)
  151. def get_data_from_set(self, key_name):
  152. """
  153. 获取set中的所有数据
  154. :param key_name: key
  155. :return: data
  156. """
  157. conn = self.connect()
  158. if not conn.exists(key_name):
  159. # key不存在
  160. return None
  161. data = []
  162. cursor = 0
  163. while True:
  164. cur, temp = conn.sscan(key_name, cursor=cursor, count=2000)
  165. data.extend(temp)
  166. if cur == 0:
  167. break
  168. cursor = cur
  169. return list(set(data))
  170. def add_data_with_set(self, key_name, values, expire_time=30*60):
  171. """
  172. 新增数据,set
  173. :param key_name: key
  174. :param values: 要添加的元素 类型-tuple
  175. :param expire_time: 过期时间,单位:s,默认0.5小时 type-int
  176. :return: None
  177. """
  178. conn = self.connect()
  179. conn.sadd(key_name, *values)
  180. # 设置过期时间
  181. conn.expire(key_name, int(expire_time))
  182. def data_exists_with_set(self, key_name, value):
  183. """
  184. 判断元素value是否在集合key_name中
  185. :param key_name: key
  186. :param value: 需判断的元素
  187. :return: 存在-True, 不存在-False
  188. """
  189. conn = self.connect()
  190. return conn.sismember(key_name, value)
  191. def remove_value_from_set(self, key_name, values):
  192. """
  193. 删除set中的指定元素
  194. :param key_name: key
  195. :param values: 元素的值, 类型-tuple
  196. :return: None
  197. """
  198. conn = self.connect()
  199. conn.srem(key_name, *values)
  200. def decr_key(self, key_name, amount=1, expire_time=30*60):
  201. """
  202. redis自减
  203. :param key_name: key
  204. :param amount: 自减数,默认为1,type-int
  205. :param expire_time: 过期时间,单位:s,默认0.5小时 type-int
  206. :return: None
  207. """
  208. conn = self.connect()
  209. conn.decr(name=key_name, amount=amount)
  210. conn.expire(key_name, int(expire_time))
  211. def incr_key(self, key_name, amount=1, expire_time=30*60):
  212. """
  213. redis自增
  214. :param key_name: key
  215. :param amount: 自减数,默认为1,type-int
  216. :param expire_time: 过期时间,单位:s,默认0.5小时 type-int
  217. :return: None
  218. """
  219. conn = self.connect()
  220. conn.incr(name=key_name, amount=amount)
  221. conn.expire(key_name, int(expire_time))
  222. def setnx_key(self, key_name, value, expire_time=5*60):
  223. """
  224. 当key不存在时,将value塞入key中,key存在时不做操作
  225. :param key_name: key
  226. :param value: value
  227. :return: 过期时间,单位:s,默认5分钟 type-int
  228. """
  229. conn = self.connect()
  230. conn.setnx(name=key_name, value=value)
  231. conn.expire(name=key_name, time=int(expire_time))
  232. #hologres_info = config_.HOLOGRES_INFO
  233. #conn = psycopg2.connect(**hologres_info)
  234. #cur = conn.cursor()
  235. # holo连接
  236. # connectPool = pgpool.SimpleConnectionPool(1, 30, **config_.HOLOGRES_INFO)
  237. # class HologresHelper(object):
  238. # def __init__(self):
  239. # """初始化hologres连接信息"""
  240. # #self.hologres_info = config_.HOLOGRES_INFO
  241. #
  242. # def get_data(self, sql):
  243. # #global conn
  244. # # 连接Hologres
  245. # #conn = psycopg2.connect(**self.hologres_info)
  246. # conn = connectPool.getconn()
  247. # # 创建游标
  248. # cur = conn.cursor()
  249. # try:
  250. # # 查询数据
  251. # cur.execute(sql)
  252. # data = cur.fetchall()
  253. # cur.close()
  254. # connectPool.putconn(conn, close=False)
  255. # # 提交事务
  256. # #conn.commit()
  257. # # 释放资源
  258. # #cur.close()
  259. # #conn.close()
  260. # except Exception as e:
  261. # log_.error('hologress error...')
  262. # log_.error(traceback.format_exc())
  263. # cur.close()
  264. # connectPool.putconn(conn, close=False)
  265. # return []
  266. # return data
  267. class MysqlHelper(object):
  268. def __init__(self):
  269. """
  270. 初始化mysql连接信息
  271. """
  272. self.mysql_info = config_.MYSQL_INFO
  273. def get_data(self, sql):
  274. """
  275. 查询数据
  276. :param sql: sql语句
  277. :return: data
  278. """
  279. # 连接数据库
  280. conn = pymysql.connect(**self.mysql_info)
  281. # 创建游标
  282. cursor = conn.cursor()
  283. try:
  284. # 执行SQL语句
  285. cursor.execute(sql)
  286. # 获取查询的所有记录
  287. data = cursor.fetchall()
  288. except Exception as e:
  289. return None
  290. # 关闭游标对象
  291. cursor.close()
  292. # 关闭数据库连接
  293. conn.close()
  294. return data
  295. if __name__ == '__main__':
  296. redis_helper = RedisHelper()
  297. # key = 'flow.video.12345.123#112'
  298. # redis_helper.decr_key(key_name=key)
  299. # res = redis_helper.get_data_from_redis(key_name=key)
  300. # print(int(res), type(int(res)))
  301. # data = redis_helper.get_data_zset_with_index(key_name=config_.BOTTOM_KEY_NAME, start=0, end=-1)
  302. # print(len(data))
  303. # key_name = 'com.weiqu.video.hot.recommend.previewed.4.weixin_openid_otjoB5VG780SB4aVjYqBBNLb - X6M'
  304. # values = (6134455, 9772930, 9912678, 9901969, 9926876, 9904203, 2384831, 9932272, 9737653, 9925240)
  305. key_name = 'com.weiqu.video.hot.recommend.previewed.4.abcd1'
  306. # values = (9902612, 9905573, 9928264, 9932148, 9809440, 9919900, 6093379, 9917093, 9793537, 9814345)
  307. # redis_helper.add_data_with_set(key_name=key_name, values=values, expire_time=30 * 60)
  308. res = redis_helper.get_data_from_set(key_name=key_name)
  309. print(res)