Spaces:
Sleeping
Sleeping
File size: 2,628 Bytes
338036b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 | from .BaseDataModel import BaseDataModel
from .db_schemes.minirag.schemes import DataChunk
from sqlalchemy.future import select
from sqlalchemy import func, delete
from bson.objectid import ObjectId
class ChunkModel(BaseDataModel):
def __init__(self, db_client: object):
super().__init__(db_client=db_client)
self.db_client = db_client
@classmethod
async def create_instance(cls, db_client: object):
instance = cls(db_client)
return instance
async def create_chunk(self, chunk: DataChunk):
async with self.db_client() as session:
async with session.begin():
session.add(chunk)
await session.commit()
await session.refresh(chunk)
return chunk
async def get_chunk(self, chunk_id: str):
async with self.db_client() as session:
result = await session.execute(select(DataChunk.chunk_id).where(DataChunk.chunk_id==chunk_id))
chunk = result.scalar_one_or_none()
return chunk
async def insert_many_chunks(self, chunks: list, batch_size: int=100):
async with self.db_client() as session:
async with session.begin():
for i in range(0, len(chunks), batch_size):
chunk = chunks[i:i+batch_size]
session.add_all(chunk)
await session.commit()
return len(chunks)
async def delete_chunks_by_project_id(self, project_id: ObjectId):
async with self.db_client() as session:
stmt = (delete(DataChunk).where(DataChunk.chunk_project_id==project_id))
result = await session.execute(stmt)
await session.commit()
return result.rowcount
async def get_poject_chunks(self, project_id: ObjectId, page_no: int=1, page_size: int=20):
async with self.db_client() as session:
stmt = select(DataChunk).where(DataChunk.chunk_project_id==project_id).order_by(DataChunk.chunk_order).offset((page_no-1)*page_size).limit(page_size)
result = await session.execute(stmt)
records = result.scalars().all()
return records
async def get_chunk_count(self,project_id:ObjectId):
total_count = 0
async with self.db_client() as session:
count_sql = select(func.count(DataChunk.chunk_id)).where(DataChunk.chunk_project_id==project_id)
result = await session.execute(count_sql)
total_count = result.scalar()
return total_count
|