zhang_crawl_bio/database/search_model.py

from sqlite3 import connect
from config.settings import OUTPUT_DIR, WORK_DIR
from datetime import datetime
from sqlmodel import Field, SQLModel, create_engine, Session, select
from sqlalchemy import text  # Add this import
from database.sqlite_engine import engine, drop_table

class SearchResult(SQLModel, table=True):
    id: int = Field(default=None, primary_key=True)
    keyword: str = Field()
    start: int = Field()
    url: str = Field()
    html_path: str = Field()
    is_last_page: bool = Field(default=False)  # 新增字段，标记是否是最后一页
    created_at: datetime = Field(default_factory=datetime.now)

class SearchDatabaseManager:
    def __init__(self):
        self.engine = engine
        
    def get_existing_result(self, keyword: str, start: int) -> SearchResult | None:
        """获取已存在的搜索结果"""
        with Session(self.engine) as session:
            return session.exec(
                select(SearchResult)
                .where(SearchResult.keyword == keyword)
                .where(SearchResult.start == start)
            ).first()
            
    def save_search_result(self, keyword: str, start: int, url: str, html_path: str, is_last_page: bool = False) -> SearchResult:
        """保存搜索结果到数据库"""
        # 检查是否已存在
        existing = self.get_existing_result(keyword, start)
        if existing:
            return existing
            
        with Session(self.engine) as session:
            result = SearchResult(
                keyword=keyword,
                start=start,
                url=url,
                html_path=html_path,
                is_last_page=is_last_page
            )
            session.add(result)
            session.commit()
            session.refresh(result)
            return result
            
    def get_search_results(self, keyword: str):
        """获取指定关键词的所有搜索结果"""
        with Session(self.engine) as session:
            results = session.exec(
                select(SearchResult)
                .where(SearchResult.keyword == keyword)
                .order_by(SearchResult.start)
            ).all()
            return results

def main():
    # 先删除SearchResult表
    drop_table(SearchResult)
    # print(SearchResult.__tablename__)
    # 再创建新表
    # create_db_and_tables()

if __name__ == "__main__":
    main()