Adjusted the code according to dependency bump

Author: kamikaze

requirements.txt

@@ -3,6 +3,7 @@ asyncpg==0.28.0
 Authlib==1.2.1
 fastapi==0.101.0
 fastapi-pagination==0.12.7
+fastapi-users-db-sqlalchemy==6.0.1
 fastapi-users[sqlalchemy]==12.1.1
 numpy==1.25.2
 orjson==3.9.4

setup.cfg

@@ -59,6 +59,7 @@ web =
     Authlib==1.2.1
     fastapi==0.101.0
     fastapi-pagination==0.12.7
+    fastapi-users-db-sqlalchemy==6.0.1
     fastapi-users[sqlalchemy]==12.1.1
     numpy==1.25.2
     orjson==3.9.4

src/dataset_image_annotator/api/http.py

@@ -1,20 +1,29 @@
 import logging
+from contextlib import asynccontextmanager
 
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi_pagination import add_pagination
+from starlette.applications import Starlette
 
 from dataset_image_annotator.api.v1.endpoints import router, auth_router, users_router
 
+
+logger = logging.getLogger(__name__)
 origins = [
     'http://localhost',
     'http://localhost:3000',
     'http://localhost:8080',
     'http://localhost:5000',
 ]
 
-logger = logging.getLogger(__name__)
-app = FastAPI(docs_url='/api/docs', openapi_url='/api/v1/openapi.json')
+
+@asynccontextmanager
+async def lifespan(app: Starlette):
+    yield
+
+
+app = FastAPI(docs_url='/api/docs', openapi_url='/api/v1/openapi.json', lifespan=lifespan)
 app.add_middleware(
     CORSMiddleware,
     allow_origins=origins,

src/dataset_image_annotator/api/v1/endpoints.py

@@ -11,6 +11,7 @@
 from fastapi_users.authentication import AuthenticationBackend, JWTStrategy, CookieTransport
 from pydantic import Json
 from python3_commons.db import connect_to_db
+from sqlalchemy.ext.asyncio import AsyncSession
 
 from dataset_image_annotator import core
 from dataset_image_annotator.api import users
@@ -20,8 +21,8 @@
 )
 from dataset_image_annotator.conf import settings
 from dataset_image_annotator.core import upload_handler
-from dataset_image_annotator.db import database
 from dataset_image_annotator.db.models import User
+from dataset_image_annotator.db.user_db_helpers import get_async_session
 
 logger = logging.getLogger(__name__)
 router = APIRouter()
@@ -44,16 +45,6 @@ def get_jwt_strategy() -> JWTStrategy:
 users_router = fastapi_users.get_users_router(UserRead, UserUpdate, requires_verification=True)
 
 
-@router.on_event('startup')
-async def startup():
-    await connect_to_db(database, settings.db_dsn)
-
-
-@router.on_event('shutdown')
-async def shutdown():
-    await database.disconnect()
-
-
 def _handle_exceptions_helper(status_code, *args):
     if args:
         raise HTTPException(status_code=status_code, detail=args[0])
@@ -81,8 +72,9 @@ async def wrapper(*args, **kwargs):
 @router.get('/users', response_class=ORJSONResponse, tags=['Admin'])
 @handle_exceptions
 async def get_users(search: Json | None = None, order_by: str | None = None,
-                    user=Depends(get_current_superuser)) -> Page[UserItem]:
-    return await core.get_users(database, search, order_by)
+                    user=Depends(get_current_superuser),
+                    session: AsyncSession = Depends(get_async_session)) -> Page[UserItem]:
+    return await core.get_users(session, search, order_by)
 
 
 @router.post('/users', response_class=ORJSONResponse, tags=['Admin'])
