feat: 增加配置项支持自由选择数据是否保存到关系型数据库中

Author: NanmiCoder

README.md

@@ -21,20 +21,22 @@
 - [x] 抖音登录（二维码、手机号、cookies）
 - [x] 抖音滑块（模拟滑动实现，准确率不太OK）
 - [x] 支持登录成功后的上下文浏览器环境保留
+- [x] 数据持久化到硬盘（关系型数据库）
 
-## 待实现
-
-- [ ] 数据持久化到硬盘
 
 ## 使用方法
 
 1. 安装依赖库
    `pip install -r requirements.txt`
 2. 安装playwright浏览器驱动
    `playwright install`
-3. 运行爬虫程序
+3. 是否选择开启保存数据到DB中   
+   如果选择开启，则需要配置数据库连接信息，`config/db_config.py` 中的 `IS_SAVED_DATABASED`和`RELATION_DB_URL` 变量
+   <br>再执行 `python db.py` 初始化数据库信息，生成相关的数据库表结构
+4.  运行爬虫程序
    `python main.py --platform xhs --lt qrcode`
-4. 打开小红书扫二维码登录
+5. 打开对应APP扫二维码登录
+
 
 ## 项目代码结构
 
@@ -67,24 +69,16 @@ MediaCrawler
 │       ├── help.py             # 辅助函数
 │       └── login.py            # 登录实现
 ├── modles 
