Upsert (#76)

Upsert

Author: roman-right

beanie/__init__.py

@@ -4,7 +4,7 @@
 from beanie.odm.utils.general import init_beanie
 from beanie.odm.documents import Document
 
-__version__ = "1.1.6"
+__version__ = "1.2.0"
 __all__ = [
     # ODM
     "Document",

beanie/odm/documents.py

@@ -118,6 +118,10 @@ async def insert_one(
         :param session: ClientSession - pymongo session
         :return: InsertOneResult
         """
+        if not isinstance(document, cls):
+            raise TypeError(
+                "Inserting document must be of the original document class"
+            )
         return await cls.get_motor_collection().insert_one(
             get_dict(document), session=session
         )

beanie/odm/queries/find.py

@@ -12,6 +12,10 @@
     cast,
 )
 
+from pydantic import BaseModel
+from pymongo.client_session import ClientSession
+from pymongo.results import UpdateResult
+
 from beanie.exceptions import DocumentNotFound
 from beanie.odm.enums import SortDirection
 from beanie.odm.interfaces.aggregate import AggregateMethods
@@ -33,9 +37,6 @@
     UpdateOne,
 )
 from beanie.odm.utils.projection import get_projection
-from pydantic import BaseModel
-from pymongo.client_session import ClientSession
-from pymongo.results import UpdateResult
 
 if TYPE_CHECKING:
     from beanie.odm.documents import DocType
@@ -83,7 +84,7 @@ def update(
         :param session: Optional[ClientSession]
         :return: UpdateMany query
         """
-        self.set_session(session=session)
+        self.set_session(session)
         return (
             self.UpdateQueryType(
                 document_model=self.document_model,
@@ -93,6 +94,32 @@ def update(
             .set_session(session=self.session)
         )
 
+    def upsert(
+        self,
+        *args: Mapping[str, Any],
+        on_insert: "DocType",
+        session: Optional[ClientSession] = None
+    ):
+        """
+        Create Update with modifications query
+        and provide search criteria there
+
+        :param args: *Mapping[str,Any] - the modifications to apply.
+        :param on_insert: DocType - document to insert if there is no matched
+        document in the collection
+        :param session: Optional[ClientSession]
+        :return: UpdateMany query
+        """
+        self.set_session(session)
+        return (
+            self.UpdateQueryType(
+                document_model=self.document_model,
+                find_query=self.get_filter_query(),
+            )
+            .upsert(*args, on_insert=on_insert)
+            .set_session(session=self.session)
+        )
+
     def delete(
         self, session: Optional[ClientSession] = None
     ) -> Union[DeleteOne, DeleteMany]:

beanie/odm/queries/update.py

@@ -1,3 +1,4 @@
+from abc import abstractmethod
 from typing import (
     List,
     Type,
@@ -6,10 +7,11 @@
     Mapping,
     Any,
     Dict,
+    Union,
 )
 
 from pymongo.client_session import ClientSession
-from pymongo.results import UpdateResult
+from pymongo.results import UpdateResult, InsertOneResult
 
 from beanie.odm.interfaces.session import SessionMethods
 from beanie.odm.interfaces.update import (
@@ -40,6 +42,8 @@ def __init__(
         self.find_query = find_query
         self.update_expressions: List[Mapping[str, Any]] = []
         self.session = None
+        self.is_upsert = False
+        self.upsert_insert_doc: Optional["DocType"] = None
 
     @property
     def update_query(self) -> Dict[str, Any]:
@@ -57,7 +61,7 @@ def update(
         self, *args: Mapping[str, Any], session: Optional[ClientSession] = None
     ) -> "UpdateQuery":
         """
-        Provide modifications to the update query. The same as `update()`
+        Provide modifications to the update query.
 
         :param args: *Union[dict, Mapping] - the modifications to apply.
         :param session: Optional[ClientSession]
@@ -67,6 +71,48 @@ def update(
         self.update_expressions += args
         return self
 
+    def upsert(
+        self,
+        *args: Mapping[str, Any],
+        on_insert: "DocType",
+        session: Optional[ClientSession] = None
+    ) -> "UpdateQuery":
+        """
+        Provide modifications to the upsert query.
+
+        :param args: *Union[dict, Mapping] - the modifications to apply.
+        :param on_insert: DocType - document to insert if there is no matched
+        document in the collection
+        :param session: Optional[ClientSession]
+        :return: UpdateMany query
+        """
+        self.upsert_insert_doc = on_insert
+        self.update(*args, session=session)
+        return self
+
+    @abstractmethod
+    async def _update(self):
+        ...
+
+    def __await__(self) -> Union[UpdateResult, InsertOneResult]:
+        """
+        Run the query
+        :return:
+        """
+
+        update_result = yield from self._update().__await__()
+        if self.upsert_insert_doc is None:
+            return update_result
+        else:
+            if update_result.matched_count == 0:
+                return (
+                    yield from self.document_model.insert_one(
+                        document=self.upsert_insert_doc, session=self.session
+                    ).__await__()
+                )
+            else:
+                return update_result
+
 
 class UpdateMany(UpdateQuery):
     """
@@ -89,12 +135,8 @@ def update_many(
         """
         return self.update(*args, session=session)
 
-    def __await__(self) -> UpdateResult:
-        """
-        Run the query
-        :return:
-        """
-        yield from self.document_model.get_motor_collection().update_many(
+    async def _update(self):
+        return await self.document_model.get_motor_collection().update_many(
             self.find_query, self.update_query, session=self.session
         )
 
@@ -120,11 +162,7 @@ def update_one(
         """
         return self.update(*args, session=session)
 
-    def __await__(self) -> UpdateResult:
-        """
-        Run the query
-        :return:
-        """
-        yield from self.document_model.get_motor_collection().update_one(
+    async def _update(self):
+        return await self.document_model.get_motor_collection().update_one(
             self.find_query, self.update_query, session=self.session
         )

docs/changelog.md

@@ -2,6 +2,16 @@
 
 Beanie project changes
 
+## [1.2.0] - 2021-06-25
+
+### Added
+
+- Upsert
+
+### Implementation
+
+- Issue - <https://github.com/roman-right/beanie/issues/64>
+
 ## [1.1.6] - 2021-06-21
 
 ### Fix

docs/tutorial/update.md

@@ -44,6 +44,15 @@ Native MongoDB syntax is also accepted:
 await Product.find_one(Product.name == "Tony's").update({"$set": {Product.price: 3.33}})
 ```
 
+## Upsert
+
+To insert a document if no one document was matched the search cruteria during update query, the `upsert` method can be used:
+```python
+await Product.find_one(Product.name == "Tony's").upsert(
+    Set({Product.price: 3.33}), 
+    on_insert=Product(name="Tony's", price=3.33, category=chocolate)
+)
+```
 
 ## Deleting documents
 

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "beanie"
-version = "1.1.6"
+version = "1.2.0"
 description = "Asynchronous Python ODM for MongoDB"
 authors = ["Roman <roman-right@protonmail.com>"]
 license = "Apache-2.0"

tests/odm/conftest.py

@@ -74,7 +74,31 @@ async def preset_documents(point):
     await Sample.insert_many(documents=docs)
 
 
-object_storage = {}
+@pytest.fixture()
+def sample_doc_not_saved(point):
+    nested = Nested(
+        integer=0,
+        option_1=Option1(s="TEST"),
+        union=Option1(s="TEST"),
+        optional=None,
+    )
+    geo = GeoObject(
+        coordinates=[
+            point["longitude"],
+            point["latitude"],
+        ]
+    )
+    return Sample(
+        timestamp=datetime.utcnow(),
+        increment=0,
+        integer=0,
+        float_num=0,
+        string="TEST_NOT_SAVED",
+        nested=nested,
+        optional=None,
+        union=Option1(s="TEST"),
+        geo=geo,
+    )
 
 
 @pytest.fixture()

tests/odm/query/test_update.py

@@ -1,3 +1,5 @@
+import asyncio
+
 import pytest
 
 from beanie.odm.operators.update.general import Set, Max
@@ -136,3 +138,59 @@ async def test_update_many_with_session(preset_documents, session):
     assert len(result) == 3
     for sample in result:
         assert sample.increment == 100
+
+
+async def test_update_many_upsert_with_insert(
+    preset_documents, sample_doc_not_saved
+):
+    await Sample.find_many(Sample.integer > 100000).upsert(
+        Set({Sample.integer: 100}), on_insert=sample_doc_not_saved
+    )
+    await asyncio.sleep(2)
+    new_docs = await Sample.find_many(
+        Sample.string == sample_doc_not_saved.string
+    ).to_list()
+    assert len(new_docs) == 1
+    doc = new_docs[0]
+    assert doc.integer == sample_doc_not_saved.integer
+
+
+async def test_update_many_upsert_without_insert(
+    preset_documents, sample_doc_not_saved
+):
+    await Sample.find_many(Sample.integer > 1).upsert(
+        Set({Sample.integer: 100}), on_insert=sample_doc_not_saved
+    )
+    await asyncio.sleep(2)
+    new_docs = await Sample.find_many(
+        Sample.string == sample_doc_not_saved.string
+    ).to_list()
+    assert len(new_docs) == 0
+
+
+async def test_update_one_upsert_with_insert(
+    preset_documents, sample_doc_not_saved
+):
+    await Sample.find_one(Sample.integer > 100000).upsert(
+        Set({Sample.integer: 100}), on_insert=sample_doc_not_saved
+    )
+    await asyncio.sleep(2)
+    new_docs = await Sample.find_many(
+        Sample.string == sample_doc_not_saved.string
+    ).to_list()
+    assert len(new_docs) == 1
+    doc = new_docs[0]
+    assert doc.integer == sample_doc_not_saved.integer
+
+
+async def test_update_one_upsert_without_insert(
+    preset_documents, sample_doc_not_saved
+):
+    await Sample.find_one(Sample.integer > 1).upsert(
+        Set({Sample.integer: 100}), on_insert=sample_doc_not_saved
+    )
+    await asyncio.sleep(2)
+    new_docs = await Sample.find_many(
+        Sample.string == sample_doc_not_saved.string
+    ).to_list()
+    assert len(new_docs) == 0

tests/test_beanie.py

@@ -2,4 +2,4 @@
 
 
 def test_version():
-    assert __version__ == "1.1.6"
+    assert __version__ == "1.2.0"