improve in memory migration reader

datastore/migrations/core/migration_reader.py

@@ -1,5 +1,5 @@
 from collections import defaultdict
-from typing import Callable, Dict, Iterable, List, Optional, Protocol
+from typing import Any, Callable, Dict, Iterable, List, Optional, Protocol
 
 from datastore.reader.core import (
     AggregateRequest,
@@ -133,9 +133,7 @@ class MigrationReaderImplementationMemory(MigrationReader):
     is_in_memory_migration: bool = True
 
     def get(self, fqid: Fqid, mapped_fields: List[Field] = []) -> Model:
-        if fqid not in self.models:
-            raise ModelDoesNotExist(fqid)
-        return self.models[fqid]
+        return self._get_deep_copy_by_fqid(fqid, mapped_fields)
 
     def get_many(
         self, requests: List[GetManyRequestPart]
@@ -145,22 +143,24 @@ def get_many(
             for id in request.ids:
                 fqid = fqid_from_collection_and_id(request.collection, id)
                 if fqid in self.models:
-                    result[request.collection][id] = self.models[fqid]
+                    result[request.collection][id] = self._deep_copy_dict(
+                        self.models[fqid], request.mapped_fields
+                    )
         return result
 
     def get_all(
         self, collection: Collection, mapped_fields: List[Field] = []
     ) -> Dict[Id, Model]:
         return {
-            model["id"]: model
+            model["id"]: self._deep_copy_dict(model, mapped_fields)
             for fqid, model in self.models.items()
             if collection_from_fqid(fqid) == collection
         }
 
     def filter(
         self, collection: Collection, filter: Filter, mapped_fields: List[Field] = []
     ) -> Dict[Id, Model]:
-        return filter_models(self.models, collection, filter)
+        return filter_models(self.models, collection, filter, mapped_fields)
 
     def exists(self, collection: Collection, filter: Filter) -> bool:
         return self.count(collection, filter) > 0
@@ -204,3 +204,32 @@ def is_deleted(self, fqid: Fqid) -> bool:
 
     def model_exists(self, fqid: Fqid) -> bool:
         return fqid in self.models
+
+    def _get_deep_copy_by_fqid(
+        self, fqid: str, mapped_fields: List[Field]
+    ) -> dict[str, Any]:
+        """
+        Creates a deep copy of given model fqid recursively. Assumes no circular references between dict and subdicts.
+        Also filters all non mapped fields out.
+        """
+        if fqid not in self.models:
+            raise ModelDoesNotExist(fqid)
+        return self._deep_copy_dict(self.models[fqid], mapped_fields)
+
+    def _deep_copy_dict(
+        self, model: dict[str, Any], mapped_fields: List[Field] = []
+    ) -> dict[str, Any]:
+        """
+        Creates a deep copy of given dict recursively. Assumes no circular references between dict and subdicts.
+        Also filters all non mapped fields out.
+        """
+        new_model: dict[str, Any] = dict()
+        for key, value in model.items():
+            if not mapped_fields or key in mapped_fields:
+                if isinstance(value, dict):
+                    new_model[key] = self._deep_copy_dict(model[key])
+                elif isinstance(value, list):
+                    new_model[key] = value.copy()
+                else:
+                    new_model[key] = value
+        return new_model

tests/migrations/system/test_migration_reader.py

@@ -9,7 +9,7 @@
 from datastore.shared.util.filter import FilterOperator
 
 
-model = {"id": 1, "f": 1, "g": "test"}
+model = {"id": 1, "f": 1, "g": "test", "sub_dict": {"sub_list": []}}
 
 
 def check_migration_reader(migration_reader: MigrationReader):