import datetime from django.core.exceptions import FieldDoesNotExist from django.db.models import F from django.db.models.functions import Lower from django.db.utils import IntegrityError from django.test import TestCase, override_settings, skipUnlessDBFeature from .models import ( Article, CustomDbColumn, CustomPk, Detail, Food, Individual, JSONFieldNullable, Member, Note, Number, Order, Paragraph, RelatedObject, SingleObject, SpecialCategory, Tag, Valid, ) class WriteToOtherRouter: def db_for_write(self, model, **hints): return "other" class BulkUpdateNoteTests(TestCase): @classmethod def setUpTestData(cls): cls.notes = [Note.objects.create(note=str(i), misc=str(i)) for i in range(10)] def create_tags(self): self.tags = [Tag.objects.create(name=str(i)) for i in range(10)] def test_simple(self): for note in self.notes: note.note = "test-%s" % note.id with self.assertNumQueries(1): Note.objects.bulk_update(self.notes, ["note"]) self.assertCountEqual( Note.objects.values_list("note", flat=True), [cat.note for cat in self.notes], ) def test_multiple_fields(self): for note in self.notes: note.note = "test-%s" % note.id note.misc = "misc-%s" % note.id with self.assertNumQueries(1): Note.objects.bulk_update(self.notes, ["note", "misc"]) self.assertCountEqual( Note.objects.values_list("note", flat=True), [cat.note for cat in self.notes], ) self.assertCountEqual( Note.objects.values_list("misc", flat=True), [cat.misc for cat in self.notes], ) def test_batch_size(self): with self.assertNumQueries(len(self.notes)): Note.objects.bulk_update(self.notes, fields=["note"], batch_size=1) def test_unsaved_models(self): objs = self.notes + [Note(note="test", misc="test")] msg = "All bulk_update() objects must have a primary key set." with self.assertRaisesMessage(ValueError, msg): Note.objects.bulk_update(objs, fields=["note"]) def test_foreign_keys_do_not_lookup(self): self.create_tags() for note, tag in zip(self.notes, self.tags): note.tag = tag with self.assertNumQueries(1): Note.objects.bulk_update(self.notes, ["tag"]) self.assertSequenceEqual(Note.objects.filter(tag__isnull=False), self.notes) def test_set_field_to_null(self): self.create_tags() Note.objects.update(tag=self.tags[0]) for note in self.notes: note.tag = None Note.objects.bulk_update(self.notes, ["tag"]) self.assertCountEqual(Note.objects.filter(tag__isnull=True), self.notes) def test_set_mixed_fields_to_null(self): self.create_tags() midpoint = len(self.notes) // 2 top, bottom = self.notes[:midpoint], self.notes[midpoint:] for note in top: note.tag = None for note in bottom: note.tag = self.tags[0] Note.objects.bulk_update(self.notes, ["tag"]) self.assertCountEqual(Note.objects.filter(tag__isnull=True), top) self.assertCountEqual(Note.objects.filter(tag__isnull=False), bottom) def test_functions(self): Note.objects.update(note="TEST") for note in self.notes: note.note = Lower("note") Note.objects.bulk_update(self.notes, ["note"]) self.assertEqual(set(Note.objects.values_list("note", flat=True)), {"test"}) # Tests that use self.notes go here, otherwise put them in another class. class BulkUpdateTests(TestCase): databases = {"default", "other"} def test_no_fields(self): msg = "Field names must be given to bulk_update()." with self.assertRaisesMessage(ValueError, msg): Note.objects.bulk_update([], fields=[]) def test_invalid_batch_size(self): msg = "Batch size must be a positive integer." with self.assertRaisesMessage(ValueError, msg): Note.objects.bulk_update([], fields=["note"], batch_size=-1) with self.assertRaisesMessage(ValueError, msg): Note.objects.bulk_update([], fields=["note"], batch_size=0) def test_nonexistent_field(self): with self.assertRaisesMessage( FieldDoesNotExist, "Note has no field named 'nonexistent'" ): Note.objects.bulk_update([], ["nonexistent"]) pk_fields_error = "bulk_update() cannot be used with primary key fields." def test_update_primary_key(self): with self.assertRaisesMessage(ValueError, self.pk_fields_error): Note.objects.bulk_update([], ["id"]) def test_update_custom_primary_key(self): with self.assertRaisesMessage(ValueError, self.pk_fields_error): CustomPk.objects.bulk_update([], ["name"]) def test_empty_objects(self): with self.assertNumQueries(0): rows_updated = Note.objects.bulk_update([], ["note"]) self.assertEqual(rows_updated, 0) def test_large_batch(self): Note.objects.bulk_create( [Note(note=str(i), misc=str(i)) for i in range(0, 2000)] ) notes = list(Note.objects.all()) rows_updated = Note.objects.bulk_update(notes, ["note"]) self.assertEqual(rows_updated, 2000) def test_updated_rows_when_passing_duplicates(self): note = Note.objects.create(note="test-note", misc="test") rows_updated = Note.objects.bulk_update([note, note], ["note"]) self.assertEqual(rows_updated, 1) # Duplicates in different batches. rows_updated = Note.objects.bulk_update([note, note], ["note"], batch_size=1) self.assertEqual(rows_updated, 2) def test_only_concrete_fields_allowed(self): obj = Valid.objects.create(valid="test") detail = Detail.objects.create(data="test") paragraph = Paragraph.objects.create(text="test") Member.objects.create(name="test", details=detail) msg = "bulk_update() can only be used with concrete fields." with self.assertRaisesMessage(ValueError, msg): Detail.objects.bulk_update([detail], fields=["member"]) with self.assertRaisesMessage(ValueError, msg): Paragraph.objects.bulk_update([paragraph], fields=["page"]) with self.assertRaisesMessage(ValueError, msg): Valid.objects.bulk_update([obj], fields=["parent"]) def test_custom_db_columns(self): model = CustomDbColumn.objects.create(custom_column=1) model.custom_column = 2 CustomDbColumn.objects.bulk_update([model], fields=["custom_column"]) model.refresh_from_db() self.assertEqual(model.custom_column, 2) def test_custom_pk(self): custom_pks = [ CustomPk.objects.create(name="pk-%s" % i, extra="") for i in range(10) ] for model in custom_pks: model.extra = "extra-%s" % model.pk CustomPk.objects.bulk_update(custom_pks, ["extra"]) self.assertCountEqual( CustomPk.objects.values_list("extra", flat=True), [cat.extra for cat in custom_pks], ) def test_falsey_pk_value(self): order = Order.objects.create(pk=0, name="test") order.name = "updated" Order.objects.bulk_update([order], ["name"]) order.refresh_from_db() self.assertEqual(order.name, "updated") def test_inherited_fields(self): special_categories = [ SpecialCategory.objects.create(name=str(i), special_name=str(i)) for i in range(10) ] for category in special_categories: category.name = "test-%s" % category.id category.special_name = "special-test-%s" % category.special_name SpecialCategory.objects.bulk_update( special_categories, ["name", "special_name"] ) self.assertCountEqual( SpecialCategory.objects.values_list("name", flat=True), [cat.name for cat in special_categories], ) self.assertCountEqual( SpecialCategory.objects.values_list("special_name", flat=True), [cat.special_name for cat in special_categories], ) def test_field_references(self): numbers = [Number.objects.create(num=0) for _ in range(10)] for number in numbers: number.num = F("num") + 1 Number.objects.bulk_update(numbers, ["num"]) self.assertCountEqual(Number.objects.filter(num=1), numbers) def test_f_expression(self): notes = [ Note.objects.create(note="test_note", misc="test_misc") for _ in range(10) ] for note in notes: note.misc = F("note") Note.objects.bulk_update(notes, ["misc"]) self.assertCountEqual(Note.objects.filter(misc="test_note"), notes) def test_booleanfield(self): individuals = [Individual.objects.create(alive=False) for _ in range(10)] for individual in individuals: individual.alive = True Individual.objects.bulk_update(individuals, ["alive"]) self.assertCountEqual(Individual.objects.filter(alive=True), individuals) def test_ipaddressfield(self): for ip in ("2001::1", "1.2.3.4"): with self.subTest(ip=ip): models = [ CustomDbColumn.objects.create(ip_address="0.0.0.0") for _ in range(10) ] for model in models: model.ip_address = ip CustomDbColumn.objects.bulk_update(models, ["ip_address"]) self.assertCountEqual( CustomDbColumn.objects.filter(ip_address=ip), models ) def test_datetime_field(self): articles = [ Article.objects.create(name=str(i), created=datetime.datetime.today()) for i in range(10) ] point_in_time = datetime.datetime(1991, 10, 31) for article in articles: article.created = point_in_time Article.objects.bulk_update(articles, ["created"]) self.assertCountEqual(Article.objects.filter(created=point_in_time), articles) @skipUnlessDBFeature("supports_json_field") def test_json_field(self): JSONFieldNullable.objects.bulk_create( [JSONFieldNullable(json_field={"a": i}) for i in range(10)] ) objs = JSONFieldNullable.objects.all() for obj in objs: obj.json_field = {"c": obj.json_field["a"] + 1} JSONFieldNullable.objects.bulk_update(objs, ["json_field"]) self.assertCountEqual( JSONFieldNullable.objects.filter(json_field__has_key="c"), objs ) def test_nullable_fk_after_related_save(self): parent = RelatedObject.objects.create() child = SingleObject() parent.single = child parent.single.save() RelatedObject.objects.bulk_update([parent], fields=["single"]) self.assertEqual(parent.single_id, parent.single.pk) parent.refresh_from_db() self.assertEqual(parent.single, child) def test_unsaved_parent(self): parent = RelatedObject.objects.create() parent.single = SingleObject() msg = ( "bulk_update() prohibited to prevent data loss due to unsaved " "related object 'single'." ) with self.assertRaisesMessage(ValueError, msg): RelatedObject.objects.bulk_update([parent], fields=["single"]) def test_unspecified_unsaved_parent(self): parent = RelatedObject.objects.create() parent.single = SingleObject() parent.f = 42 RelatedObject.objects.bulk_update([parent], fields=["f"]) parent.refresh_from_db() self.assertEqual(parent.f, 42) self.assertIsNone(parent.single) @override_settings(DATABASE_ROUTERS=[WriteToOtherRouter()]) def test_database_routing(self): note = Note.objects.create(note="create") note.note = "bulk_update" with self.assertNumQueries(1, using="other"): Note.objects.bulk_update([note], fields=["note"]) @override_settings(DATABASE_ROUTERS=[WriteToOtherRouter()]) def test_database_routing_batch_atomicity(self): f1 = Food.objects.create(name="Banana") f2 = Food.objects.create(name="Apple") f1.name = "Kiwi" f2.name = "Kiwi" with self.assertRaises(IntegrityError): Food.objects.bulk_update([f1, f2], fields=["name"], batch_size=1) self.assertIs(Food.objects.filter(name="Kiwi").exists(), False)