rag_server/applications/utils/mysql/books.py

from .base import BaseMySQLClient


class Books(BaseMySQLClient):
    async def get_books(self):
        query = """
            SELECT book_id, book_name, book_oss_path, extract_status
            FROM books
            WHERE status = 1;
        """
        return await self.pool.async_fetch(query=query)

    async def get_book_extract_detail(self, book_id):
        query = """
            SELECT book_name, book_oss_path, extract_result FROM books WHERE book_id = %s;
        """
        return await self.pool.async_fetch(query=query, params=(book_id,))

    async def update_book_chunk_status(self, book_id, ori_status, new_status):
        query = """
            UPDATE books SET chunk_status = %s WHERE book_id = %s and chunk_status = %s;
        """
        return await self.pool.async_save(
            query=query, params=(new_status, book_id, ori_status)
        )