change pymysql to aiomysql

_playwright/bilibili.py

@@ -20,8 +20,6 @@ from __init__ import (
     bilibili_img,
 )
 
-bilibili_live_rooms = get_bilibili_live_rooms_from_pi()
-
 
 # 判断登录状态
 async def is_login(page):
@@ -68,6 +66,7 @@ async def live(page):
     if not await is_login(page):
         return "bilibili_not_login"
 
+    bilibili_live_rooms = await get_bilibili_live_rooms_from_pi()
     for room_id in list(bilibili_live_rooms.keys()):
         await page.goto(f"{bilibili_live_url}/{str(room_id)}")
         # 输入框选择 xpath 绝对定位
@@ -83,10 +82,14 @@ async def live(page):
         if await send_button.is_visible():
             # 点击发送按钮
             await send_button.click()
-            room_msg = f"直播间：【{bilibili_live_rooms[room_id]}】打卡成功，打卡时间：【{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}】"
+            room_msg = (
+                f"直播间：【{bilibili_live_rooms[room_id]}】打卡成功，打卡时间：【{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}】"
+            )
             bilibili_logger.info(room_msg)
         else:
-            room_msg = f"直播间：【{bilibili_live_rooms[room_id]}】打卡失败，打卡时间：【{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}】"
+            room_msg = (
+                f"直播间：【{bilibili_live_rooms[room_id]}】打卡失败，打卡时间：【{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}】"
+            )
             bilibili_logger.error(room_msg)
         msg.append(room_msg)
 

_requests/aliyundrive.py

@@ -1,4 +1,6 @@
 # 接口参考：https://github.com/Anonym-w/autoSigninAliyun/blob/main/autoSignin.js
+
+import asyncio
 import sys
 import requests
 
@@ -29,17 +31,21 @@ query_body = {}
 
 def get_refresh_token():
     # local storage only work in browser
-    return False
+    pass
 
 
-def validate_access_token():
+async def validate_access_token():
     # token只有三小时有效期，所以定时任务刷新间隔为三小时
     # 从树莓派获取 refresh_token
-    refresh_token = get_aliyundrive_refresh_token()
+    refresh_token = await get_aliyundrive_refresh_token()
     # 查询结构体
     query_body.update({"grant_type": "refresh_token", "refresh_token": refresh_token})
     # validate
-    resp = session.post(aliyundrive_access_token_url, headers=headers, json=query_body)
+    resp = session.post(
+        aliyundrive_access_token_url,
+        headers=headers,
+        json=query_body,
+    )
     if resp.status_code != 200:
         msg = "阿里云 access_token 失效"
         aliyundrive_logger.error(msg)
@@ -47,16 +53,18 @@ def validate_access_token():
     else:
         aliyundrive_logger.info("阿里云 access_token 有效")
         new_access_token = resp.json()["access_token"]
-        update_aliyundrive_access_token(new_access_token)
+        await update_aliyundrive_access_token(new_access_token)
         headers["Authorization"] = f"Bearer {new_access_token}"
-        return new_access_token
-    return False
 
 
 def get_aliyun_nickname():
     nickname = False
     # 获取用户名
-    user_res = session.post(aliyundrive_nickname_url, headers=headers, json=query_body)
+    user_res = session.post(
+        aliyundrive_nickname_url,
+        headers=headers,
+        json=query_body,
+    )
     if user_res.status_code == 200:
         user_msg = f"获取用户名成功，stdout：{user_res.json()}"
         nickname = user_res.json()["nick_name"]
@@ -69,7 +77,11 @@ def get_aliyun_nickname():
 
 def sign():
     # 签到
-    sign_res = session.post(aliyundrive_signin_url, headers=headers, json=query_body)
+    sign_res = session.post(
+        aliyundrive_signin_url,
+        headers=headers,
+        json=query_body,
+    )
     if sign_res.status_code != 200:
         sign_msg = f"阿里云盘签到失败，stderr：{sign_res.json()}"
         aliyundrive_logger.error(sign_msg)
