Merge pull request #5 from aiogram/dev-2.x

Mongo storage cleanup (aiogram#609)
sakibguy · Jun 22, 2021 · e94a634 · e94a634
2 parents 5678df9 + e70a76f
commit e94a634
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/aiogram/contrib/fsm_storage/mongo.py b/aiogram/contrib/fsm_storage/mongo.py
@@ -142,9 +142,11 @@ async def set_data(self, *, chat: Union[str, int, None] = None, user: Union[str,
                        data: Dict = None):
         chat, user = self.check_address(chat=chat, user=user)
         db = await self.get_db()
-
-        await db[DATA].update_one(filter={'chat': chat, 'user': user},
-                                  update={'$set': {'data': data}}, upsert=True)
+        if not data:
+            await db[DATA].delete_one(filter={'chat': chat, 'user': user})
+        else:
+            await db[DATA].update_one(filter={'chat': chat, 'user': user},
+                                      update={'$set': {'data': data}}, upsert=True)
 
     async def get_data(self, *, chat: Union[str, int, None] = None, user: Union[str, int, None] = None,
                        default: Optional[dict] = None) -> Dict: