zhang_crawl_bio/worker/search_engine/drission_google_search.py

import time
import re
import logging
from pathlib import Path
from typing import Dict, Optional
from DrissionPage import ChromiumPage
from pydantic import BaseModel
from scrapling import Adaptor
from mylib.logu import logger
from mylib.base import save_to_file
from config.settings import GOOGLE_SEARCH_DIR
from mylib.drission_page import load_chrome_from_ini
from worker.search_engine.search_result_db import SearchResultManager, SearchResultItem, KeywordTask
from worker.search_engine.smart_selector import get_search_ele
from DrissionPage.common import Keys

class SearchResultEle(BaseModel):
    search_div: bool | None = None
    next_page_url: str | None = None
    current_page: int | None = None
    results: list[SearchResultItem] = []

class GoogleSearchHandlerDrission:
    def __init__(self, page: ChromiumPage):
        self.page = page
        self.db_manager = SearchResultManager()
        self.save_dir = GOOGLE_SEARCH_DIR
    
    def save_current_page(self, keyword: str, filename: str=time.strftime("%Y%m%d_%H%M%S")) -> Path:
        html_dir = self.save_dir / keyword
        html_dir.mkdir(parents=True, exist_ok=True)
        html_path = save_to_file(self.page.html, html_dir / f"{filename}.html")
        logger.info(f"save_to_file {html_path}")
        return html_path

    def _process_single_page(self, keyword: str) -> SearchResultEle:
        content = self.page.html
        result_ele = self.get_search_result_ele(content)
        
        if not result_ele.search_div:
            logger.warning(f"未找到搜索结果容器，可能遇到验证页面 keyword: {keyword}")
            return result_ele

        html_path = self.save_current_page(keyword, filename=f"{result_ele.current_page}")
        page_result = self.db_manager.save_page_results(
            keyword=keyword,
            page_number=result_ele.current_page,
            results_count=len(result_ele.results) if result_ele.results else 0,
            has_next_page=bool(result_ele.next_page_url),
            html_path=html_path
        )

        if result_ele.results:
            self.db_manager.save_result_items(
                keyword=keyword,
                page_id=page_result.id,
                items=result_ele.results,
                html_path=None, # 当前函数的 html_path 是 page 的 html_path，不是 results items 的 html_path，每个 items 有各自的 url 链接，对应各自的 html_path
            )
            
        return result_ele

    def check_keyword(self, keyword: str, skip_existing: bool) -> tuple[bool, KeywordTask]:
        key_model = self.db_manager.get_keyword_task(keyword)
        if skip_existing and key_model and key_model.is_completed:
            logger.info(f"关键词任务已完成，跳过处理: id={key_model.id} {keyword}")
            return True, key_model

        if key_model:
            self.db_manager.delete_keyword_task(keyword)
        key_model = self.db_manager.create_keyword_task(keyword)
        return False, key_model

    def process_keyword(self, keyword: str, max_result_items: int = 200, skip_existing: bool = False):
        _, key_model = self.check_keyword(keyword, skip_existing)
        self.search(keyword)
        search_result_item_count = 0
        should_complete = False
        
        while True:
            result_ele = self._process_single_page(keyword)
            if not result_ele.search_div:
                break

            search_result_item_count += len(result_ele.results) if result_ele.results else 0
            if search_result_item_count >= max_result_items or not result_ele.next_page_url:
                should_complete = True
                break

            try:
                self.page.scroll.to_bottom()
                time.sleep(3)
                next_btn = self.page.ele('#pnnext')
                if next_btn:
                    next_btn.click()
                    logger.info(f"跳转到下一页: {self.page.url}")
                    self.page.wait.load_start()
                else:
                    break
            except Exception as e:
                logger.warning(f"翻页失败: {str(e)}")
                break

        if should_complete:
            key_model = self.db_manager.mark_task_completed(keyword)
            logger.info(f"正常完成关键词处理: {keyword}")
        else:
            logger.warning(f"关键词处理被中断: {keyword}")
        return key_model

    def goto_home_page(self):
        url = "https://www.google.com"
        logger.info(f"{self.page.url}")
        if self.page.url != url:
            self.page.get(url,timeout=1)
        logger.info(f"{self.page.url}")
        if 'sorry/' in self.page.url:
            raise Exception(f"出现人机验证，正在换身份重试。。 {self.page.url}")

    def search(self, query: str):
        self.goto_home_page()
        search_ele_dict = get_search_ele(self.page.html)
        logger.info(f"smart xpath: {search_ele_dict}")
        if not search_ele_dict:
            raise Exception("未找到搜索框")
        search_box = self.page.ele(f'xpath:{search_ele_dict['xpath']}')
        search_box.input(query)
        self.page.actions.type(Keys.ENTER)
        self.page.wait.load_start()

    def get_current_page_num(self) -> int:
        if '/search?q=' in self.page.url:
            match = re.search(r'&start=(\d+)', self.page.url)
            return int(match.group(1)) // 10 + 1 if match else 1

    def get_search_result_ele(self, html_content: str):
        res = SearchResultEle(
            search_div=None,
            next_page_url=None,
            current_page=self.get_current_page_num(),
            results=[]
        )

        page = Adaptor(html_content)
        body = Adaptor(page.body)
        search_div = body.xpath('//div[@id="search"]')
        next_page_url = body.xpath_first('//a[@id="pnnext"]/@href')
        res.search_div = bool(search_div)
        res.next_page_url = f"https://www.google.com{next_page_url}" if next_page_url else None
        if not search_div:
            return res

        result_list = search_div.xpath('//*[@data-snc]')
        logger.info(f"当前页结果数量: {len(result_list)}, next_page_url: {next_page_url}")
        
        for result_item in result_list:
            if len(result_item.children) < 2:
                continue
                
            result = SearchResultItem()
            title_ele = result_item.children[0]
            if title_ele:
                result.url = title_ele.xpath_first('.//a/@href')
                result.title = title_ele.xpath_first('.//h3/text()')

            content_ele = result_item.children[1]
            if content_ele:
                content_list = content_ele.xpath('.//span/text()')
                result.content = ''.join(content_list) if content_list else None

            if any([result.url, result.title, result.content]):
                res.results.append(result)
        return res

