diff --git a/ormar/models/mixins/save_mixin.py b/ormar/models/mixins/save_mixin.py index e9512db..46ef451 100644 --- a/ormar/models/mixins/save_mixin.py +++ b/ormar/models/mixins/save_mixin.py @@ -305,7 +305,7 @@ class SavePrepareMixin(RelationMixin, AliasMixin): 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, diff --git a/ormar/models/model.py b/ormar/models/model.py index d17cafd..509983c 100644 --- a/ormar/models/model.py +++ b/ormar/models/model.py @@ -37,6 +37,15 @@ class Model(ModelRow): :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) diff --git a/tests/test_model_methods/test_save_related_uuid.py b/tests/test_model_methods/test_save_related_uuid.py new file mode 100644 index 0000000..602d1bd --- /dev/null +++ b/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 diff --git a/tests/test_model_methods/test_upsert.py b/tests/test_model_methods/test_upsert.py new file mode 100644 index 0000000..178ae72 --- /dev/null +++ b/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"