added async db io logic

audiostats/db/api.py

@@ -1,8 +1,8 @@
 import logging
-import time
 from collections.abc import Iterator
 
-from sqlalchemy.orm import sessionmaker, Session
+from sqlalchemy.ext.asyncio import async_sessionmaker, AsyncSession
+#from sqlalchemy.orm import sessionmaker, Session
 from .uof import UnitOfWork
 
 from audiostats.handlers import AlbumDTO
@@ -11,18 +11,15 @@
 logger = logging.getLogger(__name__)
 
 class DBApi:
-    def __init__(self, session_factory : sessionmaker[Session]):
+    def __init__(self, session_factory : async_sessionmaker[AsyncSession]):
         self._session_factory = session_factory
 
-    def upsert_albums(self, albums : Iterator[AlbumDTO]):
-        with UnitOfWork(self._session_factory) as uof:
-            total, success = 0, 0
+    async def upsert_albums(self, albums : Iterator[AlbumDTO]):
+        async with UnitOfWork(self._session_factory) as uow:
             for album in albums:
-                uof.albums.upsert(album)
-                total += 1
-                success += 1
+                await uow.albums.upsert(album)
 
-    def get_all_albums(self):
-        with UnitOfWork(self._session_factory) as uof:
-            return uof.albums.all()
+    async def get_all_albums(self):
+        async with UnitOfWork(self._session_factory) as uow:
+            return await uow.albums.all()
 

audiostats/db/repositories.py

@@ -1,5 +1,10 @@
+#from sqlalchemy import select
+from sqlalchemy import select, delete
 from sqlalchemy import select
-from sqlalchemy.orm import Session
+from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy.orm import selectinload
+
+#from sqlalchemy.orm import Session
 
 from .models import Album
 from audiostats.handlers import AlbumDTO
@@ -8,11 +13,11 @@
 
 
 class AlbumRepository:
-    def __init__(self, session : Session):
+    def __init__(self, session : AsyncSession):
         self._session = session
 
-    def upsert(self, album_data : AlbumDTO):
-        album = self.find_by_title_performer(album_data.title, album_data.performer)
+    async def upsert(self, album_data : AlbumDTO):
+        album = await self.find_by_title_performer(album_data.title, album_data.performer)
 
         if not album:
             album = Album()
@@ -30,11 +35,20 @@ def upsert(self, album_data : AlbumDTO):
                 album.tracks.append(track)
 
         for track in old_tracks_by_title.values():
-            #album.tracks.remove(track)
-            self._session.delete(track)
-
-    def find_by_title_performer(self, title : str, performer : str | None) -> Album | None:
-        return self._session.query(Album).filter(Album.title == title and Album.performer == performer if performer else Album.performer.is_(None)).first()
-
-    def all(self) -> list[Album]:
-        return [create_album_dto_f_orm(album) for album in self._session.scalars(select(Album)).all()]
+            await self._session.delete(track)
+
+    async def find_by_title_performer(self, title : str, performer : str | None) -> Album | None:
+        result = await self._session.execute(select(Album).where(
+            Album.title == title,
+            Album.performer == performer
+        ).options(
+        selectinload(Album.tracks)
+    ))
+        return result.scalar_one_or_none()
+        #return self._session.query(Album).filter(Album.title == title and Album.performer == performer if performer else Album.performer.is_(None)).first()
+
+    async def all(self) -> list[Album]:
+        result = await self._session.scalars(select(Album).options(
+        selectinload(Album.tracks)
+    ))
+        return [create_album_dto_f_orm(album) for album in result.all()]

audiostats/db/uof.py

@@ -1,28 +1,29 @@
-from sqlalchemy.orm import sessionmaker, Session
+from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
+#from sqlalchemy.orm import sessionmaker, Session
 import logging
 
 from .repositories import AlbumRepository
 
 logger = logging.getLogger(__name__)
 
 class UnitOfWork:
-    def __init__(self, session_factory: sessionmaker[Session]):
+    def __init__(self, session_factory: async_sessionmaker[AsyncSession]):
         self._session_factory = session_factory
-        self._session : Session | None = None
+        self._session : AsyncSession | None = None
         self.albums : AlbumRepository | None = None
         logger.debug(f'UoF initialized: {self}')
 
-    def __enter__(self):
+    async def __aenter__(self):
         self._session = self._session_factory()
         self.albums = AlbumRepository(self._session)
         logger.debug(f'AlbumRepo initialised: {self.albums}, Current Session: {self._session}')
         return self
 
-    def __exit__(self, exc_type, exc_val, exc_tb):
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
         if exc_type is not None:
-            self._session.rollback()
+            await self._session.rollback()
         else:
-            self._session.commit()
-        self._session.close()
+            await self._session.commit()
+        await self._session.close()
         logger.debug(f'Session closed')