From 94ddc6b92621e82245c7202d25206b229ce28f15 Mon Sep 17 00:00:00 2001 From: xingc Date: Wed, 16 Apr 2025 14:40:19 +0800 Subject: [PATCH] =?UTF-8?q?feat(db):=20=E8=B0=83=E6=95=B4=E7=BB=9F?= =?UTF-8?q?=E4=B8=80=E7=AE=A1=E7=90=86=E4=BA=8B=E5=8A=A1=E6=96=B9=E5=BC=8F?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- core/db/crud.py | 160 +++++++++++++++++++++++----------------------- core/db/engine.py | 3 +- 2 files changed, 81 insertions(+), 82 deletions(-) diff --git a/core/db/crud.py b/core/db/crud.py index 003a219..5070be1 100644 --- a/core/db/crud.py +++ b/core/db/crud.py @@ -23,28 +23,27 @@ class CRUD: async def get(self, filters: dict, model: Model = None) -> Optional[Model]: model = model or self.model - async with self.session.begin(): - result = await self.session.execute(select(model).filter( - *self.build_filter_conditions(filters, model)) - ) - record = result.scalar_one_or_none() + + result = await self.session.execute(select(model).filter( + *self.build_filter_conditions(filters, model)) + ) + record = result.scalar_one_or_none() return record async def get_paginated( self, filters: dict, page: int = 1, page_size: int = 10, model: Model = None ) -> Tuple[List[Model], int]: model = model or self.model - async with self.session.begin(): - filter_conditions = self.build_filter_conditions(filters) if filters else [] - query = select(model).filter(*filter_conditions).order_by(desc(model.created_at)) + filter_conditions = self.build_filter_conditions(filters, model) if filters else [] + query = select(model).filter(*filter_conditions).order_by(desc(model.created_at)) - total_query = select(func.count()).select_from(model).filter(*filter_conditions) - total_result = await self.session.execute(total_query) - total = total_result.scalar() + total_query = select(func.count()).select_from(model).filter(*filter_conditions) + total_result = await self.session.execute(total_query) + total = total_result.scalar() - query = query.limit(page_size).offset((page - 1) * page_size) - result = await self.session.execute(query) - records = result.scalars().all() + query = query.limit(page_size).offset((page - 1) * page_size) + result = await self.session.execute(query) + records = result.scalars().all() return records, total @@ -52,103 +51,102 @@ class CRUD: self, group_by_field: str, filters: Dict[str, Any] = None, model: Model = None ) -> Dict[Any, int]: model = model or self.model - filter_conditions = self.build_filter_conditions(filters) if filters else [] + filter_conditions = self.build_filter_conditions(filters, model) if filters else [] group_by = getattr(model, group_by_field) query = select(group_by, func.count()).filter(*filter_conditions).group_by(group_by) - async with self.session.begin(): - result = await self.session.execute(query) - records = result.all() + result = await self.session.execute(query) + records = result.all() return {record[0]: record[1] for record in records} async def create(self, data: dict, filters: dict, model: Model = None) -> Union[Model, bool]: model = model or self.model - async with self.session.begin(): - result = await self.session.execute(select(model).filter( - *self.build_filter_conditions(filters, model)) - ) - if result.scalars().first() is not None: - return False - record = model(**data) - self.session.add(record) - await self.session.commit() + result = await self.session.execute(select(model).filter( + *self.build_filter_conditions(filters, model)) + ) + if result.scalars().first() is not None: + return False + + record = model(**data) + self.session.add(record) + await self.session.commit() return record async def get_all(self, filters: dict = None, model: Model = None) -> List[Model]: model = model or self.model - async with self.session.begin(): - result = await self.session.execute( - select(model) if filters is None else select(model).filter( - *self.build_filter_conditions(filters, model)).order_by(desc(model.insert_time)) - ) - records = result.scalars().all() + + result = await self.session.execute( + select(model) if filters is None else select(model).filter( + *self.build_filter_conditions(filters, model)).order_by(desc(model.insert_time)) + ) + records = result.scalars().all() return records async def update(self, data: dict, filters: dict, model: Model = None) -> Optional[Model]: model = model or self.model - async with self.session.begin(): - result = await self.session.execute(select(model).filter( - *self.build_filter_conditions(filters, model)) - ) - record = result.scalars().first() - if not record: - return None + result = await self.session.execute(select(model).filter( + *self.build_filter_conditions(filters, model)) + ) + record = result.scalars().first() - count = 0 - for key, value in data.items(): - if any([ - value is None, - value == getattr(record, key), - key in ['created_at', 'updated_at'] - ]): - continue - setattr(record, key, value) - count += 1 - if count > 0: - await self.session.commit() + if not record: + return None + + count = 0 + for key, value in data.items(): + if any([ + value is None, + value == getattr(record, key), + key in ['created_at', 'updated_at'] + ]): + continue + setattr(record, key, value) + count += 1 + if count > 0: + await self.session.commit() return record async def updates(self, data: dict, filters: dict, model: Model = None) -> Optional[Model]: model = model or self.model - async with self.session.begin(): - result = await self.session.execute(select(model).filter( - *self.build_filter_conditions(filters, model)) - ) - records = result.scalars().all() - if len(records) == 0: - return None - count = 0 - for record in records: - for key, value in data.items(): - if any([ - value is None, - value == getattr(record, key), - key in ['insert_time', 'update_time'] - ]): - continue - setattr(record, key, value) - count += 1 - if count > 0: - await self.session.commit() + result = await self.session.execute(select(model).filter( + *self.build_filter_conditions(filters, model)) + ) + records = result.scalars().all() + + if len(records) == 0: + return None + count = 0 + for record in records: + for key, value in data.items(): + if any([ + value is None, + value == getattr(record, key), + key in ['insert_time', 'update_time'] + ]): + continue + setattr(record, key, value) + count += 1 + if count > 0: + await self.session.commit() return records async def delete(self, filters: dict, model: Model = None) -> bool: model = model or self.model - async with self.session.begin(): - result = await self.session.execute(select(model).filter( - *self.build_filter_conditions(filters, model)) - ) - record = result.scalars().first() - if not record: - return False + result = await self.session.execute(select(model).filter( + *self.build_filter_conditions(filters, model)) + ) + record = result.scalars().first() - await self.session.delete(record) - await self.session.commit() + if not record: + return False + + await self.session.delete(record) + await self.session.commit() return True def build_filter_conditions(self, filters: Dict[str, Any], model: Model) -> tuple: diff --git a/core/db/engine.py b/core/db/engine.py index cdfab5a..c656cd4 100644 --- a/core/db/engine.py +++ b/core/db/engine.py @@ -35,7 +35,8 @@ session_maker = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit async def get_async_session() -> AsyncGenerator[AsyncSession, None]: async with session_maker() as session: - yield session + async with session.begin(): + yield session def create_redis_pool():