Allow upserting rows with existing_pk even if they do not exist. (#889)

* Allow for __force_save__ in Model.upsert() method to save the models despite they already have a pk set. On integrity Error proceed to update the model, so in worst case two db calls will be made.

* Fix coverage

* Change implementation to checking if the row exists as postgres leaves hanging invalid transaction on integrity error. On force_save always check if row exists and then save/update (so always two queries).

Author: collerek

ormar/models/mixins/save_mixin.py

@@ -305,7 +305,7 @@ async def _upsert_model(
         if (
             save_all or not instance.pk or not instance.saved
         ) and not instance.__pk_only__:
-            await instance.upsert()
+            await instance.upsert(__force_save__=True)
             if relation_field and relation_field.is_multi:
                 await instance._upsert_through_model(
                     instance=instance,

ormar/models/model.py

@@ -37,6 +37,15 @@ async def upsert(self: T, **kwargs: Any) -> T:
         :return: saved Model
         :rtype: Model
         """
+
+        force_save = kwargs.pop("__force_save__", False)
+        if force_save:
+            expr = self.Meta.table.select().where(self.pk_column == self.pk)
+            row = await self.Meta.database.fetch_one(expr)
+            if not row:
+                return await self.save()
+            return await self.update(**kwargs)
+
         if not self.pk:
             return await self.save()
         return await self.update(**kwargs)

tests/test_model_methods/test_save_related_uuid.py

@@ -0,0 +1,83 @@
+import uuid
+from typing import Optional
+
+import databases
+import pytest
+import sqlalchemy
+
+import ormar
+from tests.settings import DATABASE_URL
+
+database = databases.Database(DATABASE_URL, force_rollback=True)
+metadata = sqlalchemy.MetaData()
+
+
+class Department(ormar.Model):
+    class Meta:
+        database = database
+        metadata = metadata
+
+    id: uuid.UUID = ormar.UUID(primary_key=True, default=uuid.uuid4)
+    department_name: str = ormar.String(max_length=100)
+
+
+class Course(ormar.Model):
+    class Meta:
+        database = database
+        metadata = metadata
+
+    id: uuid.UUID = ormar.UUID(primary_key=True, default=uuid.uuid4)
+    course_name: str = ormar.String(max_length=100)
+    completed: bool = ormar.Boolean()
+    department: Optional[Department] = ormar.ForeignKey(Department)
+
+
+class Student(ormar.Model):
+    class Meta:
+        database = database
+        metadata = metadata
+
+    id: uuid.UUID = ormar.UUID(primary_key=True, default=uuid.uuid4)
+    name: str = ormar.String(max_length=100)
+    courses = ormar.ManyToMany(Course)
+
+
+@pytest.fixture(autouse=True, scope="module")
+def create_test_database():
+    engine = sqlalchemy.create_engine(DATABASE_URL)
+    metadata.drop_all(engine)
+    metadata.create_all(engine)
+    yield
+    metadata.drop_all(engine)
+
+
+@pytest.mark.asyncio
+async def test_uuid_pk_in_save_related():
+    async with database:
+        to_save = {
+            "department_name": "Ormar",
+            "courses": [
+                {
+                    "course_name": "basic1",
+                    "completed": True,
+                    "students": [{"name": "Abi"}, {"name": "Jack"}],
+                },
+                {
+                    "course_name": "basic2",
+                    "completed": True,
+                    "students": [{"name": "Kate"}, {"name": "Miranda"}],
+                },
+            ],
+        }
+        department = Department(**to_save)
+        await department.save_related(follow=True, save_all=True)
+        department_check = (
+            await Department.objects.select_all(follow=True)
+            .order_by(Department.courses.students.name.asc())
+            .get()
+        )
+        to_exclude = {
+            "id": ...,
+            "courses": {"id": ..., "students": {"id", "studentcourse"}},
+        }
+        assert department_check.dict(exclude=to_exclude) == to_save

tests/test_model_methods/test_upsert.py

@@ -0,0 +1,64 @@
+from typing import Optional
+
+import databases
+import pytest
+import sqlalchemy
+
+import ormar
+from tests.settings import DATABASE_URL
+
+database = databases.Database(DATABASE_URL, force_rollback=True)
+metadata = sqlalchemy.MetaData()
+
+
+class Director(ormar.Model):
+    class Meta:
+        tablename = "directors"
+        metadata = metadata
+        database = database
+
+    id: int = ormar.Integer(primary_key=True)
+    name: str = ormar.String(max_length=100, nullable=False, name="first_name")
+    last_name: str = ormar.String(max_length=100, nullable=False, name="last_name")
+
+
+class Movie(ormar.Model):
+    class Meta:
+        tablename = "movies"
+        metadata = metadata
+        database = database
+
+    id: int = ormar.Integer(primary_key=True)
+    name: str = ormar.String(max_length=100, nullable=False, name="title")
+    year: int = ormar.Integer()
+    profit: float = ormar.Float()
+    director: Optional[Director] = ormar.ForeignKey(Director)
+
+
+@pytest.fixture(autouse=True, scope="module")
+def create_test_database():
+    engine = sqlalchemy.create_engine(DATABASE_URL)
+    metadata.drop_all(engine)
+    metadata.create_all(engine)
+    yield
+    metadata.drop_all(engine)
+
+
+@pytest.mark.asyncio
+async def test_updating_selected_columns():
+    async with database:
+        director1 = await Director(name="Peter", last_name="Jackson").save()
+
+        await Movie(
+            id=1, name="Lord of The Rings", year=2003, director=director1, profit=1.212
+        ).upsert()
+
+        with pytest.raises(ormar.NoMatch):
+            await Movie.objects.get()
+
+        await Movie(
+            id=1, name="Lord of The Rings", year=2003, director=director1, profit=1.212
+        ).upsert(__force_save__=True)
+        lotr = await Movie.objects.get()
+        assert lotr.year == 2003
+        assert lotr.name == "Lord of The Rings"