handled DB transaction globally

Author: dfm88

fastapi_2fa/api/deps/db.py

@@ -7,5 +7,11 @@ async def get_db() -> Generator:
     try:
         db = SessionLocal()
         yield db
+        await db.commit()
+        print('committed in db...')
+    except Exception as ex:
+        print(f'rolling db for exception {ex} ...')
+        await db.rollback()
     finally:
+        print('closing db...')
         await db.close()

fastapi_2fa/crud/base_crud.py

@@ -22,16 +22,31 @@ class CrudBase(
         UpdateSchemaType,  # Pydantic Update Schema
     ]
 ):
-    def __init__(self, model: Type[ModelType]):
+    def __init__(self, model: Type[ModelType], transaction: bool = False):
         """
         CRUD object with default methods to Create, Read, Update, Delete (CRUD).
 
-        **Parameters**
-
-        * `model`: A SQLAlchemy model class
-        * `schema`: A Pydantic model (schema) class
+        Args:
+            model (Type[ModelType]): A SQLAlchemy model class
+            transaction (bool, optional): if True, changes won't be committed.
+                Defaults to False.
         """
+
         self.model = model
+        self.transaction = transaction
+
+    def __call__(self, transaction: bool) -> 'CrudBase':
+        self.transaction = transaction
+        return self
+
+    async def handle_commit(self, db: Session) -> bool:
+        committed = False
+        if not self.transaction:
+            print(f'{self} is NOT under transaction, committing..')
+            await db.commit()
+            committed = True
+        print(f'{self} is under transaction, NOT committing..')
+        return committed
 
     async def get(self, db: Session, id: Any) -> Optional[ModelType]:
         query = select(self.model).where(self.model.id == id)
@@ -49,12 +64,8 @@ async def create(self, db: Session, *, obj_in: CreateSchemaType):
         obj_in_data = jsonable_encoder(obj_in)
         db_obj: ModelType = self.model(**obj_in_data)
         db.add(db_obj)
-        try:
-            await db.commit()
-        except Exception:
-            await db.rollback()
-            raise
-        await db.refresh(db_obj)
+        if await self.handle_commit(db):
+            await db.refresh(db_obj)
         return db_obj
 
     async def update(
@@ -76,21 +87,11 @@ async def update(
             .execution_options(synchronize_session="fetch")
         )
         await db.execute(query)
-        try:
-            # commit only if not under transaction
-            if not db.in_nested_transaction():
-                await db.commit()
-            return db_obj
-        except Exception as ex:
-            await db.rollback()
-            raise ex
+        if await self.handle_commit(db):
+            await db.refresh(db_obj)
+        return db_obj
 
     async def remove(self, db: Session, *, id: int) -> bool:
         query = delete(self.model).where(self.model.id == id)
         await db.execute(query)
-        try:
-            await db.commit()
-        except Exception as ex:
-            await db.rollback()
-            raise ex
-        return True
+        return self.handle_commit(db)

fastapi_2fa/crud/device.py

@@ -10,28 +10,26 @@
 
 class DeviceCrud(CrudBase[Device, DeviceCreate, DeviceUpdate]):
 
-    @staticmethod
-    async def create(db: Session, device: DeviceCreate, user: User) -> Device:
+    async def create(
+        self, db: Session, device: DeviceCreate, user: User
+    ) -> tuple[Device, SvgImage | None]:
         # create device
+        encoded_key: str = create_encoded_two_factor_auth_key()
         db_device = Device(
-            key=create_encoded_two_factor_auth_key(),
-            user_id=user.id,
+            key=encoded_key,
+            user=user,
             device_type=device.device_type
         )
         # create backup tokens
         for token in get_fake_otp_tokens():
             token_db = BackupToken(
-                device_id=db_device.id,
+                device=db_device,
                 token=token
             )
             db_device.backup_tokens.append(token_db)
         db.add(db_device)
-        try:
-            await db.commit()
-        except Exception:
-            await db.rollback()
-            raise
-        return db_device
+        if await self.handle_commit(db):
+            await db.refresh(db_device)
 
 
 device_crud = DeviceCrud(model=Device)

fastapi_2fa/crud/users.py

@@ -11,21 +11,17 @@
 
 
 class UserCrud(CrudBase[User, UserCreate, UserUpdate]):
-    @staticmethod
-    async def create(db: Session, user: UserCreate) -> User:
+
+    async def create(self, db: Session, user: UserCreate) -> User:
         db_obj = User(
             email=user.email,
             hashed_password=get_password_hash(user.password),
             tfa_enabled=user.tfa_enabled,
             full_name=user.full_name,
         )
         db.add(db_obj)
-        try:
-            await db.commit()
-        except Exception:
-            await db.rollback()
-            raise
-        await db.refresh(db_obj)
+        if await self.handle_commit(db):
+            await db.refresh(db_obj)
         return db_obj
 
     @staticmethod