zhang_crawl_bio/database/sql_model.py

from sqlite3 import connect
from mylib.settings import OUTPUT_DIR, WORK_DIR
from datetime import datetime
from sqlmodel import Field, SQLModel, create_engine, Session, select


sqlite_file_name = OUTPUT_DIR / "database.db"
sqlite_url = f"sqlite:///{sqlite_file_name}"
engine = create_engine(sqlite_url, echo=True)

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()
    created_at: datetime = Field(default_factory=datetime.now)

def create_db_and_tables():
    SQLModel.metadata.create_all(engine)

class DatabaseManager:
    def __init__(self):
        self.engine = engine
        
    def save_search_result(self, keyword: str, start: int, url: str, html_path: str) -> SearchResult:
        """保存搜索结果到数据库"""
        with Session(self.engine) as session:
            result = SearchResult(
                keyword=keyword,
                start=start,
                url=url,
                html_path=html_path
            )
            session.add(result)
            session.commit()
            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():
    create_db_and_tables()

if __name__ == "__main__":
    main()