State controller (#114)

.github/workflows/github-actions-tests.yml

@@ -9,7 +9,7 @@ jobs:
         python-version: [ 3.6, 3.7, 3.8, 3.9 ]
         poetry-version: [ 1.1.4 ]
         mongodb-version: [ 4.4, 5.0 ]
-        pydantic-version: [1.5, 1.6, 1.7, 1.8]
+        pydantic-version: [1.7, 1.8]
         os: [ ubuntu-18.04 ]
     runs-on: ${{ matrix.os }}
     steps:

beanie/__init__.py

@@ -1,11 +1,17 @@
 from beanie.migrations.controllers.free_fall import free_fall_migration
 from beanie.migrations.controllers.iterative import iterative_migration
-from beanie.odm.actions import before_event, after_event, Insert, Replace
+from beanie.odm.actions import (
+    before_event,
+    after_event,
+    Insert,
+    Replace,
+    SaveChanges,
+)
 from beanie.odm.fields import PydanticObjectId, Indexed
 from beanie.odm.utils.general import init_beanie
 from beanie.odm.documents import Document
 
-__version__ = "1.3.0"
+__version__ = "1.4.0"
 __all__ = [
     # ODM
     "Document",
@@ -17,6 +23,7 @@
     "after_event",
     "Insert",
     "Replace",
+    "SaveChanges",
     # Migrations
     "iterative_migration",
     "free_fall_migration",

beanie/exceptions.py

@@ -24,3 +24,11 @@ class MigrationException(Exception):
 
 class ReplaceError(Exception):
     pass
+
+
+class StateManagementIsTurnedOff(Exception):
+    pass
+
+
+class StateNotSaved(Exception):
+    pass

beanie/odm/actions.py

@@ -16,10 +16,12 @@
 class EventTypes(str, Enum):
     INSERT = "INSERT"
     REPLACE = "REPLACE"
+    SAVE_CHANGES = "SAVE_CHANGES"
 
 
 Insert = EventTypes.INSERT
 Replace = EventTypes.REPLACE
+SaveChanges = EventTypes.SAVE_CHANGES
 
 
 class ActionDirections(str, Enum):  # TODO think about this name

beanie/odm/documents.py

@@ -13,7 +13,7 @@
 
 from bson import ObjectId
 from motor.motor_asyncio import AsyncIOMotorDatabase, AsyncIOMotorCollection
-from pydantic import ValidationError, parse_obj_as
+from pydantic import ValidationError, parse_obj_as, PrivateAttr
 from pydantic.main import BaseModel
 from pymongo.client_session import ClientSession
 from pymongo.results import (
@@ -44,6 +44,7 @@
 from beanie.odm.queries.update import UpdateMany
 from beanie.odm.utils.collection import collection_factory
 from beanie.odm.utils.dump import get_dict
+from beanie.odm.utils.state import saved_state_needed, save_state_after
 
 DocType = TypeVar("DocType", bound="Document")
 DocumentProjectionType = TypeVar("DocumentProjectionType", bound=BaseModel)
@@ -66,6 +67,8 @@ class Document(BaseModel, UpdateMethods):
 
     id: Optional[PydanticObjectId] = None
 
+    _saved_state: Optional[Dict[str, Any]] = PrivateAttr(default=None)
+
     def __init__(self, *args, **kwargs):
         super(Document, self).__init__(*args, **kwargs)
         self.get_motor_collection()
@@ -84,8 +87,11 @@ async def _sync(self) -> None:
             )
         for key, value in dict(new_instance).items():
             setattr(self, key, value)
+        if self.use_state_management():
+            self._save_state()
 
     @wrap_with_actions(EventTypes.INSERT)
+    @save_state_after
     async def insert(
         self: DocType, session: Optional[ClientSession] = None
     ) -> DocType:
@@ -417,6 +423,7 @@ def all(
         )
 
     @wrap_with_actions(EventTypes.REPLACE)
+    @save_state_after
     async def replace(
         self: DocType, session: Optional[ClientSession] = None
     ) -> DocType:
@@ -470,6 +477,7 @@ async def replace_many(
         await cls.find(In(cls.id, ids_list), session=session).delete()
         await cls.insert_many(documents, session=session)
 
+    @save_state_after
     async def update(
         self, *args, session: Optional[ClientSession] = None
     ) -> None:
@@ -648,6 +656,59 @@ async def inspect_collection(
                 )
         return inspection_result
 
+    # State management
+
+    @classmethod
+    def use_state_management(cls) -> bool:
+        collection_meta = cls._get_collection_meta()
+        return collection_meta.use_state_management
+
+    def _save_state(self):
+        if self.use_state_management():
+            self._saved_state = self.dict()
+
+    def get_saved_state(self):
+        return self._saved_state
+
+    @classmethod
+    def _parse_obj_saving_state(cls: Type[DocType], obj: Any) -> DocType:
+        result: DocType = cls.parse_obj(obj)
+        result._save_state()
+        return result
+
+    @property  # type: ignore
+    @saved_state_needed
+    def is_changed(self) -> bool:
+        if self._saved_state == self.dict():
+            return False
+        return True
+
+    @saved_state_needed
+    def get_changes(self) -> Dict[str, Any]:
+        #  TODO search deeply
+        changes = {}
+        if self.is_changed:
+            current_state = self.dict()
+            for k, v in self._saved_state.items():  # type: ignore
+                if v != current_state[k]:
+                    changes[k] = current_state[k]
+        return changes
+
+    @saved_state_needed
+    @wrap_with_actions(EventTypes.SAVE_CHANGES)
+    @save_state_after
+    async def save_changes(self) -> None:
+        if not self.is_changed:
+            return None
+        changes = self.get_changes()
+        await self.set(changes)  # type: ignore #TODO fix this
+
+    @saved_state_needed
+    def rollback(self) -> None:
+        if self.is_changed:
+            for key, value in self._saved_state.items():  # type: ignore
+                setattr(self, key, value)
+
     class Config:
         json_encoders = {
             ObjectId: lambda v: str(v),

beanie/odm/queries/cursor.py

@@ -11,6 +11,7 @@
 )
 
 from pydantic.main import BaseModel
+from beanie.odm.utils.parsing import parse_obj
 
 CursorResultType = TypeVar("CursorResultType")
 
@@ -40,11 +41,9 @@ async def __anext__(self) -> CursorResultType:
             self.cursor = self.motor_cursor
         next_item = await self.cursor.__anext__()
         projection = self.get_projection_model()
-        return (
-            projection.parse_obj(next_item)
-            if projection is not None
-            else next_item
-        )  # type: ignore
+        if projection is None:
+            return next_item
+        return parse_obj(projection, next_item)  # type: ignore
 
     async def to_list(
         self, length: Optional[int] = None
@@ -64,6 +63,6 @@ async def to_list(
         if projection is not None:
             return cast(
                 List[CursorResultType],
-                [projection.parse_obj(i) for i in motor_list],
+                [parse_obj(projection, i) for i in motor_list],
             )
         return cast(List[CursorResultType], motor_list)

beanie/odm/queries/find.py

@@ -40,6 +40,7 @@
     UpdateMany,
     UpdateOne,
 )
+from beanie.odm.utils.parsing import parse_obj
 from beanie.odm.utils.projection import get_projection
 
 if TYPE_CHECKING:
@@ -620,5 +621,5 @@ def __await__(
         if document is None:
             return None
         return cast(
-            FindQueryResultType, self.projection_model.parse_obj(document)
+            FindQueryResultType, parse_obj(self.projection_model, document)
         )

beanie/odm/utils/collection.py

@@ -20,6 +20,7 @@ def validate(cls, v):
 
 class CollectionInputParameters(BaseModel):
     name: str = ""
+    use_state_management: bool = False
     indexes: List[IndexModelField] = []
 
     class Config:
@@ -95,5 +96,6 @@ class CollectionMeta:
         name: str = collection_parameters.name
         motor_collection: AsyncIOMotorCollection = collection
         indexes: List = found_indexes
+        use_state_management: bool = collection_parameters.use_state_management
 
     return CollectionMeta

beanie/odm/utils/parsing.py

@@ -0,0 +1,14 @@
+from typing import Any, Type, Union, TYPE_CHECKING
+
+from pydantic import BaseModel
+
+if TYPE_CHECKING:
+    from beanie.odm.documents import Document
+
+
+def parse_obj(
+    model: Union[Type[BaseModel], Type["Document"]], data: Any
+) -> BaseModel:
+    if hasattr(model, "_parse_obj_saving_state"):
+        return model._parse_obj_saving_state(data)  # type: ignore
+    return model.parse_obj(data)

beanie/odm/utils/state.py

@@ -0,0 +1,44 @@
+import inspect
+from functools import wraps
+from typing import Callable, TYPE_CHECKING
+
+from beanie.exceptions import StateManagementIsTurnedOff, StateNotSaved
+
+if TYPE_CHECKING:
+    from beanie.odm.documents import DocType
+
+
+def check_if_state_saved(self: "DocType"):
+    if not self.use_state_management():
+        raise StateManagementIsTurnedOff(
+            "State management is turned off for this document"
+        )
+    if self._saved_state is None:
+        raise StateNotSaved("No state was saved")
+
+
+def saved_state_needed(f: Callable):
+    @wraps(f)
+    def sync_wrapper(self: "DocType", *args, **kwargs):
+        check_if_state_saved(self)
+        return f(self, *args, **kwargs)
+
+    @wraps(f)
+    async def async_wrapper(self: "DocType", *args, **kwargs):
+        check_if_state_saved(self)
+        return await f(self, *args, **kwargs)
+
+    if inspect.iscoroutinefunction(f):
+        return async_wrapper
+    return sync_wrapper
+
+
+def save_state_after(f: Callable):
+    @wraps(f)
+    async def wrapper(self: "DocType", *args, **kwargs):
+        result = await f(self, *args, **kwargs)
+        if self.use_state_management():
+            self._save_state()
+        return result
+
+    return wrapper

docs/changelog.md

@@ -2,6 +2,16 @@
 
 Beanie project changes
 
+## [1.4.0] - 2021-09-13
+
+### Added
+
+- Document state management
+
+### Implementation
+
+- PR - <https://github.com/roman-right/beanie/pull/114>
+
 ## [1.3.0] - 2021-09-08
 
 ### Added
@@ -439,4 +449,6 @@ Beanie project changes
 
 [1.2.8]: https://pypi.org/project/beanie/1.2.8
 
-[1.3.0]: https://pypi.org/project/beanie/1.3.0
+[1.3.0]: https://pypi.org/project/beanie/1.3.0
+
+[1.4.0]: https://pypi.org/project/beanie/1.4.0

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "beanie"
-version = "1.3.0"
+version = "1.4.0"
 description = "Asynchronous Python ODM for MongoDB"
 authors = ["Roman <[email protected]>"]
 license = "Apache-2.0"
@@ -14,7 +14,7 @@ readme = "README.md"
 
 [tool.poetry.dependencies]
 python = ">=3.6.1,<4.0"
-pydantic = "^1.5"
+pydantic = "^1.7"
 motor = "^2.5"
 click = "^7.1.2"
 toml = "^0.10.2"

tests/odm/conftest.py

@@ -17,6 +17,8 @@
     DocumentWithCustomIdUUID,
     DocumentWithCustomIdInt,
     DocumentWithActions,
+    DocumentWithTurnedOnStateManagement,
+    DocumentWithTurnedOffStateManagement,
 )
 from tests.odm.models import (
     Sample,
@@ -124,6 +126,8 @@ async def init(loop, db):
         DocumentWithCustomIdInt,
         Sample,
         DocumentWithActions,
+        DocumentWithTurnedOnStateManagement,
+        DocumentWithTurnedOffStateManagement,
     ]
     await init_beanie(
         database=db,

tests/odm/models.py

@@ -155,3 +155,16 @@ def num_2_change(self):
     @after_event(Replace)
     def num_3_change(self):
         self.num_3 -= 1
+
+
+class DocumentWithTurnedOnStateManagement(Document):
+    num_1: int
+    num_2: int
+
+    class Collection:
+        use_state_management = True
+
+
+class DocumentWithTurnedOffStateManagement(Document):
+    num_1: int
+    num_2: int