@@ -94,18 +86,20 @@ async def create_user(new_user: UserCreate, user=Depends(get_current_superuser))
 @router.get('/user-groups', response_class=ORJSONResponse, tags=['Admin'])
 @handle_exceptions
 async def get_user_groups(search: Json | None = None, order_by: str | None = None,
-                          user=Depends(get_current_user)) -> Sequence[UserGroup]:
-    return await core.get_user_groups(database, search, order_by)
+                          user=Depends(get_current_user),
+                          session: AsyncSession = Depends(get_async_session)) -> Sequence[UserGroup]:
+    return await core.get_user_groups(session, search, order_by)
 
 
 @router.post('/raw-file', response_class=ORJSONResponse, tags=['Admin'])
 @handle_exceptions
-async def upload_raw_file(image_file: UploadFile = File(...), user=Depends(get_current_superuser)) -> bool:
+async def upload_raw_file(image_file: UploadFile = File(...), user=Depends(get_current_superuser),
+                          session: AsyncSession = Depends(get_async_session)) -> bool:
     if not image_file.filename:
         raise HTTPException(status_code=400, detail='Missing file')
 
     try:
-        response = await upload_handler.handle_raw_file(database, image_file)
+        response = await upload_handler.handle_raw_file(session, image_file)
     except TimeoutError as e:
         raise HTTPException(status_code=504, detail=str(e))
 