@@ -95,7 +107,11 @@ def sign():
 # 领取奖励
 def get_reward(signin_count):
     query_body.update({"signInDay": signin_count})
-    reward_res = session.post(aliyundrive_reward_url, headers=headers, json=query_body)
+    reward_res = session.post(
+        aliyundrive_reward_url,
+        headers=headers,
+        json=query_body,
+    )
     if reward_res.status_code == 200:
         aliyundrive_logger.info("领取奖励成功")
     else:
@@ -106,8 +122,8 @@ def get_reward(signin_count):
     return False
 
 
-def aliyundrive_sign():
-    validate_access_token()
+async def aliyundrive_sign():
+    await validate_access_token()
     nickname = get_aliyun_nickname()
     is_reward, signin_count, curent_reward = sign()
     if not is_reward:
@@ -115,8 +131,7 @@ def aliyundrive_sign():
     message = f"{nickname} 签到成功，本月累计签到 {signin_count} 天\n本次签到获得：{curent_reward}"
     aliyundrive_logger.info(message)
     push_msg(title=aliyundrive_title, message=message, img_url=aliyundrive_img)
-    return message
 
 
 if __name__ == "__main__":
-    aliyundrive_sign()
+    asyncio.run(aliyundrive_sign())

api.py

@@ -17,7 +17,6 @@ from _playwright.vits import vits_sign
 from _playwright.tsdm import tsdm_sign
 from _playwright.wuaipojie import wuaipojie_sign
 from _requests.aliyundrive import aliyundrive_sign
-from utils import get_cookie_from_pi
 
 
 def tick():
@@ -26,8 +25,6 @@ def tick():
 
 def run_scheduler(loop):
     asyncio.set_event_loop(loop)
-    # pi
-    get_cookie_from_pi()
     # add task
     sign_scheduler = AsyncIOScheduler()
     sign_scheduler.configure(timezone=pytz.timezone("Asia/Shanghai"))

app.py

@@ -10,7 +10,7 @@ from fastapi.responses import Response
 from loguru import logger
 
 from __init__ import DIR_PATH, api_token
-from utils import push_msg
+from utils import push_msg, read_logfile, get_cookie_from_pi
 from api import start_daemon_scheduler, kill_daemon_scheduler
 
 schedule_loop = asyncio.new_event_loop()
@@ -26,10 +26,11 @@ def home():
 
 # 开启定时任务
 @app.get("/run")
-def do_run(token: str):
+async def do_run(token: str):
     if token != api_token:
         return {"code": 422}
 
+    await get_cookie_from_pi()
     start_daemon_scheduler(schedule_loop)
     res = {"code": 0, "data": "start scheduler"}
     return res
@@ -53,9 +54,7 @@ async def do_log(token: str, name: str):
         return {"code": 422}
 
     log_file = join(DIR_PATH, f"log/{name}.log")
-    with open(log_file, "r", encoding="utf-8") as f:
-        data = f.read()
-
+    data = await read_logfile(log_file)
     res = {"code": 200, "data": data}
     return res
 

config/settings_prod.py

@@ -13,7 +13,9 @@ from option_logger import register_logger
 # playwright
 browser_headless = True
 browser_proxy = {"server": None}
-user_agent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/114.0.0.0 Safari/537.36"
+user_agent = (
+    "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/114.0.0.0 Safari/537.36"
+)
 ignore_https_errors = True
 
 # api
@@ -74,9 +76,7 @@ bilibili_title = "Bilibili"
 bilibili_logger = register_logger("bilibili", logfile_mapping["bilibili"])
 bilibili_cookie_file = join(PROD_PATH, "cookie/bilibili.json")
 bilibili_url = "https://account.bilibili.com/account/home"
