From 86170e1cb91cce4967118f16e38eb263197caef0 Mon Sep 17 00:00:00 2001
From: kievzenit <kir.backender@gmail.com>
Date: Mon, 11 Aug 2025 17:37:38 +0300
Subject: [PATCH] fixed PyMongoStorage update_data method implementation

---
 aiogram/fsm/storage/pymongo.py | 25 +++++++++++++++++--------
 1 file changed, 17 insertions(+), 8 deletions(-)

diff --git a/aiogram/fsm/storage/pymongo.py b/aiogram/fsm/storage/pymongo.py
index d547bb2f..5b891679 100644
--- a/aiogram/fsm/storage/pymongo.py
+++ b/aiogram/fsm/storage/pymongo.py
@@ -123,14 +123,23 @@ class PyMongoStorage(BaseStorage):
 
     async def update_data(self, key: StorageKey, data: Mapping[str, Any]) -> Dict[str, Any]:
         document_id = self._key_builder.build(key)
-        update_with = {f"data.{key}": value for key, value in data.items()}
-        update_result = await self._collection.find_one_and_update(
-            filter={"_id": document_id},
-            update={"$set": update_with},
-            upsert=True,
-            return_document=True,
-            projection={"_id": 0},
-        )
+        update_result = {}
+        if not data:
+            update_result = await self._collection.find_one_and_update(
+                filter={"_id": document_id},
+                update={"$unset": {"data": 1}},
+                upsert=True,
+                return_document=True,
+                projection={"_id": 0},
+            )
+        else:
+            update_result = await self._collection.find_one_and_update(
+                filter={"_id": document_id},
+                update={"$set": {"data": data}},
+                upsert=True,
+                return_document=True,
+                projection={"_id": 0},
+            )
         if not update_result:
             await self._collection.delete_one({"_id": document_id})
         return update_result.get("data", {})