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)
        )

    async def insert_book(self, book_id, book_name, book_oss_path):
        query = """
            INSERT INTO books (book_id, book_name, book_oss_path)
             VALUES (%s, %s, %s);
        """
        return await self.pool.async_save(
            query=query, params=(book_id, book_name, book_oss_path)
        )

    async def select_init_books(self):
        query = """
            SELECT book_id, book_name, book_oss_path, extract_status
            FROM books
            WHERE extract_status = 0;
        """
        return await self.pool.async_fetch(query=query)

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

    async def update_book_extract_status(self, book_id, status):
        query = """
            UPDATE books SET extract_status = %s WHERE book_id = %s;
            """
        return await self.pool.async_save(query=query, params=(status, book_id))

    async def update_book_extract_result(self, book_id, extract_result):
        query = """
            UPDATE books SET extract_result = %s, extract_status = 2 WHERE book_id = %s;
            """
        return await self.pool.async_save(query=query, params=(extract_result, book_id))