bugfix: 修改数据库连接错误

cron_tasks.py

@@ -9,6 +9,9 @@ from lottery import Lottery
 session = HTMLSession()
 
 def get_data_job(lottery_type='pls'):
+    """
+    :param lottery_type
+    """
     print("pls: I'm working...")
     url = f"https://kaijiang.500.com/{lottery_type}.shtml"
     r = session.get(url)
@@ -22,10 +25,12 @@ def get_data_job(lottery_type='pls'):
         draw_issue = HTML(html=issue_elem[0]).text
         draw_code = "".join(result)
         print(f"result:{draw_code}")
-        lottery = Lottery(lottery_type=lottery_type)
+        with Lottery(lottery_type=lottery_type) as lottery:
             last_id = lottery.insert(draw_issue, draw_date, draw_code)
             if last_id:
                 print(f"issue:{draw_issue}数据写入完成。。。")
+            else:
+                print(f'issue:{draw_issue}已经存在')
 
 @repeat(every().day.at("00:01"))
 def pls_job():

get_data.py

@@ -8,14 +8,12 @@ session = HTMLSession()
 
 @retry(stop_max_attempt_number=3)
 def _get_data(url, lottery_type):
-    lottery = Lottery(lottery_type=lottery_type)
     r = session.get(url)
     table_list = r.html.find("table.kj_tablelist02", first=True)
     draw_issue = table_list.find("td.td_title01 span.span_left strong")[0].text
     draw_date = table_list.find("td.td_title01 span.span_right")[0].text
     draw_code = table_list.find("div.ball_box01")[0].text.replace('\n', '')
-    print(draw_code)
-    print(draw_date)
+    with Lottery(lottery_type=lottery_type) as lottery:
         last_id = lottery.insert(draw_issue, draw_date, draw_code)
         if last_id:
             print(f"issue:{draw_issue}数据写入完成。。。")

lottery.py

@@ -33,6 +33,12 @@ class Lottery(object):
         _DbSession = sessionmaker(bind=_engine)
         return _DbSession()
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, type, value, trace):
+        self.db.close()
+
     def insert(self, draw_issue, draw_date,  draw_code, **kwargs):
         result = self.db.query(self._Model).filter_by(draw_issue=draw_issue).first()
         if result is None:
@@ -83,4 +89,4 @@ class Lottery(object):
             self.db.commit()
             return record.id
         else:
-            return result.id
+            return None