diff --git a/sat/memory/sqla.py b/sat/memory/sqla.py index 9fbb9cc1..9b975208 100644 --- a/sat/memory/sqla.py +++ b/sat/memory/sqla.py @@ -30,6 +30,7 @@ from sqlalchemy.orm import ( sessionmaker, subqueryload, joinedload, selectinload, contains_eager ) from sqlalchemy.orm.decl_api import DeclarativeMeta +from sqlalchemy.orm.attributes import Mapped from sqlalchemy.future import select from sqlalchemy.engine import Engine, Connection from sqlalchemy import update, delete, and_, or_, event, func @@ -215,6 +216,27 @@ class Storage: ## Generic + @aio + async def get( + self, + client: SatXMPPEntity, + db_cls: DeclarativeMeta, + db_id_col: Mapped, + id_value: Any, + joined_loads=None + ) -> Optional[DeclarativeMeta]: + stmt = select(db_cls).where(db_id_col==id_value) + if client is not None: + stmt = stmt.filter_by(profile_id=self.profiles[client.profile]) + if joined_loads is not None: + for joined_load in joined_loads: + stmt = stmt.options(joinedload(joined_load)) + async with self.session() as session: + result = await session.execute(stmt) + if joined_loads is not None: + result = result.unique() + return result.scalar_one_or_none() + @aio async def add(self, db_obj: DeclarativeMeta) -> None: """Add an object to database""" @@ -223,11 +245,29 @@ class Storage: session.add(db_obj) @aio - async def delete(self, db_obj: DeclarativeMeta) -> None: - """Delete an object from database""" + async def delete( + self, + db_obj: Union[DeclarativeMeta, List[DeclarativeMeta]], + session_add: Optional[List[DeclarativeMeta]] + ) -> None: + """Delete an object from database + + @param db_obj: object to delete or list of objects to delete + @param session_add: other objects to add to session. + This is useful when parents of deleted objects needs to be updated too, or if + other objects needs to be updated in the same transaction. + """ + if not db_obj: + return + if not isinstance(db_obj, list): + db_obj = [db_obj] async with self.session() as session: async with session.begin(): - await session.delete(db_obj) + if session_add is not None: + for obj in session_add: + session.add(obj) + for obj in db_obj: + await session.delete(obj) await session.commit() ## Profiles