def search_keyword_drission(keyword: str, max_result_items: int = 200, skip_existing: bool = True, browser_config: Optional[Dict] = None):
    ret = {'error': 0, 'msg': '', 'data': None}
    logger.info(f"keyword {keyword} max_result_items: {max_result_items} skip_existing: {skip_existing}")
    search_handler = GoogleSearchHandlerDrission(None)
    exist, keyword_model = search_handler.check_keyword(keyword, skip_existing)
    if exist and keyword_model.is_completed:
        ret['data'] = keyword_model.model_dump()
        return ret
    
    try:
        page = load_chrome_from_ini(**browser_config) if browser_config else load_chrome_from_ini()
        logger.info(f"proxy {page._browser._chromium_options.proxy}")
        search_handler = GoogleSearchHandlerDrission(page)
        kw = search_handler.process_keyword(keyword, max_result_items=max_result_items, skip_existing=skip_existing)
        if not kw:
            ret['error'] = 1
            html_path = search_handler.save_current_page(keyword, filename=f"warning_{time.strftime('%Y%m%d_%H%M%S')}")
            logger.warning(f"关键词任务未完成: {keyword} html_path: {html_path}")
            ret['msg'] = f"关键词任务未完成: {keyword}"
            ret['data'] = html_path
        else:
            ret['data'] = kw.model_dump()
        page.quit()
        return ret
    except Exception as e:
        html_path = search_handler.save_current_page(keyword, filename=f"error_{time.strftime('%Y%m%d_%H%M%S')}")
        logger.exception(f"失败: {str(e)} html_path: {html_path}")
        ret['error'] = 1
        ret['msg'] = f"失败: {str(e)}"
        ret['data'] = html_path
        return ret
    finally:
        page.quit()
def main():
    browser_config = {
        'proxy': 'http://127.0.0.1:9360',
        'headless': False
    }
    res = search_keyword_drission("Acalypha nusbaumeri essential oil", max_result_items=200, skip_existing = False,browser_config=browser_config)
    logger.info(f"{res}")

if __name__ == "__main__":
    main()