From 90ef582f1068f78470a7e3b9733545ff7f626d71 Mon Sep 17 00:00:00 2001
From: chenwj113 <chenwj113@gmail.com>
Date: Tue, 2 May 2023 11:16:07 +0800
Subject: [PATCH] =?UTF-8?q?bugfix:=20=E4=BF=AE=E6=94=B9=E6=95=B0=E6=8D=AE?=
 =?UTF-8?q?=E5=BA=93=E8=BF=9E=E6=8E=A5=E9=94=99=E8=AF=AF?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 cron_tasks.py | 13 +++++++++----
 get_data.py   | 14 ++++++--------
 lottery.py    |  8 +++++++-
 3 files changed, 22 insertions(+), 13 deletions(-)

diff --git a/cron_tasks.py b/cron_tasks.py
index b2f9197..e1cf44e 100644
--- a/cron_tasks.py
+++ b/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)
-        last_id = lottery.insert(draw_issue, draw_date, draw_code)
-        if last_id:
-            print(f"issue:{draw_issue}数据写入完成。。。")
+        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():
diff --git a/get_data.py b/get_data.py
index 9fc456c..a53e21e 100644
--- a/get_data.py
+++ b/get_data.py
@@ -8,19 +8,17 @@ 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)
-    last_id = lottery.insert(draw_issue, draw_date, draw_code)
-    if last_id:
-        print(f"issue:{draw_issue}数据写入完成。。。")
-    else:
-        print(f'issue:{draw_issue}已经存在')
+    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}已经存在')
 
 def get_data(url, lottery_type):
     try:
diff --git a/lottery.py b/lottery.py
index 2213c09..2479e7b 100644
--- a/lottery.py
+++ b/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
\ No newline at end of file
+            return None
\ No newline at end of file