From 49e9b34b3f78c8db0f8fed3154a5188bee828138 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=D0=92=D0=B0=D0=B4=D0=B8=D0=BC=20=D0=A5=D1=80=D0=B8=D1=81?=
 =?UTF-8?q?=D1=82=D0=B5=D0=BD=D0=BA=D0=BE?= <vadim.khristenko@gmail.com>
Date: Sat, 5 Apr 2025 19:45:22 +0000
Subject: [PATCH] feat(storage): add get_value method to MongoStorage

Implement get_value method for MongoStorage that retrieves
individual values from storage data using MongoDB projections.

Add comprehensive test coverage for the new method.
---
 .gitignore                           |  1 +
 aiogram/fsm/storage/mongo.py         | 19 ++++++-
 tests/test_fsm/storage/test_mongo.py | 82 ++++++++++++++++++++++++++++
 3 files changed, 101 insertions(+), 1 deletion(-)

diff --git a/.gitignore b/.gitignore
index d1dbf6d6..1ce11ad0 100644
--- a/.gitignore
+++ b/.gitignore
@@ -22,3 +22,4 @@ reports
 
 dev/
 .venv/
+.conda/
diff --git a/aiogram/fsm/storage/mongo.py b/aiogram/fsm/storage/mongo.py
index b4b1eeaa..420eea5e 100644
--- a/aiogram/fsm/storage/mongo.py
+++ b/aiogram/fsm/storage/mongo.py
@@ -1,4 +1,4 @@
-from typing import Any, Dict, Optional, cast
+from typing import Any, Dict, Optional, cast, overload
 
 from motor.motor_asyncio import AsyncIOMotorClient
 
@@ -115,6 +115,23 @@ class MongoStorage(BaseStorage):
             return {}
         return cast(Dict[str, Any], document["data"])
 
+    @overload
+    async def get_value(self, storage_key: StorageKey, dict_key: str) -> Optional[Any]: ...
+
+    @overload
+    async def get_value(self, storage_key: StorageKey, dict_key: str, default: Any) -> Any: ...
+
+    async def get_value(
+        self, storage_key: StorageKey, dict_key: str, default: Optional[Any] = None
+    ) -> Optional[Any]:
+        document_id = self._key_builder.build(storage_key)
+        projection = {"_id": 0, f"data.{dict_key}": 1}
+        document = await self._collection.find_one({"_id": document_id}, projection=projection)
+        if not document or "data" not in document:
+            return default
+        nested_data = document.get("data", {})
+        return nested_data.get(dict_key, default)
+
     async def update_data(self, key: StorageKey, data: Dict[str, Any]) -> Dict[str, Any]:
         document_id = self._key_builder.build(key)
         update_with = {f"data.{key}": value for key, value in data.items()}
diff --git a/tests/test_fsm/storage/test_mongo.py b/tests/test_fsm/storage/test_mongo.py
index 95a476d6..23c8acb5 100644
--- a/tests/test_fsm/storage/test_mongo.py
+++ b/tests/test_fsm/storage/test_mongo.py
@@ -149,6 +149,88 @@ class TestStateAndDataDoNotAffectEachOther:
         }
 
 
+class TestGetValue:
+    async def test_get_existing_value(
+        self,
+        mongo_storage: MongoStorage,
+        storage_key: StorageKey,
+    ):
+        await mongo_storage.set_data(
+            storage_key, {"key": "value", "number": 42, "list": [1, 2, 3]}
+        )
+
+        assert await mongo_storage.get_value(storage_key, "key") == "value"
+        assert await mongo_storage.get_value(storage_key, "number") == 42
+        assert await mongo_storage.get_value(storage_key, "list") == [1, 2, 3]
+
+    async def test_get_non_existing_value(
+        self,
+        mongo_storage: MongoStorage,
+        storage_key: StorageKey,
+    ):
+        await mongo_storage.set_data(storage_key, {"key": "value"})
+
+        assert await mongo_storage.get_value(storage_key, "non_existing_key") is None
+        assert (
+            await mongo_storage.get_value(storage_key, "non_existing_key", default="default")
+            == "default"
+        )
+
+    async def test_get_value_from_non_existing_document(
+        self,
+        mongo_storage: MongoStorage,
+        storage_key: StorageKey,
+    ):
+        assert await mongo_storage._collection.find_one({}) is None
+
+        assert await mongo_storage.get_value(storage_key, "any_key") is None
+        assert (
+            await mongo_storage.get_value(storage_key, "any_key", default="default") == "default"
+        )
+
+    async def test_get_value_with_document_containing_only_state(
+        self,
+        mongo_storage: MongoStorage,
+        storage_key: StorageKey,
+    ):
+        await mongo_storage.set_state(storage_key, "test")
+
+        document = await mongo_storage._collection.find_one({})
+        assert document is not None
+        assert "data" not in document
+
+        assert await mongo_storage.get_value(storage_key, "any_key") is None
+        assert (
+            await mongo_storage.get_value(storage_key, "any_key", default="default") == "default"
+        )
+
+    async def test_get_value_uses_projection(
+        self,
+        mongo_storage: MongoStorage,
+        storage_key: StorageKey,
+        monkeypatch,
+    ):
+        await mongo_storage.set_data(
+            storage_key, {"key1": "value1", "key2": "value2", "key3": {"nested": "data"}}
+        )
+
+        original_find_one = mongo_storage._collection.find_one
+        calls = []
+
+        async def mock_find_one(*args, **kwargs):
+            calls.append(kwargs)
+            return await original_find_one(*args, **kwargs)
+
+        monkeypatch.setattr(mongo_storage._collection, "find_one", mock_find_one)
+
+        value = await mongo_storage.get_value(storage_key, "key2")
+
+        assert len(calls) == 1
+        assert "projection" in calls[0]
+        assert calls[0]["projection"] == {"_id": 0, "data.key2": 1}
+        assert value == "value2"
+
+
 @pytest.mark.parametrize(
     "value,result",
     [