@@ -114,5 +108,6 @@ async def upload_raw_file(image_file: UploadFile = File(...), user=Depends(get_c
 
 @router.get('/image-samples', response_class=ORJSONResponse, tags=['Images'])
 @handle_exceptions
-async def get_image_samples(search: Json | None = None, order_by: str | None = None) -> Page[ImageSampleItem]:
-    return await core.get_image_samples(search, order_by)
+async def get_image_samples(session: AsyncSession = Depends(get_async_session), search: Json | None = None,
+                            order_by: str | None = None) -> Page[ImageSampleItem]:
+    return await core.get_image_samples(session, search, order_by)

src/dataset_image_annotator/core.py

@@ -4,9 +4,9 @@
 from typing import Mapping, Sequence
 
 import sqlalchemy as sa
-from databases.backends.postgres import Record
 from fastapi_pagination import Page
-from fastapi_pagination.ext.databases import paginate
+from fastapi_pagination.ext.sqlalchemy import paginate
+from sqlalchemy.ext.asyncio import AsyncSession
 
 from dataset_image_annotator.api.v1.schemas import UserItem, ImageSampleItem
 from dataset_image_annotator.db.helpers import get_query
@@ -17,28 +17,32 @@
 _ = t.gettext
 
 
-async def get_users(database, search: Mapping[str, str] | None = None, order_by: str | None = None) -> Page[UserItem]:
+async def get_users(session: AsyncSession, search: Mapping[str, str] | None = None,
+                    order_by: str | None = None) -> Page[UserItem]:
     query = sa.select([User])
-    result = await paginate(database, query)
+    result = await paginate(session, query)
 
     return result
 
 
-async def get_user(database, user_id: str) -> Record:
-    query = sa.select([User]).where(User.id == user_id)
+async def get_user(session: AsyncSession, user_id: str) -> UserItem:
+    query = sa.select(User).where(User.id == user_id)
+    cursor = await session.execute(query)
+    result = cursor.scalar_one()
 
-    return await database.fetch_row(query)
+    return result
 
 
-async def get_user_groups(database, search: Mapping[str, str] | None = None,
+async def get_user_groups(session: AsyncSession, search: Mapping[str, str] | None = None,
                           order_by: str | None = None) -> Sequence[UserGroup]:
-    query = sa.select([UserGroup]).order_by(UserGroup.name)
-    result = await database.fetch_all(query)
+    query = sa.select(UserGroup).order_by(UserGroup.name)
+    cursor = await session.execute(query)
+    result = cursor.scalars()
 
     return result
 
 
-async def get_image_samples(database, search: Mapping[str, str] | None = None,
+async def get_image_samples(session: AsyncSession, search: Mapping[str, str] | None = None,
                             order_by: str | None = None) -> Page[ImageSampleItem]:
     columns = {
         'id': (ImageSample.id, False, int, True),
@@ -47,11 +51,9 @@ async def get_image_samples(database, search: Mapping[str, str] | None = None,
         'location': (ImageSample.location, True, str, False),
     }
     where_clause, order_by_clause = get_query(search, order_by, columns)
-    query = sa.select([ImageSample.id, ImageSample.location]).order_by(order_by_clause)
+    query = sa.select(ImageSample.id, ImageSample.location).order_by(order_by_clause)
 
     if where_clause is not None:
         query = query.where(where_clause)
 
-    result = await database.fetch_all(query)
-
-    return result
+    return await paginate(session, query)

src/dataset_image_annotator/core/upload_handler.py

@@ -1,15 +1,15 @@
 import logging
 from zoneinfo import ZoneInfo
 
-from databases import Database
+from sqlalchemy.ext.asyncio import AsyncSession
 
 from dataset_image_annotator.conf import settings
 
 logger = logging.getLogger(__name__)
 timezone = ZoneInfo(settings.timezone)
 
 
-async def handle_raw_file(database: Database, image_file):
+async def handle_raw_file(session: AsyncSession, image_file):
     image_file_body = await image_file.read()
 
     return True

src/dataset_image_annotator/db/__init__.py

@@ -1,12 +1,8 @@
-import databases
 from sqlalchemy import MetaData
 from sqlalchemy.ext.declarative import declarative_base
 
-from dataset_image_annotator.conf import settings
-
 metadata = MetaData()
 Base = declarative_base(metadata=metadata)
-database = databases.Database(settings.db_dsn)
 
 
 async def is_healthy(pg) -> bool:

src/dataset_image_annotator/db/models.py

@@ -61,7 +61,6 @@ class ImageSampleAnnotation(BaseDBModel, Base):
     color = Column(String(32))
     votes = Column(Integer, nullable=False, default=0)
 
-
     __table_args__ = (
         UniqueConstraint('user_id', 'image_sample_id', name='uq_image_sample_annotation_item'),
     )

src/dataset_image_annotator/db/user_db_helpers.py

@@ -4,16 +4,16 @@
 from fastapi import Depends
 from fastapi_users_db_sqlalchemy import SQLAlchemyUserDatabase
 from sqlalchemy import MetaData
-from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession
-from sqlalchemy.orm import declarative_base, sessionmaker
+from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker
+from sqlalchemy.orm import declarative_base
 
 from dataset_image_annotator.conf import settings
 from dataset_image_annotator.db.models import User
 
 metadata = MetaData()
 Base = declarative_base(metadata=metadata)
 engine = create_async_engine(settings.db_dsn)
-async_session_maker = sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
+async_session_maker = async_sessionmaker(engine, expire_on_commit=False)
 
 
 async def get_async_session() -> AsyncGenerator[AsyncSession, None]:

src/dataset_image_annotator/jobs.py

@@ -7,7 +7,6 @@
 from dataset_image_annotator.api.users import get_user_manager_context
 from dataset_image_annotator.api.v1.schemas import UserCreate
 from dataset_image_annotator.conf import settings
-from dataset_image_annotator.db import database
 from dataset_image_annotator.db.user_db_helpers import get_async_session_context, get_user_db_context
 
 logging.config.dictConfig({
@@ -55,8 +54,6 @@
 
 
 async def create_superuser():
-    await database.connect()
-
     try:
         async with get_async_session_context() as session:
             async with get_user_db_context(session) as user_db:
@@ -73,8 +70,6 @@ async def create_superuser():
                     logger.info(f'User created: {settings.bootstrap_user_email}')
     except UserAlreadyExists:
         logger.warning(f'User already exists: {settings.bootstrap_user_email}')
-    finally:
-        await database.disconnect()
 
 
 def get_parsed_args():