-bilibili_sign_url = (
-    "https://api.live.bilibili.com/xlive/web-ucenter/v1/sign/DoSign"  # 签到地址
-)
+bilibili_sign_url = "https://api.live.bilibili.com/xlive/web-ucenter/v1/sign/DoSign"  # 签到地址
 bilibili_live_url = "https://live.bilibili.com"  # 直播间地址
 bilibili_live_rooms = os.environ.get("bilibili_live_rooms")
 bilibili_img = "https://img.evimo.top/bilibili.png"

option_mysql.py

@@ -1,54 +1,67 @@
-import pymysql
-
-
-class OptionMysql(object):
-    def __init__(self, options):
-        host = options["HOST"]
-        user = options["USERNAME"]
-        password = options["PASSWORD"]
-        database = options["DATABASE"]
-        port = options["PORT"]
-        charset = "utf8"
-        # 连接数据库
-        self.conn = pymysql.connect(
-            host=host,
-            port=port,
-            user=user,
-            password=password,
-            database=database,
-            charset=charset,
-        )
-        # 创建游标
-        self.cur = self.conn.cursor()
-        self.dict_cur = self.conn.cursor(cursor=pymysql.cursors.DictCursor)
-
-    def __exit__(self):
-        self.cur.close()
-        self.dict_cur.close()
-        self.conn.close()
-
-    def get_dict_data(self, sql, params=[]):
-        """查询，返回字典类型"""
+import aiomysql
+
+from __init__ import mysql_option
+
+
+class AIOMYSQL:
+    def __init__(self) -> None:
+        self.pool = None
+
+    async def init_pool(self):
+        try:
+            __pool = await aiomysql.create_pool(
+                host=mysql_option["HOST"],
+                port=mysql_option["PORT"],
+                user=mysql_option["USERNAME"],
+                password=mysql_option["PASSWORD"],
+                db=mysql_option["DATABASE"],
+                charset="utf8",
+                autocommit=False,
+                minsize=5,
+                maxsize=10,
+                cursorclass=aiomysql.DictCursor,
+            )
+            return __pool
+        except:
+            raise ("aiomysql create_pool error")
+
+    async def get_cursor(self):
+        conn = await self.pool.acquire()
+        cur = await conn.cursor()
+        return conn, cur
+
+    async def query(self, sql, param=None):
+        conn, cur = await self.get_cursor()
         try:
-            if params:
-                self.dict_cur.execute(sql, params)
-            else:
-                self.dict_cur.execute(sql)
-            data = self.dict_cur.fetchall()
-            return data
-        except Exception as e:
-            self.conn.rollback()
-            raise e
-
-    def update_data(self, sql, params=[]):
-        """更新"""
+            await cur.execute(sql, param)
+        except:
+            await cur.rollback()
+            print("aiomysql query error")
+        else:
+            return await cur.fetchall()
+        finally:
+            if cur:
+                await cur.close()
+            # 释放 conn 连接回连接池
+            await self.pool.release(conn)
+
+    async def update(self, sql, param=None):
+        conn, cur = await self.get_cursor()
         try:
-            if not params:
-                self.cur.execute(sql)
-            else:
-                self.cur.execute(sql, params)
-            self.conn.commit()
-        except Exception as e:
-            self.conn.rollback()
-            raise e
-        return True
+            await cur.execute(sql, param)
+        except:
+            await cur.rollback()
+            print("aiomysql query error")
+        else:
+            await conn.commit()
+        finally:
+            if cur:
+                await cur.close()
+            # 释放 conn 连接回连接池
+            await self.pool.release(conn)
+
+
+async def get_aiomysql_instance():
+    instance = AIOMYSQL()
+    instance.pool = await instance.init_pool()
+    return instance

requirements.txt

@@ -10,6 +10,8 @@ uvicorn
 fastapi
 pytz
 pymysql
+aiomysql
+aiofiles
 #ddddocr
 #Pillow
 #selenium

utils.py

@@ -1,8 +1,7 @@
 # coding:utf-8
-import os
-import sys
 import socket
 import asyncio
+import aiofiles
 import psutil
 import requests
 
@@ -12,13 +11,12 @@ from loguru import logger
 from playwright.async_api import async_playwright, Page
 from playwright_stealth import stealth_async
 