-│   ├── douyin
-│   │   └── m_douyin.py
-│   └── xhs
-│       └── m_xhs.py
+│   ├── douyin.py               # 抖音数据模型
+│   └── xiaohongshu.py          # 小红书数据模型 
 ├── tools
 │   └── utils.py                # 工具函数
 ├── main.py                     # 程序入口
 └── recv_sms_notification.py    # 短信转发器的HTTP SERVER接口
 ```
+## 数据持久化
 
-## 小红书运行截图
-
-![小红书运行截图](https://s2.loli.net/2023/06/09/PVBe3X5vf4yncrd.gif)
-
-## 抖音运行截图
-
-- ![抖音运行截图](https://s2.loli.net/2023/06/25/GXfkeLhpTyNiAqH.gif)
-
+![数据持久化](https://s2.loli.net/2023/07/24/ZTcGWz8jPAy7b5M.png)
 
 ## 支持一下
 

config/__init__.py

@@ -1,2 +1,3 @@
 from .base_config import *
 from .account_config import *
+from .db_config import *

config/base_config.py

@@ -2,11 +2,7 @@
 PLATFORM = "xhs"
 KEYWORDS = "健身,旅游"
 LOGIN_TYPE = "qrcode"  # qrcode or phone or cookies
-COOKIES = "web_session=xxxxcfed1566xxxxxxxxxxxxxxxxxxx;" # if platform is xhs, pleas set only web_session cookie attr
-
-# redis config
-REDIS_DB_HOST = "redis://127.0.0.1"  # your redis host
-REDIS_DB_PWD = "123456"  # your redis password
+COOKIES = "web_session=xxxxcfed1566xxxxxxxxxxxxxxxxxxx;"  # if platform is xhs, pleas set only web_session cookie attr
 
 # enable ip proxy
 ENABLE_IP_PROXY = False
@@ -18,7 +14,7 @@
 HEADLESS = True
 
 # save login state
-SAVE_LOGIN_STATE = False
+SAVE_LOGIN_STATE = True
 
 # save user data dir
 USER_DATA_DIR = "%s_user_data_dir"  # %s will be replaced by platform name

config/db_config.py

@@ -0,0 +1,9 @@
+# redis config
+REDIS_DB_HOST = "redis://127.0.0.1"  # your redis host
+REDIS_DB_PWD = "123456"  # your redis password
+
+# mysql config
+RELATION_DB_URL = "mysql://root:youdbpassword@localhost:3306/media_crawler"
+
+# save data to database option
+IS_SAVED_DATABASED = True  # if you want to save data to database, set True

db.py

@@ -0,0 +1,24 @@
+from tortoise import Tortoise
+from tortoise import run_async
+
+from config.db_config import *
+
+from tools import utils
+
+
+async def init_db(create_db: bool = False) -> None:
+    await Tortoise.init(
+        db_url=RELATION_DB_URL,
+        modules={'models': ['models']},
+        _create_db=create_db
+    )
+
+
+async def init():
+    await init_db(create_db=True)
+    await Tortoise.generate_schemas()
+    utils.logger.info("Init DB Success!")
+
+
+if __name__ == '__main__':
+    run_async(init())

images/douyin.gif

@@ -2,8 +2,8 @@
 import asyncio
 import argparse
 
+import db
 import config
-from tools import utils
 from base import proxy_account_pool
 from media_platform.douyin import DouYinCrawler
 from media_platform.xhs import XiaoHongShuCrawler
@@ -29,6 +29,10 @@ async def main():
     # init account pool
     account_pool = proxy_account_pool.create_account_pool()
 
+    # init db
+    if config.IS_SAVED_DATABASED:
+        await db.init_db()
+
     args = parser.parse_args()
     crawler = CrawlerFactory().create_crawler(platform=args.platform)
     crawler.init_config(

images/xiaoshongshu.gif

@@ -23,12 +23,12 @@ class DouYinCrawler(AbstractCrawler):
     dy_client: DOUYINClient
 
     def __init__(self) -> None:
-        self.browser_context: Optional[BrowserContext] = None # type: ignore
-        self.context_page: Optional[Page] = None # type: ignore
+        self.browser_context: Optional[BrowserContext] = None  # type: ignore
+        self.context_page: Optional[Page] = None  # type: ignore
         self.user_agent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/109.0.0.0 Safari/537.36"  # fixed
         self.index_url = "https://www.douyin.com"
-        self.command_args: Optional[Namespace] = None # type: ignore
-        self.account_pool: Optional[AccountPool] = None # type: ignore
+        self.command_args: Optional[Namespace] = None  # type: ignore
+        self.account_pool: Optional[AccountPool] = None  # type: ignore
 
     def init_config(self, **kwargs):
         for key, value in kwargs.items():
@@ -53,7 +53,7 @@ async def start(self) -> None:
             self.dy_client = await self.create_douyin_client(httpx_proxy)
             if not await self.dy_client.ping(browser_context=self.browser_context):
                 login_obj = DouYinLogin(
-                    login_type=self.command_args.lt, # type: ignore
+                    login_type=self.command_args.lt,  # type: ignore
                     login_phone=account_phone,
                     browser_context=self.browser_context,
                     context_page=self.context_page,
@@ -88,35 +88,37 @@ async def search_posts(self) -> None:
                                            post_item.get("aweme_mix_info", {}).get("mix_items")[0]
                     except TypeError:
                         continue
-                    aweme_list.append(aweme_info.get("aweme_id",""))
+                    aweme_list.append(aweme_info.get("aweme_id", ""))
                     await douyin.update_douyin_aweme(aweme_item=aweme_info)
             utils.logger.info(f"keyword:{keyword}, aweme_list:{aweme_list}")
-            # await self.batch_get_note_comments(aweme_list)
+            await self.batch_get_note_comments(aweme_list)
 
     async def batch_get_note_comments(self, aweme_list: List[str]):
         task_list: List[Task] = []
+        _semaphore = asyncio.Semaphore(config.MAX_CONCURRENCY_NUM)
         for aweme_id in aweme_list:
-            task = asyncio.create_task(self.get_comments(aweme_id), name=aweme_id)
+            task = asyncio.create_task(self.get_comments(aweme_id, _semaphore), name=aweme_id)
             task_list.append(task)
         await asyncio.wait(task_list)
 
-    async def get_comments(self, aweme_id: str):
-        try:
-            await self.dy_client.get_aweme_all_comments(
-                aweme_id=aweme_id,
-                callback=douyin.batch_update_dy_aweme_comments
-            )
-            utils.logger.info(f"aweme_id: {aweme_id} comments have all been obtained completed ...")
-        except DataFetchError as e:
-            utils.logger.error(f"aweme_id: {aweme_id} get comments failed, error: {e}")
+    async def get_comments(self, aweme_id: str, semaphore: "asyncio.Semaphore"):
+        async with semaphore:
+            try:
+                await self.dy_client.get_aweme_all_comments(
+                    aweme_id=aweme_id,
+                    callback=douyin.batch_update_dy_aweme_comments
+                )
+                utils.logger.info(f"aweme_id: {aweme_id} comments have all been obtained completed ...")
+            except DataFetchError as e:
+                utils.logger.error(f"aweme_id: {aweme_id} get comments failed, error: {e}")
 
     def create_proxy_info(self) -> Tuple[Optional[str], Optional[Dict], Optional[str]]:
         """Create proxy info for playwright and httpx"""
         if not config.ENABLE_IP_PROXY:
             return None, None, None
 
         # phone: 13012345671  ip_proxy: 111.122.xx.xx1:8888
-        phone, ip_proxy = self.account_pool.get_account() # type: ignore
+        phone, ip_proxy = self.account_pool.get_account()  # type: ignore
         playwright_proxy = {
             "server": f"{config.IP_PROXY_PROTOCOL}{ip_proxy}",
             "username": config.IP_PROXY_USER,
@@ -127,7 +129,7 @@ def create_proxy_info(self) -> Tuple[Optional[str], Optional[Dict], Optional[str
 
     async def create_douyin_client(self, httpx_proxy: Optional[str]) -> DOUYINClient:
         """Create douyin client"""
-        cookie_str, cookie_dict = utils.convert_cookies(await self.browser_context.cookies()) # type: ignore
+        cookie_str, cookie_dict = utils.convert_cookies(await self.browser_context.cookies())  # type: ignore
         douyin_client = DOUYINClient(
             proxies=httpx_proxy,
             headers={
@@ -152,18 +154,19 @@ async def launch_browser(
     ) -> BrowserContext:
         """Launch browser and create browser context"""
         if config.SAVE_LOGIN_STATE:
-            user_data_dir = os.path.join(os.getcwd(), "browser_data", config.USER_DATA_DIR % self.command_args.platform) # type: ignore
+            user_data_dir = os.path.join(os.getcwd(), "browser_data",
+                                         config.USER_DATA_DIR % self.command_args.platform)  # type: ignore
             browser_context = await chromium.launch_persistent_context(
                 user_data_dir=user_data_dir,
                 accept_downloads=True,
                 headless=headless,
-                proxy=playwright_proxy, # type: ignore
+                proxy=playwright_proxy,  # type: ignore
                 viewport={"width": 1920, "height": 1080},
                 user_agent=user_agent
-            ) # type: ignore
+            )  # type: ignore
             return browser_context
         else:
-            browser = await chromium.launch(headless=headless, proxy=playwright_proxy) # type: ignore
+            browser = await chromium.launch(headless=headless, proxy=playwright_proxy)  # type: ignore
             browser_context = await browser.new_context(
                 viewport={"width": 1920, "height": 1080},
                 user_agent=user_agent

main.py

@@ -83,12 +83,15 @@ async def post(self, uri: str, data: dict) -> Dict:
     async def ping(self) -> bool:
         """get a note to check if login state is ok"""
         utils.logger.info("begin to ping xhs...")
-        note_id = "5e5cb38a000000000100185e"
+        ping_flag = False
         try:
-            note_card: Dict = await self.get_note_by_id(note_id)
-            return note_card.get("note_id") == note_id
-        except Exception:
-            return False
+            note_card: Dict = await self.get_note_by_keyword(keyword="小红书")
+            if note_card.get("items"):
+                ping_flag = True
+        except Exception as e:
+            utils.logger.error(f"ping xhs failed: {e}")
+            ping_flag = False
+        return ping_flag
 
     async def update_cookies(self, browser_context: BrowserContext):
         cookie_str, cookie_dict = utils.convert_cookies(await browser_context.cookies())

media_platform/douyin/core.py

@@ -15,7 +15,7 @@
 from .exception import *
 from .login import XHSLogin
 from .client import XHSClient
-from models import xhs as xhs_model
+from models import xiaohongshu as xhs_model
 from base.base_crawler import AbstractCrawler
 from base.proxy_account_pool import AccountPool
 

media_platform/xhs/client.py

@@ -0,0 +1,3 @@
+from .douyin import *
+from .xiaohongshu import *
+