-from option_mysql import OptionMysql
+from option_mysql import get_aiomysql_instance
 from __init__ import (
     gotify_host,
     gotify_port,
     gotify_token,
     gotify_img,
-    mysql_option,
     DIR_PATH,
     DEBUG,
     browser_headless,
@@ -52,7 +50,7 @@ async def init_page(storage_state=None, browser_type: str = "chromium") -> Page:
         page = await context.new_page()
         # 隐藏 webdriver 特征
         await stealth_async(page)
-        # await page.pause()  # async 跳出上下文就会 playwright._impl._api_types.Error: Connection closed
+        # await page.pause()  # async_playwright 跳出上下文就会 playwright._impl._api_types.Error: Connection closed
         await goto_github(page)
         await dump_cookie(page)  # 创建函数：一个 browser 对应一个 context、page，在 with 内完成多个任务
         await browser.close()
@@ -63,22 +61,14 @@ async def goto_github(page):
 
 
 async def dump_cookie(page):
-    # url = "https://www.52pojie.cn/home.php"
-    # url = "https://v2ex.com"
-    # url = "https://www.bilibili.com"
-    # url = "https://www.tsdm39.com/forum.php?mobile=no"
-    # url = "https://www.aliyundrive.com/drive"
     url = "https://huggingface.co/login"
     save_cookie = join(DIR_PATH, f"{get_domain(url)}.json")
     await page.goto(url)
     logger.debug(DIR_PATH)
-    await page.screenshot(
-        path=join(DIR_PATH, "screenshot.png"),
-        full_page=True,
-    )  # 截图
-    while input("input c to continue: ") != "c":  # 阻塞等待填写登录信息
+    await page.screenshot(path=join(DIR_PATH, "screenshot.png"), full_page=True)  # 截图
+    while input("input c to continue: ") != "c":  # 阻塞等待填完登录信息
         continue
-    await page.context.storage_state(path=save_cookie)  # 保存 Cookie 到文件
+    await page.context.storage_state(path=save_cookie)  # 保存 Cookie 到本地
 
 
 # 获取容器内 IP
@@ -101,91 +91,88 @@ def get_domain(url: str) -> str:
 
 
 # 从域名获取 IP
-def get_domain_ip(host: str) -> str:
+def get_pi_ip(host: str) -> str:
     domain_ip = socket.gethostbyname(host)
     return domain_ip
 
 
 # gotify 推送
-def push_msg(title: str = "无标题", message: str = "无内容", img_url: str = gotify_img):
-    gotify_ip = get_domain_ip(gotify_host)
+def push_msg(title: str = "无标题", message: str = "无内容", img_url: str = gotify_img) -> dict:
+    gotify_ip = get_pi_ip(gotify_host)
     url = f"http://{gotify_ip}:{gotify_port}/message?token={gotify_token}"
-    resp = requests.post(
-        url,
-        json={
-            "title": title,
-            "message": message,
-            "priority": 10,
-            "extras": {
-                "client::display": {"contentType": "text/markdown"},
-                "client::notification": {"bigImageUrl": img_url},
-            },
+    data = {
+        "title": title,
+        "message": message,
+        "priority": 10,
+        "extras": {
+            "client::display": {"contentType": "text/markdown"},
+            "client::notification": {"bigImageUrl": img_url},
         },
-    )
-    return resp.json()
+    }
+    resp = requests.post(url=url, json=data, timeout=20)
+    return {"code": resp.status_code, "res": resp.content}
 
 
-def get_bilibili_live_rooms_from_pi() -> dict:
-    logger.info("getting bilibili_live_rooms from pi")
-    bilibili_live_rooms = {}
-    mysql = OptionMysql(mysql_option)
+async def get_bilibili_live_rooms_from_pi() -> dict:
     sql = f"SELECT room_id, username FROM live_rooms"
-    data = mysql.get_dict_data(sql)
-    if data:
-        for d in data:
-            bilibili_live_rooms[d["room_id"]] = d["username"]
+    pool = await get_aiomysql_instance()
+    data = await pool.query(sql)
+    bilibili_live_rooms = data[0]
     return bilibili_live_rooms
 
 
-def get_cookie_from_pi() -> dict:
-    logger.info("getting cookie from pi")
-    mysql = OptionMysql(mysql_option)
+async def get_cookie_from_pi() -> dict:
     sql = f"SELECT cookie_name, cookie_value FROM cookie"
-    data = mysql.get_dict_data(sql)
+    pool = await get_aiomysql_instance()
+    data = await pool.query(sql)
     for d in data:
         log_file = join(DIR_PATH, f"cookie/{d['cookie_name']}.json")
-        with open(log_file, "w") as lf:
-            lf.write(d["cookie_value"])
+        async with aiofiles.open(log_file, mode="w") as handle:
+            await handle.write(d["cookie_value"])
     return data[0]["cookie_name"]
 
 
-def get_aliyundrive_refresh_token():
-    mysql = OptionMysql(mysql_option)
+async def get_aliyundrive_refresh_token():
     sql = f"SELECT token FROM aliyundrive WHERE id=1"
-    data = mysql.get_dict_data(sql)
+    pool = await get_aiomysql_instance()
+    data = await pool.query(sql)
     if data:
         logger.info("获取阿里云盘 refresh_token 成功")
         return data[0]["token"]
     return False
 
 
-def update_aliyundrive_access_token(access_token):
-    mysql = OptionMysql(mysql_option)
+async def update_aliyundrive_access_token(access_token):
     sql = f"UPDATE aliyundrive SET token=%s WHERE id=2"
-    data = mysql.update_data(sql, [access_token])
+    pool = await get_aiomysql_instance()
+    data = await pool.update(sql, [access_token])
     if data:
         logger.info("update aliyundrive access token")
         return True
     return False
 
 
-def get_log(log_file: str) -> str:
-    with open(log_file, "r", encoding="utf-8") as f:
-        data = f.read()
+async def read_logfile(log_file: str) -> str:
+    async with aiofiles.open(log_file, mode="r", encoding="utf-8") as f:
+        data = await f.read()
         return data
 
 
+async def main():
+    # await init_page()
+    # res = await get_bilibili_live_rooms_from_pi()
+    # res = await get_cookie_from_pi()
+    # res = await get_aliyundrive_refresh_token()
+    # res = await update_aliyundrive_access_token("no use")
+    # res = await get_bilibili_live_rooms_from_pi()
+    res = await read_logfile("log/bilibili.log")
+    logger.debug(f"utils test: {res}")
+
+
 if __name__ == "__main__":
     logger.debug(f"DEBUG Mode: {DEBUG}")
     logger.debug(f"DIR_PATH: {DIR_PATH}")
-    # asyncio.run(init_page())
-    # logger.debug(f"get_domain_ip: {get_domain_ip(gotify_host)}")
-    # logger.debug(f'get_domain: {get_domain("https://github.com")}')
-    # logger.debug(f"list_path: {list_path()}")
+    logger.debug(f"get_pi_ip: {get_pi_ip(gotify_host)}")
+    logger.debug(f'get_domain: {get_domain("https://github.com")}')
     logger.debug(f'push_msg: {push_msg("utils.py", "hello world")}')
-    # logger.debug(f"get_aliyun_token_from_pi(): {get_aliyun_token_from_pi()}")
-    # logger.debug(f"get_bilibili_live_rooms_from_pi(): {get_bilibili_live_rooms_from_pi()}")
-    # logger.debug(f"get_cookie_from_pi(): {get_cookie_from_pi()}")
-    # logger.debug(f'get_cron_log(): {get_cron_log("log/tsdm.log")}')
-    # logger.debug(f"get_aliyundrive_access_token: {get_aliyundrive_access_token()}")
-    # logger.debug(f"update_aliyundrive_access_token: {update_aliyundrive_access_token('123')}")
+    # asyncio.run(main())