From 246405288d699ee365a4badc9eda8f9cc930c199 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?St=C3=A9phan=20Mestach?= Date: Mon, 5 Oct 2026 21:33:14 +0200 Subject: [PATCH 1/5] IA-5485 Improve bulk zip upload code coverage --- .../tasks/test_process_mobile_bulk_upload.py | 682 +++++++++++++++++- 1 file changed, 676 insertions(+), 6 deletions(-) diff --git a/iaso/tests/tasks/test_process_mobile_bulk_upload.py b/iaso/tests/tasks/test_process_mobile_bulk_upload.py index 92faafdc44..6138bbe0fb 100644 --- a/iaso/tests/tasks/test_process_mobile_bulk_upload.py +++ b/iaso/tests/tasks/test_process_mobile_bulk_upload.py @@ -13,7 +13,9 @@ from django.contrib.contenttypes.models import ContentType from django.core.files import File from django.core.files.storage import default_storage +from django.core.files.uploadedfile import SimpleUploadedFile from django.test import SimpleTestCase, TestCase +from django.utils import timezone from beanstalk_worker.services import TestTaskService from hat.api_import.models import APIImport @@ -23,6 +25,7 @@ from iaso.models.common import ValidationWorkflowArtefactStatus from iaso.models.forms import CR_MODE_IF_REFERENCE_FORM from iaso.models.instances import instance_file_upload_to, instance_upload_to +from iaso.models.stocks import StockImpacts from iaso.tasks.process_mobile_bulk_upload import ( get_directory_handlers, process_instance_attachments, @@ -56,7 +59,11 @@ ] DEFAULT_CREATED_AT = datetime.datetime(2024, 4, 1, 0, 0, 5, tzinfo=pytz.UTC) +# When the scale zips' registrations were last updated on the device (their `updated_at`, 1712326150.005) +SCALE_REGISTRATION_UPDATED_AT = datetime.datetime(2024, 4, 5, 14, 9, 10, 5000, tzinfo=pytz.UTC) DEFAULT_CREATED_AT_STR = "2024-04-01" +# A submission XML, for the zips built by the tests +SUBMISSION_XML = "iaso/fixtures/instance_form_1_1.xml" DISASI_MAKULO_INSTANCE_FILE_NAME = ( "a5362052-408f-44f8-8abc-2a520c01ea10/16_12_127775b2-06a2-4ae6-b2bd-cf64143a9dfe_2024-04-05_16-09-42.xml" @@ -118,6 +125,40 @@ def create_entity_with_registration( return entity +def instance_state(instance_uuid): + """ + What the tests check of the instance `instance_uuid`, as a dict (None if there is none): compared in one + assertion, a failure shows every field that differs. + """ + instance = m.Instance.objects.filter(uuid=instance_uuid).first() + if instance is None: + return None + return { + "id": instance.id, + "entity": str(instance.entity.uuid) if instance.entity else None, + "form": instance.form_id, + "project": instance.project_id, + "org unit": instance.org_unit_id, + "created by": instance.created_by_id, + "file": instance.file.name, + "has data": instance.json is not None, + "updated at": instance.source_updated_at, + "validation status": instance.general_validation_status, + "deleted": instance.deleted, + } + + +def entity_state(entity_uuid): + """What the tests check of the entity `entity_uuid` (even soft-deleted), as a dict - see `instance_state()`.""" + entity = m.Entity.objects_include_deleted.get(uuid=entity_uuid) + return { + "deleted": entity.deleted_at is not None, + "merged to": str(entity.merged_to.uuid) if entity.merged_to else None, + "reference instance": entity.attributes.uuid if entity.attributes else None, + "instances": sorted(entity.instances.values_list("uuid", flat=True)), + } + + class ProcessMobileBulkUploadTest(TestCase): fixtures = ["user.yaml", "orgunit.yaml"] @@ -144,6 +185,8 @@ def setUp(self): self.default_entity_type = m.EntityType.objects.create( id=1, name="Participant", reference_form=self.form_registration ) + # From the orgunit.yaml fixture: an org unit without uuid, that submissions refer to by its id + self.org_unit = m.OrgUnit.objects.get(name="LaLaland") # Removing all InMemoryFileNodes inside the storage to avoid name conflicts - some can be kept by previous test classes default_storage._root._children.clear() # see InMemoryFileStorage in django/core/files/storage/memory.py @@ -167,7 +210,14 @@ def _create_zip_file(self): add_to_zip(zipf, zip_fixture_dir(CATT_TABLET_DIR), CORRECT_FILES_FOR_ZIP) save_file_to_api_import(self.api_import, zip_path) - def _create_zip_file_at_scale(self, num_patients=25): + def _create_zip_file_at_scale( + self, + num_patients=25, + registration_uuids=None, + with_followups=True, + api_import=None, + registrations_updated_at=SCALE_REGISTRATION_UPDATED_AT, + ): """ `num_patients` distinct *new* patients, each with exactly one registration + one CATT follow-up - the realistic shape of a large bulk sync from one facility: broad (many @@ -175,6 +225,12 @@ def _create_zip_file_at_scale(self, num_patients=25): rarer "lots of follow-ups for the same patient" case). Same org unit and same 2 forms/ versions throughout, reusing the base fixture's registration/CATT xml content as byte templates under fresh uuids. + + `registration_uuids`: the registrations' uuids (one patient each), instead of `num_patients` fresh ones, e.g. + to send registrations already uploaded by a previous zip again. `with_followups=False` leaves the CATT + follow-ups out. The zip is saved to `api_import` (`self.api_import` by default). `registrations_updated_at`: the + registrations' `updated_at`, e.g. later than when they were first sent, as when an entity workflow updated them. + Returns the registrations' uuids. """ base_dir = zip_fixture_dir(CATT_TABLET_DIR) with open( @@ -206,22 +262,22 @@ def _create_zip_file_at_scale(self, num_patients=25): zipf.writestr("orgUnits.json", json.dumps(org_units_data)) new_instances = [] - for _ in range(num_patients): + if registration_uuids is None: + registration_uuids = [str(uuid.uuid4()) for _ in range(num_patients)] + for registration_uuid in registration_uuids: # Matches the base fixture's convention: entityUuid == the registration # instance's own uuid. - registration_uuid = str(uuid.uuid4()) catt_uuid = str(uuid.uuid4()) reg_file_name = f"registration_{registration_uuid}.xml" catt_file_name = f"followup_{catt_uuid}.xml" zipf.writestr(f"{registration_uuid}/{reg_file_name}", registration_xml_bytes) - zipf.writestr(f"{catt_uuid}/{catt_file_name}", catt_xml_bytes) new_instances.append( { "id": registration_uuid, "created_at": 1.712326150005e9, - "updated_at": 1.712326150005e9, + "updated_at": registrations_updated_at.timestamp(), "file": f"/storage/emulated/0/Android/data/org.bluesquare/files/Documents/instances/{registration_uuid}/{reg_file_name}", "name": "Enregistrement", "formId": "1", @@ -234,6 +290,9 @@ def _create_zip_file_at_scale(self, num_patients=25): "accuracy": 14.929, } ) + if not with_followups: + continue + zipf.writestr(f"{catt_uuid}/{catt_file_name}", catt_xml_bytes) new_instances.append( { "id": catt_uuid, @@ -254,7 +313,8 @@ def _create_zip_file_at_scale(self, num_patients=25): zipf.writestr("instances.json", json.dumps(new_instances)) - save_file_to_api_import(self.api_import, zip_path) + save_file_to_api_import(api_import or self.api_import, zip_path) + return registration_uuids def test_success(self): self._create_zip_file() @@ -709,6 +769,265 @@ def test_form_version_query_count_at_scale(self): # 100 when the XML paths were over the field's `max_length`, each also checking its truncated name. profiler.assertStorageCounts({"_open": 1, "exists": 50, "_save": 50}) + def test_audit_trail_at_scale(self): + """ + A large batch (520 instances): one audit row per new instance, with the instance as created by `import_data()` + (no file nor json yet) as past value, and the fully processed instance as new value. + """ + m.FormVersion.objects.create(form=self.form_registration, version_id="2024032701") + m.FormVersion.objects.create(form=self.form_catt, version_id="2024031801") + self._create_zip_file_at_scale(num_patients=260) + + with QueryProfiler(trace_tables=["audit_modification"]) as profiler: + process_mobile_bulk_upload( + api_import_id=self.api_import.id, + project_id=self.project.id, + task=self.task, + _immediate=True, + ) + + # One INSERT per audit row + self.assertEqual(profiler.table_counts()["audit_modification"], 520) + + instances = m.Instance.objects.all() + self.assertEqual(instances.count(), 520) + content_type = ContentType.objects.get_for_model(m.Instance) + modifications = Modification.objects.filter(source=BULK_UPLOAD) + self.assertEqual(modifications.count(), 520) + self.assertEqual( + sorted(modifications.values_list("object_id", flat=True)), + sorted(str(instance_id) for instance_id in instances.values_list("id", flat=True)), + ) + self.assertEqual(set(modifications.values_list("content_type", flat=True)), {content_type.id}) + self.assertEqual(set(modifications.values_list("user", flat=True)), {self.user.id}) + + for instance in [instances.order_by("id").first(), instances.order_by("id").last()]: + modification = modifications.get(object_id=instance.id) + self.assertIsNotNone(modification.created_at) + past_fields = modification.past_value[0]["fields"] + new_fields = modification.new_value[0]["fields"] + self.assertEqual(modification.past_value[0]["pk"], instance.id) + self.assertEqual(modification.new_value[0]["pk"], instance.id) + self.assertEqual(past_fields["uuid"], instance.uuid) + self.assertEqual(past_fields["file"], "") + self.assertIsNone(past_fields["json"]) + self.assertEqual(new_fields["file"], instance.file.name) + self.assertEqual(new_fields["json"], instance.json) + self.assertEqual(new_fields["form_version"], instance.form_version_id) + + def test_campaign_registrations_then_followups_query_count(self): + """ + A campaign: a first zip registers 25 patients, then a second zip sends the same 25 registrations again, updated + a day later (as an entity workflow would), along with a new CATT follow-up for each - half of its instances + already exist. The existing registrations are updated in place. + """ + m.FormVersion.objects.create(form=self.form_registration, version_id="2024032701") + m.FormVersion.objects.create(form=self.form_catt, version_id="2024031801") + registration_uuids = self._create_zip_file_at_scale(num_patients=25, with_followups=False) + process_mobile_bulk_upload( + api_import_id=self.api_import.id, project_id=self.project.id, task=self.task, _immediate=True + ) + registrations = {instance.uuid: (instance.id, instance.file.name) for instance in m.Instance.objects.all()} + self.assertEqual(set(registrations), set(registration_uuids)) + + followups_import = APIImport.objects.create( + user=self.user, import_type="bulk", json_body={}, app_version="1.2.3" + ) + self._create_zip_file_at_scale( + registration_uuids=registration_uuids, + api_import=followups_import, + registrations_updated_at=SCALE_REGISTRATION_UPDATED_AT + datetime.timedelta(days=1), + ) + followups_task = m.Task.objects.create( + name="process_mobile_bulk_upload", launcher=self.user, account=self.account + ) + with QueryProfiler(trace_tables=["iaso_instance"]) as profiler: + process_mobile_bulk_upload( + api_import_id=followups_import.id, project_id=self.project.id, task=followups_task, _immediate=True + ) + + # The batch's pre-existing uuids lookup and `import_data()`'s file names lookup (2), then 6 per new follow-up + # (see test_form_version_query_count_baseline) and 5 per updated registration: looked up by uuid 3 times (twice + # by `import_data()`, once by `process_mobile_bulk_upload()`), then its 2 saves (with the new file, then with + # its json). + self.assertEqual(profiler.table_counts()["iaso_instance"], 2 + 25 * 6 + 25 * 5) + + self.assertEqual(m.Instance.objects.count(), 50) + self.assertEqual(m.Entity.objects.count(), 25) + for registration_uuid in registration_uuids: + entity = m.Entity.objects.get(uuid=registration_uuid) + registration_id, registration_file = registrations[registration_uuid] + self.assertEqual(entity.attributes_id, registration_id) + # Updated in place: same instance, with the new upload's file and timestamp + self.assertNotEqual(entity.attributes.file.name, registration_file) + self.assertEqual(entity.attributes.source_updated_at.date().isoformat(), "2024-04-06") + self.assertEqual( + sorted(entity.instances.values_list("form_id", flat=True)), + [self.form_registration.id, self.form_catt.id], + ) + + def test_failed_validation_workflow_start_is_not_saved(self): + """ + Re-sent (and updated) registrations whose validation workflow fails to start: `ValidationWorkflowEngine.start()` + sets them as PENDING before failing, and its savepoint rolls that back in the database: the rolled back status + must not be saved when they're updated later in the import. + """ + registration_uuids = self._create_zip_file_at_scale(num_patients=2, with_followups=False) + process_mobile_bulk_upload( + api_import_id=self.api_import.id, project_id=self.project.id, task=self.task, _immediate=True + ) + m.Instance.objects.update(general_validation_status=ValidationWorkflowArtefactStatus.REJECTED) + # 2 starting nodes: `get_starting_node()` fails, after the instance was set as PENDING + workflow = m.ValidationWorkflow.objects.create(name="Broken workflow", account=self.account) + m.ValidationNodeTemplate.objects.create(name="Start 1", workflow=workflow) + m.ValidationNodeTemplate.objects.create(name="Start 2", workflow=workflow) + self.form_registration.validation_workflow = workflow + self.form_registration.save() + + resend_import = APIImport.objects.create(user=self.user, import_type="bulk", json_body={}, app_version="1.2.3") + self._create_zip_file_at_scale( + registration_uuids=registration_uuids, + with_followups=False, + api_import=resend_import, + registrations_updated_at=SCALE_REGISTRATION_UPDATED_AT + datetime.timedelta(days=1), + ) + resend_task = m.Task.objects.create(name="process_mobile_bulk_upload", launcher=self.user, account=self.account) + process_mobile_bulk_upload( + api_import_id=resend_import.id, project_id=self.project.id, task=resend_task, _immediate=True + ) + + for registration in m.Instance.objects.filter(uuid__in=registration_uuids): + self.assertEqual(registration.source_updated_at.date().isoformat(), "2024-04-06") + self.assertEqual(registration.general_validation_status, ValidationWorkflowArtefactStatus.REJECTED) + + def _submission(self, instance_uuid, form, entity_uuid, updated_at=DEFAULT_CREATED_AT): + """One entry of a zip's instances.json, as the mobile app sends it. Its XML is `.xml`.""" + return { + "id": instance_uuid, + "created_at": int(DEFAULT_CREATED_AT.timestamp()), + "updated_at": int(updated_at.timestamp()), + "file": f"/storage/emulated/0/odk/instances/{instance_uuid}.xml", + "name": form.name, + "formId": str(form.id), + "orgUnitId": str(self.org_unit.id), + "entityUuid": str(entity_uuid), + "entityTypeId": str(self.default_entity_type.id), + } + + def _prepare_upload(self, submissions, without_xml=(), stock_ledger_items=None): + """ + A new bulk upload (`APIImport` and its `Task`) of a zip with `submissions`, each with its XML except the ones + whose uuid is in `without_xml`, and the `stock_ledger_items` given. + """ + with open(SUBMISSION_XML, "rb") as xml_file: + xml_content = xml_file.read() + zip_path = "/tmp/bulk_upload_submissions.zip" + with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zipf: + zipf.writestr("instances.json", json.dumps(submissions)) + if stock_ledger_items is not None: + zipf.writestr("stockLedgerItems.json", json.dumps(stock_ledger_items)) + for submission in submissions: + if submission["id"] not in without_xml: + zipf.writestr(f"{submission['id']}/{os.path.basename(submission['file'])}", xml_content) + api_import = APIImport.objects.create(user=self.user, import_type="bulk", json_body={}, app_version="1.2.3") + save_file_to_api_import(api_import, zip_path) + task = m.Task.objects.create(name="process_mobile_bulk_upload", launcher=self.user, account=self.account) + return api_import, task + + def _process(self, upload, project_id=None): + api_import, task = upload + process_mobile_bulk_upload( + api_import_id=api_import.id, project_id=project_id or self.project.id, task=task, _immediate=True + ) + + def _create_registration_zip(self, entity_uuid, instance_uuid): + """A zip with a single new registration (its own uuid) for the entity `entity_uuid`.""" + zip_path = "/tmp/bulk_upload_new_registration.zip" + file_name = f"registration_{instance_uuid}.xml" + instances_data = [ + { + "id": instance_uuid, + "created_at": int(DEFAULT_CREATED_AT.timestamp()), + "updated_at": int(DEFAULT_CREATED_AT.timestamp()), + "file": f"/storage/test/{file_name}", + "name": "Enregistrement", + "formId": str(self.form_registration.id), + "orgUnitId": "1", + "entityUuid": str(entity_uuid), + "entityTypeId": str(self.default_entity_type.id), + }, + ] + with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zipf: + zipf.writestr("instances.json", json.dumps(instances_data)) + with open("iaso/fixtures/instance_form_1_1.xml", "rb") as xml_file: + zipf.writestr(f"{instance_uuid}/{file_name}", xml_file.read()) + save_file_to_api_import(self.api_import, zip_path) + + def test_new_reference_instance_for_existing_entity(self): + """ + A new registration (its own uuid) for an entity that already has one: it becomes the entity's reference + instance (`attributes`), the previous one is kept. + """ + # Setup: a patient with a registration, then a zip with a new registration for them + patient = create_entity_with_registration(self, name="Patient", uuid=str(uuid.uuid4())) + previous_registration_uuid = patient.attributes.uuid + new_registration_uuid = str(uuid.uuid4()) + self._create_registration_zip(patient.uuid, new_registration_uuid) + + # Exercise + process_mobile_bulk_upload( + api_import_id=self.api_import.id, project_id=self.project.id, task=self.task, _immediate=True + ) + + # Assert + self.task.refresh_from_db() + self.assertEqual(self.task.status, m.SUCCESS) + self.assertEqual( + entity_state(patient.uuid), + { + "deleted": False, + "merged to": None, + "reference instance": new_registration_uuid, + "instances": sorted([previous_registration_uuid, new_registration_uuid]), + }, + ) + self.assertFalse(instance_state(previous_registration_uuid)["deleted"]) + + def test_new_reference_instance_for_merged_entity(self): + """ + A new registration (its own uuid) for an entity soft-deleted by a merge: it's attached to, and becomes the + reference instance (`attributes`) of, the entity it was merged into. The merged one is left as is. + """ + # Setup: patients A and B merged into C, then a zip with a new registration for A + ent_a = create_entity_with_registration(self, name="Patient A", uuid=str(uuid.uuid4())) + ent_b = create_entity_with_registration(self, name="Patient B", uuid=str(uuid.uuid4())) + ent_c = merge_entities(ent_a, ent_b, {}, self.user) + merged_a = entity_state(ent_a.uuid) + self.assertEqual((merged_a["deleted"], merged_a["merged to"]), (True, str(ent_c.uuid))) + ent_c_instances = entity_state(ent_c.uuid)["instances"] + new_registration_uuid = str(uuid.uuid4()) + self._create_registration_zip(ent_a.uuid, new_registration_uuid) + + # Exercise + process_mobile_bulk_upload( + api_import_id=self.api_import.id, project_id=self.project.id, task=self.task, _immediate=True + ) + + # Assert + self.task.refresh_from_db() + self.assertEqual(self.task.status, m.SUCCESS) + self.assertEqual( + entity_state(ent_c.uuid), + { + "deleted": False, + "merged to": None, + "reference instance": new_registration_uuid, + "instances": sorted(ent_c_instances + [new_registration_uuid]), + }, + ) + self.assertFalse(instance_state(new_registration_uuid)["deleted"]) + self.assertEqual(entity_state(ent_a.uuid), merged_a) + def test_org_unit_already_exists(self): self._create_zip_file() @@ -1726,6 +2045,357 @@ def test_bulk_upload_missing_file_key_error(self): self.assertIsNotNone(entity2.deleted_at) self.assertIsNone(entity2.attributes) + def test_submission_already_received_by_one_by_one_upload(self): + """ + A project uses either one-by-one uploads (POST /api/instances/ then POST /sync/form_upload/) or zips. While it + moves from the former to the latter, a device can have sent a submission's XML alone (its POST /api/instances/ + having failed), then send that submission again in a zip. The instance, found by its file name, is completed by + the zip instead of duplicated. + """ + # Setup: the XML alone, as an instance without uuid, entity nor data (no form to parse it with) + registration_uuid = str(uuid.uuid4()) + with open(SUBMISSION_XML, "rb") as xml_file: + xml_upload = SimpleUploadedFile(f"{registration_uuid}.xml", xml_file.read()) + self.client.post("/sync/form_upload/", {"xml_submission_file": xml_upload}) + received_alone = m.Instance.objects.get(file_name=f"{registration_uuid}.xml") + self.assertEqual((received_alone.uuid, received_alone.entity, received_alone.json), (None, None, None)) + upload = self._prepare_upload( + [self._submission(registration_uuid, self.form_registration, entity_uuid=registration_uuid)] + ) + + # Exercise + with QueryProfiler() as profiler: + self._process(upload) + + # Assert: the same instance, completed with the zip's metadata, and data parsed from its XML (stored again, in + # the account's folder). Its entity is created, with it as reference instance. + self.assertEqual(m.Instance.objects.count(), 1) + account_folder = f"{self.account.short_sanitized_name}_{self.account.id}" + self.assertEqual( + instance_state(registration_uuid), + { + "id": received_alone.id, + "entity": registration_uuid, + "form": self.form_registration.id, + "project": self.project.id, + "org unit": self.org_unit.id, + "created by": self.user.id, + "file": f"{account_folder}/instances/{timezone.now():%Y_%m}/{registration_uuid}/{registration_uuid}.xml", + "has data": True, + "updated at": DEFAULT_CREATED_AT, + "validation status": "", + "deleted": False, + }, + ) + self.assertEqual( + entity_state(registration_uuid), + { + "deleted": False, + "merged to": None, + "reference instance": registration_uuid, + "instances": [registration_uuid], + }, + ) + # Queries: the batch's uuids and file names lookups (2), its uuid looked up (1), found by its file name (1), + # saved by `import_data()` (1), then twice with its XML and data (2). S3: the zip, then the XML stored again. + self.assertEqual(profiler.table_counts()["iaso_instance"], 7) + profiler.assertStorageCounts({"_open": 1, "exists": 1, "_save": 1}) + + def test_approved_instance_sent_again(self): + """ + An instance already approved, sent again with newer data (e.g. edited on the device): this documents the current + behaviour - it's updated with the new data, and stays approved, without its validation workflow being restarted. + """ + # Setup: a registration, approved since, then sent again a day later + registration_uuid = str(uuid.uuid4()) + submission = self._submission(registration_uuid, self.form_registration, entity_uuid=registration_uuid) + self._process(self._prepare_upload([submission])) + m.Instance.objects.filter(uuid=registration_uuid).update( + general_validation_status=ValidationWorkflowArtefactStatus.APPROVED + ) + approved = instance_state(registration_uuid) + a_day_later = DEFAULT_CREATED_AT + datetime.timedelta(days=1) + upload = self._prepare_upload( + [self._submission(registration_uuid, self.form_registration, registration_uuid, updated_at=a_day_later)] + ) + + # Exercise + with QueryProfiler() as profiler: + self._process(upload) + + # Assert: updated in place with the new XML (stored next to the previous one), still approved + updated = instance_state(registration_uuid) + self.assertNotEqual(updated["file"], approved["file"]) + self.assertEqual(updated, {**approved, "file": updated["file"], "updated at": a_day_later}) + # Queries: the batch's uuids and file names lookups (2), its uuid looked up by `import_data()` (1) - which + # leaves it out, as approved - and by the task (1), then saved twice with its new XML and data (2). S3: the zip, + # then the new XML: its name is taken by the previous one, so a 2nd HEAD for the name it's stored under. + self.assertEqual(profiler.table_counts()["iaso_instance"], 6) + profiler.assertStorageCounts({"_open": 1, "exists": 2, "_save": 1}) + + def test_existing_instance_sent_again_without_its_xml(self): + """ + An instance already received, sent again in a zip that lacks its XML: it's kept as it was, with its entity. + Unlike a new instance whose XML is missing, which is deleted (see test_bulk_upload_missing_file_key_error). + """ + # Setup: a registration received with its XML, then sent again (updated) without it + registration_uuid = str(uuid.uuid4()) + self._process( + self._prepare_upload( + [self._submission(registration_uuid, self.form_registration, entity_uuid=registration_uuid)] + ) + ) + received = instance_state(registration_uuid) + patient = entity_state(registration_uuid) + a_day_later = DEFAULT_CREATED_AT + datetime.timedelta(days=1) + upload = self._prepare_upload( + [self._submission(registration_uuid, self.form_registration, registration_uuid, updated_at=a_day_later)], + without_xml=[registration_uuid], + ) + + # Exercise + with QueryProfiler() as profiler: + self._process(upload) + + # Assert: nothing changed + self.assertEqual(instance_state(registration_uuid), received) + self.assertEqual(entity_state(registration_uuid), patient) + # Queries: the batch's uuids and file names lookups (2), its uuid looked up twice by `import_data()` (2) and once + # by the task (1). S3: only the zip, nothing stored. + self.assertEqual(profiler.table_counts()["iaso_instance"], 5) + profiler.assertStorageCounts({"_open": 1}) + + def test_new_followup_without_its_xml(self): + """ + A new follow-up for a known patient, in a zip that lacks its XML: the follow-up is deleted, the patient's + entity and registration are left untouched (it isn't the entity's reference instance). + """ + # Setup: a patient registered, then a follow-up without its XML + registration_uuid = str(uuid.uuid4()) + self._process( + self._prepare_upload( + [self._submission(registration_uuid, self.form_registration, entity_uuid=registration_uuid)] + ) + ) + registration = instance_state(registration_uuid) + patient = entity_state(registration_uuid) + followup_uuid = str(uuid.uuid4()) + upload = self._prepare_upload( + [self._submission(followup_uuid, self.form_catt, entity_uuid=registration_uuid)], + without_xml=[followup_uuid], + ) + + # Exercise + with QueryProfiler() as profiler: + self._process(upload) + + # Assert: no follow-up, the rest unchanged + self.assertIsNone(instance_state(followup_uuid)) + self.assertEqual(instance_state(registration_uuid), registration) + self.assertEqual(entity_state(registration_uuid), patient) + # Queries: the batch's uuids and file names lookups (2), its uuid looked up (1), created by `import_data()` (2) + # and saved (1), then deleted (1). S3: only the zip, nothing stored. + self.assertEqual(profiler.table_counts()["iaso_instance"], 7) + profiler.assertStorageCounts({"_open": 1}) + + def _sku(self): + """A stock keeping unit of the project, as the mobile app gets them.""" + sku = m.StockKeepingUnit.objects.create( + name="Vaccine", short_name="VAC", account=self.user.iaso_profile.account + ) + sku.projects.add(self.project) + return sku + + def _stock_ledger_item(self, sku, submission_uuid, value, impact=StockImpacts.ADD): + """One entry of a zip's stockLedgerItems.json: a stock movement recorded by the submission `submission_uuid`.""" + return { + "id": str(uuid.uuid4()), + "sku": sku.id, + "org_unit": self.org_unit.id, + "submission_id": submission_uuid, + "question": "quantity", + "value": value, + "impact": impact, + } + + def _stock_state(self, sku): + """The stock of `sku` in `self.org_unit`: its ledger (sorted by value) and computed balance.""" + stock_item = m.StockItem.objects.filter(org_unit=self.org_unit, sku=sku).first() + return { + "ledger": sorted( + m.StockLedgerItem.objects.filter(org_unit=self.org_unit, sku=sku).values_list( + "submission__uuid", "impact", "value", "created_by" + ), + key=lambda item: item[2], + ), + "balance": stock_item.value if stock_item else None, + } + + def test_stock_ledger_items(self): + """ + Stock movements recorded by submissions of the same zip: imported once the submissions are, each referring to + its submission, and the org unit's stock balance computed. + """ + # Setup: a registration and its follow-up, recording 10 vaccines received then 3 used + sku = self._sku() + registration_uuid, followup_uuid = str(uuid.uuid4()), str(uuid.uuid4()) + upload = self._prepare_upload( + [ + self._submission(registration_uuid, self.form_registration, entity_uuid=registration_uuid), + self._submission(followup_uuid, self.form_catt, entity_uuid=registration_uuid), + ], + stock_ledger_items=[ + self._stock_ledger_item(sku, registration_uuid, 10), + self._stock_ledger_item(sku, followup_uuid, 3, impact=StockImpacts.SUBTRACT), + ], + ) + + # Exercise + with QueryProfiler() as profiler: + self._process(upload) + + # Assert + api_import, task = upload + task.refresh_from_db() + self.assertEqual(task.status, m.SUCCESS) + self.assertEqual( + self._stock_state(sku), + { + "ledger": [ + (followup_uuid, StockImpacts.SUBTRACT, 3, self.user.id), + (registration_uuid, StockImpacts.ADD, 10, self.user.id), + ], + "balance": 7, + }, + ) + # Queries, for each movement: is it already imported (1), saved - its `save()` checks again (2), then the + # balance computed from the ledger: last reset, then sum (2) + self.assertEqual(profiler.table_counts()["iaso_stockledgeritem"], 2 * 5) + + def test_stock_ledger_item_sent_again(self): + """A stock movement already imported, sent again in another zip: it's skipped, the balance doesn't change.""" + # Setup: 10 vaccines received, imported, then the same zip sent again + sku = self._sku() + registration_uuid = str(uuid.uuid4()) + submissions = [self._submission(registration_uuid, self.form_registration, entity_uuid=registration_uuid)] + stock_ledger_items = [self._stock_ledger_item(sku, registration_uuid, 10)] + self._process(self._prepare_upload(submissions, stock_ledger_items=stock_ledger_items)) + imported = self._stock_state(sku) + upload = self._prepare_upload(submissions, stock_ledger_items=stock_ledger_items) + + # Exercise + with QueryProfiler() as profiler: + self._process(upload) + + # Assert + api_import, task = upload + task.refresh_from_db() + self.assertEqual(task.status, m.SUCCESS) + self.assertEqual(self._stock_state(sku), imported) + self.assertEqual(imported["balance"], 10) + # Queries: is it already imported (1) - yes, skipped + self.assertEqual(profiler.table_counts()["iaso_stockledgeritem"], 1) + + def test_updated_followup_of_merged_entity(self): + """ + A follow-up of a patient since merged into another, updated on a device that still has the merged patient: the + follow-up is updated, the active patient's registration isn't (the follow-up isn't the reference instance). + """ + # Setup: patient A registered with a follow-up, merged with B into C, then A's follow-up updated a day later + registration_uuid, followup_uuid = str(uuid.uuid4()), str(uuid.uuid4()) + followup = self._submission(followup_uuid, self.form_catt, entity_uuid=registration_uuid) + self._process( + self._prepare_upload( + [self._submission(registration_uuid, self.form_registration, entity_uuid=registration_uuid), followup] + ) + ) + ent_b = create_entity_with_registration(self, name="Patient B", uuid=str(uuid.uuid4())) + ent_c = merge_entities(m.Entity.objects.get(uuid=registration_uuid), ent_b, {}, self.user) + active_registration_uuid = entity_state(ent_c.uuid)["reference instance"] + active_registration = instance_state(active_registration_uuid) + received_followup = instance_state(followup_uuid) + a_day_later = DEFAULT_CREATED_AT + datetime.timedelta(days=1) + upload = self._prepare_upload([{**followup, "updated_at": int(a_day_later.timestamp())}]) + + # Exercise + with QueryProfiler() as profiler: + self._process(upload) + + # Assert: the follow-up updated, C's registration unchanged + updated_followup = instance_state(followup_uuid) + self.assertEqual( + updated_followup, {**received_followup, "file": updated_followup["file"], "updated at": a_day_later} + ) + self.assertEqual(instance_state(active_registration_uuid), active_registration) + # Queries: the batch's uuids and file names lookups (2), its uuid looked up twice by `import_data()` (2) and once + # by the task (1), saved twice with its new XML and data (2), then the merged patient's reference instance + # loaded, to see it's not this follow-up (1). S3: the zip, then the new XML (its name is taken: a 2nd HEAD). + self.assertEqual(profiler.table_counts()["iaso_instance"], 8) + profiler.assertStorageCounts({"_open": 1, "exists": 2, "_save": 1}) + + def test_older_registration_update_of_merged_entity(self): + """ + A registration of a patient since merged into another, updated on a device that still has the merged patient, + but less recently than the active patient's registration: it's updated, the active patient's registration is + kept (the most recent data wins). + """ + # Setup: patient A merged with B into C, C's registration edited on the 5th, then A's registration updated on + # the 2nd + registration_uuid = str(uuid.uuid4()) + registration = self._submission(registration_uuid, self.form_registration, entity_uuid=registration_uuid) + self._process(self._prepare_upload([registration])) + ent_b = create_entity_with_registration(self, name="Patient B", uuid=str(uuid.uuid4())) + ent_c = merge_entities(m.Entity.objects.get(uuid=registration_uuid), ent_b, {}, self.user) + active_registration_uuid = entity_state(ent_c.uuid)["reference instance"] + m.Instance.objects.filter(uuid=active_registration_uuid).update( + source_updated_at=DEFAULT_CREATED_AT + datetime.timedelta(days=4) + ) + active_registration = instance_state(active_registration_uuid) + merged_registration = instance_state(registration_uuid) + a_day_later = DEFAULT_CREATED_AT + datetime.timedelta(days=1) + upload = self._prepare_upload([{**registration, "updated_at": int(a_day_later.timestamp())}]) + + # Exercise + with QueryProfiler() as profiler: + self._process(upload) + + # Assert: A's registration updated, C's kept + updated = instance_state(registration_uuid) + self.assertEqual(updated, {**merged_registration, "file": updated["file"], "updated at": a_day_later}) + self.assertEqual(instance_state(active_registration_uuid), active_registration) + # Queries: as in test_updated_followup_of_merged_entity (8), plus the active patient's registration loaded, to + # compare their dates (1). S3: the zip, then the new XML (its name is taken: a 2nd HEAD). + self.assertEqual(profiler.table_counts()["iaso_instance"], 9) + profiler.assertStorageCounts({"_open": 1, "exists": 2, "_save": 1}) + + def test_unknown_project(self): + """ + A bulk upload for a project that doesn't exist: the task fails and the import is flagged, nothing is imported. + """ + # Setup + registration_uuid = str(uuid.uuid4()) + api_import, task = self._prepare_upload( + [self._submission(registration_uuid, self.form_registration, entity_uuid=registration_uuid)] + ) + unknown_project_id = m.Project.objects.order_by("-id").first().id + 1 + + # Exercise + with QueryProfiler() as profiler: + self._process((api_import, task), project_id=unknown_project_id) + + # Assert + task.refresh_from_db() + api_import.refresh_from_db() + self.assertEqual((task.status, api_import.has_problem), (m.ERRORED, True)) + self.assertIn("Project matching query does not exist", api_import.exception) + self.assertEqual( + {"instances": m.Instance.objects.count(), "entities": m.Entity.objects.count()}, + {"instances": 0, "entities": 0}, + ) + # Stopped before reading the zip + self.assertEqual(profiler.table_counts()["iaso_instance"], 0) + profiler.assertStorageCounts({}) + def in_memory_zip(entries): """`entries`: {zip entry name: bytes}, written in the given order - a name ending with "/" is a directory entry.""" From 1a54ba0b5abbfd4e0b12b13118e83eafb9338320 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?St=C3=A9phan=20Mestach?= Date: Mon, 5 Oct 2026 21:44:20 +0200 Subject: [PATCH 2/5] reduce number of records to keep it fast --- .../tasks/test_process_mobile_bulk_upload.py | 45 +++++++++++-------- 1 file changed, 27 insertions(+), 18 deletions(-) diff --git a/iaso/tests/tasks/test_process_mobile_bulk_upload.py b/iaso/tests/tasks/test_process_mobile_bulk_upload.py index 6138bbe0fb..578e6410e0 100644 --- a/iaso/tests/tasks/test_process_mobile_bulk_upload.py +++ b/iaso/tests/tasks/test_process_mobile_bulk_upload.py @@ -771,13 +771,15 @@ def test_form_version_query_count_at_scale(self): def test_audit_trail_at_scale(self): """ - A large batch (520 instances): one audit row per new instance, with the instance as created by `import_data()` - (no file nor json yet) as past value, and the fully processed instance as new value. + A batch of 50 new instances: one audit row each, with the instance as created by `import_data()` (no file nor + json yet) as past value, and the fully processed instance as new value. """ + # Setup: 25 patients, each with a registration and a follow-up m.FormVersion.objects.create(form=self.form_registration, version_id="2024032701") m.FormVersion.objects.create(form=self.form_catt, version_id="2024031801") - self._create_zip_file_at_scale(num_patients=260) + self._create_zip_file_at_scale(num_patients=25) + # Exercise with QueryProfiler(trace_tables=["audit_modification"]) as profiler: process_mobile_bulk_upload( api_import_id=self.api_import.id, @@ -786,14 +788,12 @@ def test_audit_trail_at_scale(self): _immediate=True, ) - # One INSERT per audit row - self.assertEqual(profiler.table_counts()["audit_modification"], 520) - + # Assert instances = m.Instance.objects.all() - self.assertEqual(instances.count(), 520) + self.assertEqual(instances.count(), 50) content_type = ContentType.objects.get_for_model(m.Instance) modifications = Modification.objects.filter(source=BULK_UPLOAD) - self.assertEqual(modifications.count(), 520) + self.assertEqual(modifications.count(), 50) self.assertEqual( sorted(modifications.values_list("object_id", flat=True)), sorted(str(instance_id) for instance_id in instances.values_list("id", flat=True)), @@ -804,16 +804,25 @@ def test_audit_trail_at_scale(self): for instance in [instances.order_by("id").first(), instances.order_by("id").last()]: modification = modifications.get(object_id=instance.id) self.assertIsNotNone(modification.created_at) - past_fields = modification.past_value[0]["fields"] - new_fields = modification.new_value[0]["fields"] - self.assertEqual(modification.past_value[0]["pk"], instance.id) - self.assertEqual(modification.new_value[0]["pk"], instance.id) - self.assertEqual(past_fields["uuid"], instance.uuid) - self.assertEqual(past_fields["file"], "") - self.assertIsNone(past_fields["json"]) - self.assertEqual(new_fields["file"], instance.file.name) - self.assertEqual(new_fields["json"], instance.json) - self.assertEqual(new_fields["form_version"], instance.form_version_id) + past, new = modification.past_value[0], modification.new_value[0] + self.assertEqual( + { + "past": {"pk": past["pk"], **{k: past["fields"][k] for k in ["uuid", "file", "json"]}}, + "new": {"pk": new["pk"], **{k: new["fields"][k] for k in ["file", "json", "form_version"]}}, + }, + { + # As created by `import_data()`: no file nor data yet + "past": {"pk": instance.id, "uuid": instance.uuid, "file": "", "json": None}, + "new": { + "pk": instance.id, + "file": instance.file.name, + "json": instance.json, + "form_version": instance.form_version_id, + }, + }, + ) + # One INSERT per audit row + self.assertEqual(profiler.table_counts()["audit_modification"], 50) def test_campaign_registrations_then_followups_query_count(self): """ From 39b0e63cc183e73988b7863283816fb09863ec91 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?St=C3=A9phan=20Mestach?= Date: Mon, 5 Oct 2026 21:50:44 +0200 Subject: [PATCH 3/5] IA-5486 Bulk zip upload: cache instances and org units per batch, bulk insert audit rows --- hat/audit/models.py | 7 +- iaso/api/instances/import_cache.py | 42 +++++++ iaso/api/instances/views.py | 22 +++- iaso/tasks/process_mobile_bulk_upload.py | 36 +++--- iaso/tests/api/instances/test_import_cache.py | 51 +++++++++ .../tasks/test_process_mobile_bulk_upload.py | 104 +++++++++--------- 6 files changed, 186 insertions(+), 76 deletions(-) create mode 100644 iaso/api/instances/import_cache.py create mode 100644 iaso/tests/api/instances/test_import_cache.py diff --git a/hat/audit/models.py b/hat/audit/models.py index 56f6436d73..f138bda088 100644 --- a/hat/audit/models.py +++ b/hat/audit/models.py @@ -186,7 +186,11 @@ def log_modification( source: Optional[str], user: User = None, org_unit_change_request_id: int = None, + save: bool = True, ) -> Modification: + """ + `save=False` returns the modification without saving it, e.g. to `bulk_create` many of them at once. + """ modification = Modification() modification.past_value = [] modification.new_value = [] @@ -231,5 +235,6 @@ def log_modification( if not any([added, removed]) and len(modified.keys()) == 1 and "updated_at" in modified: logger.warning("log_modification() called with only `updated_at`.", extra={"modification": modification}) - modification.save() + if save: + modification.save() return modification diff --git a/iaso/api/instances/import_cache.py b/iaso/api/instances/import_cache.py new file mode 100644 index 0000000000..4d3541ed8b --- /dev/null +++ b/iaso/api/instances/import_cache.py @@ -0,0 +1,42 @@ +from collections import defaultdict +from typing import Dict, Iterable, List, Optional, Tuple + +from iaso.models import Instance, OrgUnit + + +class InstanceImportCache: + """ + Lookups shared by all the instances of an import batch (e.g. a mobile bulk upload zip): done once for the whole + batch instead of once per instance. + + - Instances: the ones already in the database for the batch's uuids, in a single query, plus the ones created + while importing it (see `add_instance()`), so that a later entry of the batch with the same uuid finds them. + - Org units given by uuid: looked up the first time they're seen (a batch typically covers a few facilities). + """ + + def __init__(self, instance_uuids: Iterable[Optional[str]]): + # Lists: older instances can share a uuid (the unique index only covers recent ones). Ordered by id, so the + # first one is the one `.first()` would pick. + self._instances_by_uuid: Dict[str, List[Instance]] = defaultdict(list) + for instance in Instance.objects.filter(uuid__in=[uuid for uuid in instance_uuids if uuid]).order_by("id"): + self._instances_by_uuid[instance.uuid].append(instance) + self._org_units: Dict[Tuple[str, Optional[int]], OrgUnit] = {} + + def instances(self, uuid: str) -> List[Instance]: + """All the instances with this uuid, oldest first.""" + return self._instances_by_uuid.get(uuid, []) + + def instance(self, uuid: str) -> Optional[Instance]: + """The instance with this uuid, or None if there is none, or several.""" + instances = self.instances(uuid) + return instances[0] if len(instances) == 1 else None + + def add_instance(self, instance: Instance) -> None: + """Record an instance created (or given its uuid) while importing the batch.""" + self._instances_by_uuid[instance.uuid].append(instance) + + def org_unit(self, uuid: str, version_id: Optional[int]) -> OrgUnit: + key = (uuid, version_id) + if key not in self._org_units: + self._org_units[key] = OrgUnit.objects.get(uuid=uuid, version_id=version_id) + return self._org_units[key] diff --git a/iaso/api/instances/views.py b/iaso/api/instances/views.py index b8c0a6f218..b70540c794 100644 --- a/iaso/api/instances/views.py +++ b/iaso/api/instances/views.py @@ -39,6 +39,7 @@ safe_api_import, ) from iaso.api.instances.filters import get_form_from_instance_filters, parse_instance_filters +from iaso.api.instances.import_cache import InstanceImportCache from iaso.api.instances.json import JsonbPathQueryFirst, JsonPathField, RegexpReplace from iaso.api.instances.permissions import PERMISSION_CLASSES_RW, HasInstanceBulkPermission, HasInstancePermission from iaso.api.instances.serializers import ( @@ -1121,11 +1122,13 @@ def find_entity(account: Account, entity_uuid: str, entity_type_id: Optional[int return sorted(existing_entities, key=_entity_correctness_score, reverse=True)[0] -def import_data(instances, user, app_id, api_import=None): +def import_data(instances, user, app_id, api_import=None, cache: Optional[InstanceImportCache] = None): """ This function creates empty instances (without files) and should be called first when uploading new instances. Sometimes, due to some network issues, this function might not properly be called and the instances are created by the second endpoint (POST /sync/form_upload/). + + `cache`: the batch's `InstanceImportCache`, when the caller shares it - the instances created here are added to it. """ project = Project.objects.get_for_user_and_app_id(user, app_id) rtn_instances = [] @@ -1145,10 +1148,14 @@ def import_data(instances, user, app_id, api_import=None): if existing_uuid: uuids_by_file_name[existing_file_name].add(existing_uuid) + if cache is None: + cache = InstanceImportCache([instance_data.get("id") for instance_data in instances]) + for instance_data in instances: uuid = instance_data.get("id", None) - existing_instances = Instance.objects.filter(uuid=uuid) + # A missing or empty uuid keeps its own lookup, as before + existing_instances = cache.instances(uuid) if uuid else Instance.objects.filter(uuid=uuid).order_by("id") if existing_instances: if all( existing_instance.general_validation_status @@ -1156,7 +1163,7 @@ def import_data(instances, user, app_id, api_import=None): + Instance._meta.get_field("general_validation_status").empty_values for existing_instance in existing_instances ): - rtn_instances.append(existing_instances.first()) + rtn_instances.append(existing_instances[0]) continue # Get or create instance based on file_name - this "get or create" logic is important: @@ -1173,10 +1180,11 @@ def import_data(instances, user, app_id, api_import=None): file_name = f"{base}_dup_{uuid}{ext}" instance, _ = Instance.objects.get_or_create(file_name=file_name) + instance.uuid = uuid if uuid: uuids_by_file_name[file_name].add(uuid) + cache.add_instance(instance) - instance.uuid = uuid instance.project = project instance.name = instance_data.get("name", None) instance.period = instance_data.get("period", None) @@ -1193,8 +1201,7 @@ def import_data(instances, user, app_id, api_import=None): if str(tentative_org_unit_id).isdigit(): instance.org_unit_id = tentative_org_unit_id else: - org_unit = OrgUnit.objects.get(uuid=tentative_org_unit_id, version_id=project.account.default_version_id) - instance.org_unit = org_unit + instance.org_unit = cache.org_unit(tentative_org_unit_id, project.account.default_version_id) raw_form_id = instance_data.get("formId") # Normalize to int: the mobile app sends this as a JSON string (e.g. "1"), which would make @@ -1285,6 +1292,9 @@ def import_data(instances, user, app_id, api_import=None): except Exception as e: # so we avoid the whole instance creation crashing logger.error(e) + # `start()` may have set the status before failing: rolled back in the database, but not on this + # in-memory instance, which callers keep using and saving (see `InstanceImportCache`). + instance.refresh_from_db(fields=["general_validation_status"]) return rtn_instances diff --git a/iaso/tasks/process_mobile_bulk_upload.py b/iaso/tasks/process_mobile_bulk_upload.py index 0b89046bdc..aaebf9494a 100644 --- a/iaso/tasks/process_mobile_bulk_upload.py +++ b/iaso/tasks/process_mobile_bulk_upload.py @@ -29,8 +29,9 @@ from beanstalk_worker.throttle import Throttle from hat.api.export_utils import timestamp_to_utc_datetime from hat.api_import.models import APIImport -from hat.audit.models import BULK_UPLOAD, log_modification +from hat.audit.models import BULK_UPLOAD, Modification, log_modification from hat.sync.views import create_instance_file, process_instance_file +from iaso.api.instances.import_cache import InstanceImportCache from iaso.api.instances.views import import_data as import_instances from iaso.api.org_unit_change_requests.serializers import OrgUnitChangeRequestWriteSerializer from iaso.api.org_units import import_org_units @@ -44,6 +45,8 @@ STORAGE_LOGS_JSON = "storageLogs.json" CHANGE_REQUESTS_JSON = "changeRequests.json" STOCK_LEDGER_ITEMS_JSON = "stockLedgerItems.json" +# Audit rows are inserted in chunks: their past/new values hold each instance's whole json +MODIFICATIONS_BATCH_SIZE = 500 logger = logging.getLogger(__name__) @@ -88,30 +91,22 @@ def process_mobile_bulk_upload(api_import_id, project_id, task=None): if INSTANCES_JSON in zip_ref.namelist(): log_progress(the_task, 20, "Processing forms and files") instances_data = read_json_file_from_zip(zip_ref, INSTANCES_JSON) - existing_uuids = set( - Instance.objects.filter(uuid__in=[d["id"] for d in instances_data]).values_list( - "uuid", flat=True - ) - ) - imported_instances = import_instances(instances_data, user, project.app_id, api_import=api_import) - # `import_instances` already built and saved the new instances in memory - reuse them - # instead of re-querying by uuid below. Instances that already existed still go - # through `.get()`: `import_data` returns (or skips) them without checking for - # duplicate uuids, and `.get()` fails the import on duplicates instead of silently - # picking one. - instances_by_uuid = { - instance.uuid: instance - for instance in imported_instances - if instance.uuid not in existing_uuids - } + # Shared with `import_instances`, which adds the instances it creates: below, both new and + # pre-existing instances are reused from it instead of being re-queried by uuid. Except for + # duplicated uuids: `import_data` returns (or skips) them without checking for duplicates, and + # `.get()` fails the import on duplicates instead of silently picking one. + cache = InstanceImportCache([d["id"] for d in instances_data]) + import_instances(instances_data, user, project.app_id, api_import=api_import, cache=cache) new_instance_files = [] dirs = get_directory_handlers(zip_ref) # The batch's instances share a few (form, version) pairs: look each up once. form_versions_cache = {} + # One INSERT per chunk of audit rows instead of one per instance + modifications = [] for instance_data in instances_data: uuid = instance_data["id"] - instance = instances_by_uuid.get(uuid) or Instance.objects.get(uuid=uuid) + instance = cache.instance(uuid) or Instance.objects.get(uuid=uuid) original = copy(instance) instance = process_instance_xml(instance, instance_data, zip_ref, user, form_versions_cache) if instance is None: @@ -119,7 +114,10 @@ def process_mobile_bulk_upload(api_import_id, project_id, task=None): stats["new_instances"] += 1 if uuid in dirs: new_instance_files += process_instance_attachments(zip_ref, dirs[uuid], instance) - log_modification(v1=original, v2=instance, source=BULK_UPLOAD, user=user) + modifications.append( + log_modification(v1=original, v2=instance, source=BULK_UPLOAD, user=user, save=False) + ) + Modification.objects.bulk_create(modifications, batch_size=MODIFICATIONS_BATCH_SIZE) duplicated_count = duplicate_instance_files(new_instance_files) stats["new_instance_files"] = len(new_instance_files) + duplicated_count diff --git a/iaso/tests/api/instances/test_import_cache.py b/iaso/tests/api/instances/test_import_cache.py new file mode 100644 index 0000000000..394718a73b --- /dev/null +++ b/iaso/tests/api/instances/test_import_cache.py @@ -0,0 +1,51 @@ +from iaso import models as m +from iaso.api.instances.import_cache import InstanceImportCache +from iaso.test import TestCase + + +class InstanceImportCacheTestCase(TestCase): + @classmethod + def setUpTestData(cls): + cls.account = m.Account.objects.create(name="Account") + cls.data_source = m.DataSource.objects.create(name="Source") + cls.version = m.SourceVersion.objects.create(data_source=cls.data_source, number=1) + cls.org_unit = m.OrgUnit.objects.create(name="Facility", uuid="facility-uuid", version=cls.version) + cls.instance = m.Instance.objects.create(uuid="instance-uuid", file_name="instance.xml") + # Older data can have several instances with the same uuid + cls.duplicate_1 = m.Instance.objects.create(uuid="duplicated-uuid", file_name="duplicate_1.xml") + cls.duplicate_2 = m.Instance.objects.create(uuid="duplicated-uuid", file_name="duplicate_2.xml") + + def test_existing_instances_in_a_single_query(self): + with self.assertNumQueries(1): + cache = InstanceImportCache(["instance-uuid", "duplicated-uuid", "new-uuid", None, ""]) + + with self.assertNumQueries(0): + self.assertEqual(cache.instances("instance-uuid"), [self.instance]) + self.assertEqual(cache.instance("instance-uuid"), self.instance) + self.assertEqual(cache.instances("new-uuid"), []) + self.assertIsNone(cache.instance("new-uuid")) + + def test_duplicated_uuid(self): + cache = InstanceImportCache(["duplicated-uuid"]) + + # Oldest first, as `.first()` would pick, and no single instance to pick + self.assertEqual(cache.instances("duplicated-uuid"), [self.duplicate_1, self.duplicate_2]) + self.assertIsNone(cache.instance("duplicated-uuid")) + + def test_add_instance(self): + cache = InstanceImportCache(["new-uuid"]) + instance = m.Instance.objects.create(uuid="new-uuid", file_name="new.xml") + + cache.add_instance(instance) + + self.assertEqual(cache.instance("new-uuid"), instance) + + def test_org_unit_looked_up_once(self): + cache = InstanceImportCache([]) + + with self.assertNumQueries(1): + self.assertEqual(cache.org_unit("facility-uuid", self.version.id), self.org_unit) + self.assertEqual(cache.org_unit("facility-uuid", self.version.id), self.org_unit) + + with self.assertRaises(m.OrgUnit.DoesNotExist): + cache.org_unit("facility-uuid", None) diff --git a/iaso/tests/tasks/test_process_mobile_bulk_upload.py b/iaso/tests/tasks/test_process_mobile_bulk_upload.py index 578e6410e0..1c93262d91 100644 --- a/iaso/tests/tasks/test_process_mobile_bulk_upload.py +++ b/iaso/tests/tasks/test_process_mobile_bulk_upload.py @@ -637,21 +637,23 @@ def test_form_version_query_count_baseline(self): # by every instance of the batch (so the later `instance.form` accesses in import_data(), # xml_file_to_json and the conversions don't re-query it). `iaso_entitytype`: 1 query per # distinct entity type, whose reference form import_data() caches for the batch. - # import_data()'s own org-unit/entity lookups are still done once per instance (no batch - # prefetch/caching yet), so `iaso_orgunit`/`iaso_entity` are O(instances) - bounded at - # their exact observed value, to be tightened once those lookups get cached. `iaso_formversion`: 1 query per distinct (form, version) pair (2 + # `iaso_orgunit`: 5 to import orgUnits.json, then 1 per distinct org unit, which import_data() + # caches for the batch. Its entity lookups are still done once per instance (no batch + # prefetch/caching yet), so `iaso_entity` is O(instances) - bounded at its exact observed + # value, to be tightened once those lookups get cached. `iaso_formversion`: 1 query per distinct (form, version) pair (2 # here), from `xml_file_to_json` via the batch's `form_versions_cache` - # `get_and_save_json_of_xml` reuses the FormVersion it already found there instead of - # looking it up again via `resolve_form_version()`. `iaso_instance`: 6/instance - - # `import_data()`'s dedup filter (1) + get_or_create (2) + final save (1), then + # looking it up again via `resolve_form_version()`. `iaso_instance`: 5/instance - + # `import_data()`'s get_or_create (2) + final save (1), then # `process_instance_file()`'s file-persisting save (1, required before # `get_and_save_json_of_xml()` can fetch the file back on S3 storage) + one merged save (1, # covers json/form_version and the location/device/correlation conversions together) - # down from 8/instance, since the previous separate uuid re-fetch and the previous 3rd - # save are both gone - plus 1 per batch for `process_mobile_bulk_upload()`'s pre-existing - # uuids lookup. `iaso_instance`/`audit_modification` scale 1:1 with the batch (6 and - # 1 per instance) - a regression would push these to a multiple of the bounds below, not a - # small overshoot. The rest (`iaso_task`/`iaso_tasklog`/`iaso_project`/ + # save are both gone, then 6/instance, since import_data() reuses + # `process_mobile_bulk_upload()`'s pre-existing uuids lookup (1 per batch) instead of its own + # per-instance dedup filter. `iaso_instance` scales 1:1 with the batch (5 per instance) - a + # regression would push it to a multiple of the bound below, not a small overshoot. + # `audit_modification`: a single bulk INSERT for the whole batch, instead of 1 per instance. The rest (`iaso_task`/`iaso_tasklog`/`iaso_project`/ # `vector_control_apiimport`/`auth_user`/`iaso_profile`/`iaso_account`/`iaso_datasource`/ # `iaso_instancefile`) are fixed per-run bookkeeping overhead, unrelated to instance count - # bounded at their exact observed value so a new query pattern on any of them still gets @@ -664,12 +666,12 @@ def test_form_version_query_count_baseline(self): ): profiler.assertLessEqualQueryCount( { - "iaso_orgunit": 9, + "iaso_orgunit": 6, "iaso_form": 1, "iaso_entity": 8, "iaso_formversion": 2, - "iaso_instance": 27, - "audit_modification": 4, + "iaso_instance": 23, + "audit_modification": 1, "iaso_entitytype": 1, "iaso_task": 7, "iaso_tasklog": 4, @@ -683,7 +685,7 @@ def test_form_version_query_count_baseline(self): }, exclude=["django_content_type"], ) - self.assertLessEqual(profiler.total_queries(exclude=["django_content_type"]), 96) + self.assertLessEqual(profiler.total_queries(exclude=["django_content_type"]), 86) # S3 round trips: the zip download, then 1 upload per instance XML + 1 for the attachment (its duplicate # reuses the stored file), each preceded by an `exists()` HEAD as `AWS_S3_FILE_OVERWRITE = False`. It was 9 # `exists()` when the 128-char XML paths were over Django's default `max_length` of 100, each also checking @@ -733,10 +735,11 @@ def test_form_version_query_count_at_scale(self): self.assertEqual(m.Entity.objects.count(), 25) # Same per-instance import_data() lookups as the baseline test (see comment there), at - # 12.5x the instance count (50 vs 4): `iaso_orgunit`/`iaso_entity` scale with the batch. + # 12.5x the instance count (50 vs 4): `iaso_entity` scales with the batch, `iaso_orgunit` stays at + # the baseline's 6 (a single org unit). # `iaso_form`/`iaso_entitytype` stay at 1 query per distinct form/entity type, as in the baseline. `iaso_formversion`: same 2 distinct (form, version) pairs as the baseline, so still 2. - # `iaso_instance`: 6/instance as the baseline test (see comment there), at scale: 300, + 1 per batch. - # `audit_modification` scales 1:1 with the batch, at scale: 50. The rest is the same fixed + # `iaso_instance`: 5/instance as the baseline test (see comment there), at scale: 250, + 1 per batch. + # `audit_modification`: still a single bulk INSERT at scale. The rest is the same fixed # per-run bookkeeping overhead as the baseline test (see comment there), still O(1) rather # than scaling with the 12.5x larger batch - this zip has no attachments so # `iaso_instancefile` never fires, unlike the baseline test. @@ -745,12 +748,12 @@ def test_form_version_query_count_at_scale(self): ): profiler.assertLessEqualQueryCount( { - "iaso_orgunit": 55, + "iaso_orgunit": 6, "iaso_form": 1, "iaso_entity": 100, "iaso_formversion": 2, - "iaso_instance": 302, - "audit_modification": 50, + "iaso_instance": 252, + "audit_modification": 1, "iaso_entitytype": 1, "iaso_task": 7, "iaso_tasklog": 4, @@ -763,7 +766,7 @@ def test_form_version_query_count_at_scale(self): }, exclude=["django_content_type"], ) - self.assertLessEqual(profiler.total_queries(exclude=["django_content_type"]), 643) + self.assertLessEqual(profiler.total_queries(exclude=["django_content_type"]), 495) # S3 round trips: the zip download, then 1 upload per instance XML (no attachment in this zip), each # preceded by an `exists()` HEAD as `AWS_S3_FILE_OVERWRITE = False` makes it look for an available name - # 100 when the XML paths were over the field's `max_length`, each also checking its truncated name. @@ -771,13 +774,14 @@ def test_form_version_query_count_at_scale(self): def test_audit_trail_at_scale(self): """ - A batch of 50 new instances: one audit row each, with the instance as created by `import_data()` (no file nor - json yet) as past value, and the fully processed instance as new value. + More new instances (520) than `MODIFICATIONS_BATCH_SIZE`: still one audit row each, bulk inserted in 2 chunks, + with the instance as created by `import_data()` (no file nor json yet) as past value, and the fully processed + instance as new value. """ - # Setup: 25 patients, each with a registration and a follow-up + # Setup: 260 patients, each with a registration and a follow-up m.FormVersion.objects.create(form=self.form_registration, version_id="2024032701") m.FormVersion.objects.create(form=self.form_catt, version_id="2024031801") - self._create_zip_file_at_scale(num_patients=25) + self._create_zip_file_at_scale(num_patients=260) # Exercise with QueryProfiler(trace_tables=["audit_modification"]) as profiler: @@ -790,10 +794,10 @@ def test_audit_trail_at_scale(self): # Assert instances = m.Instance.objects.all() - self.assertEqual(instances.count(), 50) + self.assertEqual(instances.count(), 520) content_type = ContentType.objects.get_for_model(m.Instance) modifications = Modification.objects.filter(source=BULK_UPLOAD) - self.assertEqual(modifications.count(), 50) + self.assertEqual(modifications.count(), 520) self.assertEqual( sorted(modifications.values_list("object_id", flat=True)), sorted(str(instance_id) for instance_id in instances.values_list("id", flat=True)), @@ -821,8 +825,8 @@ def test_audit_trail_at_scale(self): }, }, ) - # One INSERT per audit row - self.assertEqual(profiler.table_counts()["audit_modification"], 50) + # One INSERT per chunk of `MODIFICATIONS_BATCH_SIZE` audit rows + self.assertEqual(profiler.table_counts()["audit_modification"], 2) def test_campaign_registrations_then_followups_query_count(self): """ @@ -855,11 +859,11 @@ def test_campaign_registrations_then_followups_query_count(self): api_import_id=followups_import.id, project_id=self.project.id, task=followups_task, _immediate=True ) - # The batch's pre-existing uuids lookup and `import_data()`'s file names lookup (2), then 6 per new follow-up - # (see test_form_version_query_count_baseline) and 5 per updated registration: looked up by uuid 3 times (twice - # by `import_data()`, once by `process_mobile_bulk_upload()`), then its 2 saves (with the new file, then with - # its json). - self.assertEqual(profiler.table_counts()["iaso_instance"], 2 + 25 * 6 + 25 * 5) + # The batch's instances (`InstanceImportCache`) and file names lookups (2), then 5 per new follow-up (see + # test_form_version_query_count_baseline) and 2 per updated registration (its saves, with the new file then with + # its json): the registrations are found in the batch's instances. It was 277: 3 more per existing registration + # and 1 more per new follow-up, each looked up by uuid on its own. + self.assertEqual(profiler.table_counts()["iaso_instance"], 2 + 25 * 5 + 25 * 2) self.assertEqual(m.Instance.objects.count(), 50) self.assertEqual(m.Entity.objects.count(), 25) @@ -2105,9 +2109,9 @@ def test_submission_already_received_by_one_by_one_upload(self): "instances": [registration_uuid], }, ) - # Queries: the batch's uuids and file names lookups (2), its uuid looked up (1), found by its file name (1), + # Queries: the batch's instances (`InstanceImportCache`) and file names lookups (2), found by its file name (1), # saved by `import_data()` (1), then twice with its XML and data (2). S3: the zip, then the XML stored again. - self.assertEqual(profiler.table_counts()["iaso_instance"], 7) + self.assertEqual(profiler.table_counts()["iaso_instance"], 6) profiler.assertStorageCounts({"_open": 1, "exists": 1, "_save": 1}) def test_approved_instance_sent_again(self): @@ -2136,10 +2140,10 @@ def test_approved_instance_sent_again(self): updated = instance_state(registration_uuid) self.assertNotEqual(updated["file"], approved["file"]) self.assertEqual(updated, {**approved, "file": updated["file"], "updated at": a_day_later}) - # Queries: the batch's uuids and file names lookups (2), its uuid looked up by `import_data()` (1) - which - # leaves it out, as approved - and by the task (1), then saved twice with its new XML and data (2). S3: the zip, + # Queries: the batch's instances (`InstanceImportCache`) and file names lookups (2) - where `import_data()` (which + # leaves it out, as approved) and the task find it - then saved twice with its new XML and data (2). S3: the zip, # then the new XML: its name is taken by the previous one, so a 2nd HEAD for the name it's stored under. - self.assertEqual(profiler.table_counts()["iaso_instance"], 6) + self.assertEqual(profiler.table_counts()["iaso_instance"], 4) profiler.assertStorageCounts({"_open": 1, "exists": 2, "_save": 1}) def test_existing_instance_sent_again_without_its_xml(self): @@ -2169,9 +2173,9 @@ def test_existing_instance_sent_again_without_its_xml(self): # Assert: nothing changed self.assertEqual(instance_state(registration_uuid), received) self.assertEqual(entity_state(registration_uuid), patient) - # Queries: the batch's uuids and file names lookups (2), its uuid looked up twice by `import_data()` (2) and once - # by the task (1). S3: only the zip, nothing stored. - self.assertEqual(profiler.table_counts()["iaso_instance"], 5) + # Queries: the batch's instances (`InstanceImportCache`) and file names lookups (2) - where it's found - and + # nothing else. S3: only the zip, nothing stored. + self.assertEqual(profiler.table_counts()["iaso_instance"], 2) profiler.assertStorageCounts({"_open": 1}) def test_new_followup_without_its_xml(self): @@ -2202,9 +2206,9 @@ def test_new_followup_without_its_xml(self): self.assertIsNone(instance_state(followup_uuid)) self.assertEqual(instance_state(registration_uuid), registration) self.assertEqual(entity_state(registration_uuid), patient) - # Queries: the batch's uuids and file names lookups (2), its uuid looked up (1), created by `import_data()` (2) - # and saved (1), then deleted (1). S3: only the zip, nothing stored. - self.assertEqual(profiler.table_counts()["iaso_instance"], 7) + # Queries: the batch's instances (`InstanceImportCache`) and file names lookups (2), created by `import_data()` + # (2) and saved (1), then deleted (1). S3: only the zip, nothing stored. + self.assertEqual(profiler.table_counts()["iaso_instance"], 6) profiler.assertStorageCounts({"_open": 1}) def _sku(self): @@ -2336,10 +2340,10 @@ def test_updated_followup_of_merged_entity(self): updated_followup, {**received_followup, "file": updated_followup["file"], "updated at": a_day_later} ) self.assertEqual(instance_state(active_registration_uuid), active_registration) - # Queries: the batch's uuids and file names lookups (2), its uuid looked up twice by `import_data()` (2) and once - # by the task (1), saved twice with its new XML and data (2), then the merged patient's reference instance - # loaded, to see it's not this follow-up (1). S3: the zip, then the new XML (its name is taken: a 2nd HEAD). - self.assertEqual(profiler.table_counts()["iaso_instance"], 8) + # Queries: the batch's instances (`InstanceImportCache`) and file names lookups (2) - where it's found - saved + # twice with its new XML and data (2), then the merged patient's reference instance loaded, to see it's not this + # follow-up (1). S3: the zip, then the new XML (its name is taken: a 2nd HEAD). + self.assertEqual(profiler.table_counts()["iaso_instance"], 5) profiler.assertStorageCounts({"_open": 1, "exists": 2, "_save": 1}) def test_older_registration_update_of_merged_entity(self): @@ -2372,9 +2376,9 @@ def test_older_registration_update_of_merged_entity(self): updated = instance_state(registration_uuid) self.assertEqual(updated, {**merged_registration, "file": updated["file"], "updated at": a_day_later}) self.assertEqual(instance_state(active_registration_uuid), active_registration) - # Queries: as in test_updated_followup_of_merged_entity (8), plus the active patient's registration loaded, to + # Queries: as in test_updated_followup_of_merged_entity (5), plus the active patient's registration loaded, to # compare their dates (1). S3: the zip, then the new XML (its name is taken: a 2nd HEAD). - self.assertEqual(profiler.table_counts()["iaso_instance"], 9) + self.assertEqual(profiler.table_counts()["iaso_instance"], 6) profiler.assertStorageCounts({"_open": 1, "exists": 2, "_save": 1}) def test_unknown_project(self): From 7801409289b81d16b59e1f14b952663483132152 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?St=C3=A9phan=20Mestach?= Date: Mon, 5 Oct 2026 22:44:41 +0200 Subject: [PATCH 4/5] IA-5486 Bulk zip upload: cache entities per batch --- iaso/api/instances/import_cache.py | 53 ++++++++++++++++++- iaso/api/instances/views.py | 29 ++++++++-- iaso/tests/api/instances/test_import_cache.py | 50 +++++++++++++++++ .../tasks/test_process_mobile_bulk_upload.py | 20 +++---- 4 files changed, 137 insertions(+), 15 deletions(-) diff --git a/iaso/api/instances/import_cache.py b/iaso/api/instances/import_cache.py index 4d3541ed8b..1dab812a8d 100644 --- a/iaso/api/instances/import_cache.py +++ b/iaso/api/instances/import_cache.py @@ -1,7 +1,12 @@ +import uuid as uuid_lib + from collections import defaultdict from typing import Dict, Iterable, List, Optional, Tuple -from iaso.models import Instance, OrgUnit +from iaso.models import Account, Entity, Instance, OrgUnit + + +EntityKey = Tuple[str, int] class InstanceImportCache: @@ -12,15 +17,22 @@ class InstanceImportCache: - Instances: the ones already in the database for the batch's uuids, in a single query, plus the ones created while importing it (see `add_instance()`), so that a later entry of the batch with the same uuid finds them. - Org units given by uuid: looked up the first time they're seen (a batch typically covers a few facilities). + - Entities: the ones already in the database for the batch's (entity uuid, entity type), in a single query (see + `prefetch_entities()`), plus the ones created while importing it (see `add_entity()`). """ def __init__(self, instance_uuids: Iterable[Optional[str]]): # Lists: older instances can share a uuid (the unique index only covers recent ones). Ordered by id, so the # first one is the one `.first()` would pick. + # With their entity, read (never saved) when updating them: the entities the import saves are the ones from + # `prefetch_entities()`. self._instances_by_uuid: Dict[str, List[Instance]] = defaultdict(list) - for instance in Instance.objects.filter(uuid__in=[uuid for uuid in instance_uuids if uuid]).order_by("id"): + existing = Instance.objects.filter(uuid__in=[uuid for uuid in instance_uuids if uuid]).select_related("entity") + for instance in existing.order_by("id"): self._instances_by_uuid[instance.uuid].append(instance) self._org_units: Dict[Tuple[str, Optional[int]], OrgUnit] = {} + # Lists too: older entities can share a uuid. Only for the keys given to `prefetch_entities()`. + self._entities: Dict[EntityKey, List[Entity]] = {} def instances(self, uuid: str) -> List[Instance]: """All the instances with this uuid, oldest first.""" @@ -40,3 +52,40 @@ def org_unit(self, uuid: str, version_id: Optional[int]) -> OrgUnit: if key not in self._org_units: self._org_units[key] = OrgUnit.objects.get(uuid=uuid, version_id=version_id) return self._org_units[key] + + def prefetch_entities(self, account: Account, entity_uuids_and_type_ids: Iterable[Tuple[str, str]]) -> None: + """Look up, in a single query, the entities (even soft-deleted) of `account` for these (uuid, entity type).""" + keys = {_entity_key(entity_uuid, entity_type_id) for entity_uuid, entity_type_id in entity_uuids_and_type_ids} + keys.discard(None) + if not keys: + return + for key in keys: + self._entities[key] = [] + found = Entity.objects_include_deleted.filter( + account=account, + uuid__in={entity_uuid for entity_uuid, _ in keys}, + entity_type_id__in={entity_type_id for _, entity_type_id in keys}, + ) + for entity in found: + key = _entity_key(entity.uuid, entity.entity_type_id) + if key in self._entities: + self._entities[key].append(entity) + + def entities(self, entity_uuid: str, entity_type_id: str) -> Optional[List[Entity]]: + """The entities for this (uuid, entity type) - see `find_entity()` - or None if they weren't prefetched.""" + entities = self._entities.get(_entity_key(entity_uuid, entity_type_id)) + return list(entities) if entities is not None else None + + def add_entity(self, entity: Entity) -> None: + """Record an entity created while importing the batch.""" + key = _entity_key(entity.uuid, entity.entity_type_id) + if key: + self._entities.setdefault(key, []).append(entity) + + +def _entity_key(entity_uuid, entity_type_id) -> Optional[EntityKey]: + """(uuid, entity type id) normalized - the mobile app sends them as strings - or None if they aren't valid.""" + try: + return str(uuid_lib.UUID(str(entity_uuid))), int(entity_type_id) + except (TypeError, ValueError): + return None diff --git a/iaso/api/instances/views.py b/iaso/api/instances/views.py index b70540c794..a4d5918536 100644 --- a/iaso/api/instances/views.py +++ b/iaso/api/instances/views.py @@ -7,7 +7,7 @@ from collections import defaultdict from copy import copy from time import gmtime, strftime -from typing import Any, Dict, Optional, Union +from typing import Any, Dict, List, Optional, Union import pandas as pd @@ -1086,7 +1086,16 @@ def download_attachments(self, request, pk=None) -> StreamingHttpResponse: ) -def find_entity(account: Account, entity_uuid: str, entity_type_id: Optional[int] = None) -> Entity: +def find_entity( + account: Account, + entity_uuid: str, + entity_type_id: Optional[int] = None, + existing_entities: Optional[List[Entity]] = None, +) -> Entity: + """ + `existing_entities`: the entities matching (even soft-deleted), when the caller already looked them up - see + `InstanceImportCache.entities()`. + """ # In case of duplicate UUIDs in the database, only allow 1 non-deleted one. # If a non-deleted entity was found, ignore potential duplicates. if entity_type_id is not None: @@ -1100,7 +1109,8 @@ def find_entity(account: Account, entity_uuid: str, entity_type_id: Optional[int "uuid": entity_uuid, "account": account, } - existing_entities = list(Entity.objects_include_deleted.filter(**filters)) + if existing_entities is None: + existing_entities = list(Entity.objects_include_deleted.filter(**filters)) if len(existing_entities) == 0: if entity_type_id is not None: @@ -1150,6 +1160,14 @@ def import_data(instances, user, app_id, api_import=None, cache: Optional[Instan if cache is None: cache = InstanceImportCache([instance_data.get("id") for instance_data in instances]) + cache.prefetch_entities( + project.account, + [ + (instance_data["entityUuid"], instance_data["entityTypeId"]) + for instance_data in instances + if instance_data.get("entityUuid") and instance_data.get("entityTypeId") + ], + ) for instance_data in instances: uuid = instance_data.get("id", None) @@ -1216,7 +1234,10 @@ def import_data(instances, user, app_id, api_import=None, cache: Optional[Instan entity_uuid = instance_data.get("entityUuid", None) entity_type_id = instance_data.get("entityTypeId", None) if entity_uuid and entity_type_id: - entity = find_entity(project.account, entity_uuid, entity_type_id) + existing_entities = cache.entities(entity_uuid, entity_type_id) + entity = find_entity(project.account, entity_uuid, entity_type_id, existing_entities=existing_entities) + if existing_entities == []: # none yet: `find_entity()` created it + cache.add_entity(entity) if entity.deleted_at: logger.info( diff --git a/iaso/tests/api/instances/test_import_cache.py b/iaso/tests/api/instances/test_import_cache.py index 394718a73b..3aaca75ad8 100644 --- a/iaso/tests/api/instances/test_import_cache.py +++ b/iaso/tests/api/instances/test_import_cache.py @@ -1,3 +1,7 @@ +import uuid + +from django.utils import timezone + from iaso import models as m from iaso.api.instances.import_cache import InstanceImportCache from iaso.test import TestCase @@ -49,3 +53,49 @@ def test_org_unit_looked_up_once(self): with self.assertRaises(m.OrgUnit.DoesNotExist): cache.org_unit("facility-uuid", None) + + def test_existing_entities_in_a_single_query(self): + entity_type = m.EntityType.objects.create(name="Patient", account=self.account) + patient = m.Entity.objects.create(uuid=uuid.uuid4(), entity_type=entity_type, account=self.account) + # Older data can have several entities with the same uuid, even soft-deleted ones + duplicated_uuid = uuid.uuid4() + duplicate = m.Entity.objects.create(uuid=duplicated_uuid, entity_type=entity_type, account=self.account) + deleted_duplicate = m.Entity.objects.create( + uuid=duplicated_uuid, entity_type=entity_type, account=self.account, deleted_at=timezone.now() + ) + other_account = m.Account.objects.create(name="Other account") + m.Entity.objects.create(uuid=patient.uuid, entity_type=entity_type, account=other_account) + new_uuid = str(uuid.uuid4()) + cache = InstanceImportCache([]) + + with self.assertNumQueries(1): + cache.prefetch_entities( + self.account, + [ + # As sent by the mobile app: strings, the uuid possibly in capitals + (str(patient.uuid).upper(), str(entity_type.id)), + (str(duplicated_uuid), str(entity_type.id)), + (new_uuid, str(entity_type.id)), + ("not-a-uuid", str(entity_type.id)), + ], + ) + + with self.assertNumQueries(0): + self.assertEqual(cache.entities(str(patient.uuid), entity_type.id), [patient]) + self.assertCountEqual(cache.entities(str(duplicated_uuid), entity_type.id), [duplicate, deleted_duplicate]) + self.assertEqual(cache.entities(new_uuid, entity_type.id), []) + # Not prefetched: up to `find_entity()` to look them up + self.assertIsNone(cache.entities(str(uuid.uuid4()), entity_type.id)) + self.assertIsNone(cache.entities("not-a-uuid", entity_type.id)) + + def test_add_entity(self): + entity_type = m.EntityType.objects.create(name="Patient", account=self.account) + new_uuid = str(uuid.uuid4()) + cache = InstanceImportCache([]) + cache.prefetch_entities(self.account, [(new_uuid, str(entity_type.id))]) + # Created as `find_entity()` does, with the payload's strings + entity = m.Entity.objects.create(uuid=new_uuid, entity_type_id=str(entity_type.id), account=self.account) + + cache.add_entity(entity) + + self.assertEqual(cache.entities(new_uuid, str(entity_type.id)), [entity]) diff --git a/iaso/tests/tasks/test_process_mobile_bulk_upload.py b/iaso/tests/tasks/test_process_mobile_bulk_upload.py index 1c93262d91..7326c5ff8e 100644 --- a/iaso/tests/tasks/test_process_mobile_bulk_upload.py +++ b/iaso/tests/tasks/test_process_mobile_bulk_upload.py @@ -638,9 +638,8 @@ def test_form_version_query_count_baseline(self): # xml_file_to_json and the conversions don't re-query it). `iaso_entitytype`: 1 query per # distinct entity type, whose reference form import_data() caches for the batch. # `iaso_orgunit`: 5 to import orgUnits.json, then 1 per distinct org unit, which import_data() - # caches for the batch. Its entity lookups are still done once per instance (no batch - # prefetch/caching yet), so `iaso_entity` is O(instances) - bounded at its exact observed - # value, to be tightened once those lookups get cached. `iaso_formversion`: 1 query per distinct (form, version) pair (2 + # caches for the batch. `iaso_entity`: the batch's entities looked up once (1), then for each + # new entity, its creation (1) and its reference instance set (1). `iaso_formversion`: 1 query per distinct (form, version) pair (2 # here), from `xml_file_to_json` via the batch's `form_versions_cache` - # `get_and_save_json_of_xml` reuses the FormVersion it already found there instead of # looking it up again via `resolve_form_version()`. `iaso_instance`: 5/instance - @@ -668,7 +667,7 @@ def test_form_version_query_count_baseline(self): { "iaso_orgunit": 6, "iaso_form": 1, - "iaso_entity": 8, + "iaso_entity": 5, "iaso_formversion": 2, "iaso_instance": 23, "audit_modification": 1, @@ -685,7 +684,7 @@ def test_form_version_query_count_baseline(self): }, exclude=["django_content_type"], ) - self.assertLessEqual(profiler.total_queries(exclude=["django_content_type"]), 86) + self.assertLessEqual(profiler.total_queries(exclude=["django_content_type"]), 83) # S3 round trips: the zip download, then 1 upload per instance XML + 1 for the attachment (its duplicate # reuses the stored file), each preceded by an `exists()` HEAD as `AWS_S3_FILE_OVERWRITE = False`. It was 9 # `exists()` when the 128-char XML paths were over Django's default `max_length` of 100, each also checking @@ -735,8 +734,8 @@ def test_form_version_query_count_at_scale(self): self.assertEqual(m.Entity.objects.count(), 25) # Same per-instance import_data() lookups as the baseline test (see comment there), at - # 12.5x the instance count (50 vs 4): `iaso_entity` scales with the batch, `iaso_orgunit` stays at - # the baseline's 6 (a single org unit). + # 12.5x the instance count (50 vs 4): `iaso_entity` scales with the new entities only (1 + 2 x 25), + # `iaso_orgunit` stays at the baseline's 6 (a single org unit). # `iaso_form`/`iaso_entitytype` stay at 1 query per distinct form/entity type, as in the baseline. `iaso_formversion`: same 2 distinct (form, version) pairs as the baseline, so still 2. # `iaso_instance`: 5/instance as the baseline test (see comment there), at scale: 250, + 1 per batch. # `audit_modification`: still a single bulk INSERT at scale. The rest is the same fixed @@ -750,7 +749,7 @@ def test_form_version_query_count_at_scale(self): { "iaso_orgunit": 6, "iaso_form": 1, - "iaso_entity": 100, + "iaso_entity": 51, "iaso_formversion": 2, "iaso_instance": 252, "audit_modification": 1, @@ -766,7 +765,7 @@ def test_form_version_query_count_at_scale(self): }, exclude=["django_content_type"], ) - self.assertLessEqual(profiler.total_queries(exclude=["django_content_type"]), 495) + self.assertLessEqual(profiler.total_queries(exclude=["django_content_type"]), 446) # S3 round trips: the zip download, then 1 upload per instance XML (no attachment in this zip), each # preceded by an `exists()` HEAD as `AWS_S3_FILE_OVERWRITE = False` makes it look for an available name - # 100 when the XML paths were over the field's `max_length`, each also checking its truncated name. @@ -864,6 +863,9 @@ def test_campaign_registrations_then_followups_query_count(self): # its json): the registrations are found in the batch's instances. It was 277: 3 more per existing registration # and 1 more per new follow-up, each looked up by uuid on its own. self.assertEqual(profiler.table_counts()["iaso_instance"], 2 + 25 * 5 + 25 * 2) + # Entities: the batch's looked up once (1), the registrations' coming with them (read when updating them). It + # was 50: each follow-up's patient looked up on its own, each updated registration's entity loaded on its own. + self.assertEqual(profiler.table_counts()["iaso_entity"], 1) self.assertEqual(m.Instance.objects.count(), 50) self.assertEqual(m.Entity.objects.count(), 25) From 60c7d6c65e4476903be470e8674c9500fb6256e0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?St=C3=A9phan=20Mestach?= Date: Mon, 5 Oct 2026 22:45:06 +0200 Subject: [PATCH 5/5] IA-5486 Move forms, entity types and form versions into InstanceImportCache --- iaso/api/instances/import_cache.py | 24 ++++++++++++++- iaso/api/instances/views.py | 24 ++++++--------- iaso/tasks/process_mobile_bulk_upload.py | 4 +-- iaso/tests/api/instances/test_import_cache.py | 29 +++++++++++++++++++ 4 files changed, 62 insertions(+), 19 deletions(-) diff --git a/iaso/api/instances/import_cache.py b/iaso/api/instances/import_cache.py index 1dab812a8d..39848065d3 100644 --- a/iaso/api/instances/import_cache.py +++ b/iaso/api/instances/import_cache.py @@ -3,7 +3,7 @@ from collections import defaultdict from typing import Dict, Iterable, List, Optional, Tuple -from iaso.models import Account, Entity, Instance, OrgUnit +from iaso.models import Account, Entity, Form, Instance, OrgUnit EntityKey = Tuple[str, int] @@ -19,6 +19,8 @@ class InstanceImportCache: - Org units given by uuid: looked up the first time they're seen (a batch typically covers a few facilities). - Entities: the ones already in the database for the batch's (entity uuid, entity type), in a single query (see `prefetch_entities()`), plus the ones created while importing it (see `add_entity()`). + - Forms: the batch's, in a single query (see `prefetch_forms()`), and the reference form of each entity type. + - Form versions: `form_versions`, the dict `FormVersionManager.find_for_form()` fills, passed down to the models. """ def __init__(self, instance_uuids: Iterable[Optional[str]]): @@ -33,6 +35,10 @@ def __init__(self, instance_uuids: Iterable[Optional[str]]): self._org_units: Dict[Tuple[str, Optional[int]], OrgUnit] = {} # Lists too: older entities can share a uuid. Only for the keys given to `prefetch_entities()`. self._entities: Dict[EntityKey, List[Entity]] = {} + self._forms: Dict[int, Form] = {} + self._reference_form_ids: Dict[int, Optional[int]] = {} + # See `FormVersionManager.find_for_form()` + self.form_versions: dict = {} def instances(self, uuid: str) -> List[Instance]: """All the instances with this uuid, oldest first.""" @@ -82,6 +88,22 @@ def add_entity(self, entity: Entity) -> None: if key: self._entities.setdefault(key, []).append(entity) + def prefetch_forms(self, form_ids: Iterable[int]) -> None: + """Look up these forms in a single query.""" + self._forms.update((form.id, form) for form in Form.objects.filter(id__in=set(form_ids))) + + def form(self, form_id: Optional[int]) -> Optional[Form]: + """The form, if prefetched: shared by the batch's instances, so that `instance.form` doesn't query it again.""" + return self._forms.get(form_id) + + def reference_form_id(self, entity: Entity) -> Optional[int]: + """The reference form of the entity's type, looked up once per entity type (a batch typically has one).""" + # int(): entities created by `find_entity()` keep the payload's string entityTypeId + entity_type_id = int(entity.entity_type_id) + if entity_type_id not in self._reference_form_ids: + self._reference_form_ids[entity_type_id] = entity.entity_type.reference_form_id + return self._reference_form_ids[entity_type_id] + def _entity_key(entity_uuid, entity_type_id) -> Optional[EntityKey]: """(uuid, entity type id) normalized - the mobile app sends them as strings - or None if they aren't valid.""" diff --git a/iaso/api/instances/views.py b/iaso/api/instances/views.py index a4d5918536..b191ace4ba 100644 --- a/iaso/api/instances/views.py +++ b/iaso/api/instances/views.py @@ -1143,13 +1143,6 @@ def import_data(instances, user, app_id, api_import=None, cache: Optional[Instan project = Project.objects.get_for_user_and_app_id(user, app_id) rtn_instances = [] - # A batch commonly has several instances (e.g. a registration + follow-up forms) pointing at - # the same few forms - fetch those once up front instead of lazily once per instance. - form_ids = {int(instance_data["formId"]) for instance_data in instances if instance_data.get("formId") is not None} - forms_by_id = {form.id: form for form in Form.objects.filter(id__in=form_ids)} - # Same for entity types (typically a single one per batch): only their reference form is needed. - reference_form_ids_by_entity_type_id = {} - file_names = {ntpath.basename(instance_data["file"]) for instance_data in instances if instance_data.get("file")} uuids_by_file_name = defaultdict(set) for existing_file_name, existing_uuid in Instance.objects.filter(file_name__in=file_names).values_list( @@ -1160,6 +1153,10 @@ def import_data(instances, user, app_id, api_import=None, cache: Optional[Instan if cache is None: cache = InstanceImportCache([instance_data.get("id") for instance_data in instances]) + # A batch commonly has several instances (e.g. a registration + follow-up forms) pointing at the same few forms + cache.prefetch_forms( + int(instance_data["formId"]) for instance_data in instances if instance_data.get("formId") is not None + ) cache.prefetch_entities( project.account, [ @@ -1225,9 +1222,10 @@ def import_data(instances, user, app_id, api_import=None, cache: Optional[Instan # Normalize to int: the mobile app sends this as a JSON string (e.g. "1"), which would make # id comparisons (e.g. against `entity_type.reference_form_id`) silently fail ("1" != 1). instance.form_id = int(raw_form_id) if raw_form_id is not None else None - if instance.form_id in forms_by_id: - # Share the prefetched Form across instances, so `instance.form` doesn't re-query it. - instance.form = forms_by_id[instance.form_id] + prefetched_form = cache.form(instance.form_id) + if prefetched_form: + # Shared by the batch's instances, so `instance.form` doesn't query it again + instance.form = prefetched_form # TODO: check that planning_id is valid instance.planning_id = instance_data.get("planningId", None) @@ -1263,11 +1261,7 @@ def import_data(instances, user, app_id, api_import=None, cache: Optional[Instan instance.entity = entity # If instance's form is the same as the type reference form, set the instance as reference_instance - # int(): entities created above by `find_entity()` keep the payload's string entityTypeId. - entity_type_id = int(entity.entity_type_id) - if entity_type_id not in reference_form_ids_by_entity_type_id: - reference_form_ids_by_entity_type_id[entity_type_id] = entity.entity_type.reference_form_id - if reference_form_ids_by_entity_type_id[entity_type_id] == instance.form_id: + if cache.reference_form_id(entity) == instance.form_id: entity.attributes = instance entity.save() diff --git a/iaso/tasks/process_mobile_bulk_upload.py b/iaso/tasks/process_mobile_bulk_upload.py index aaebf9494a..46add5d1e6 100644 --- a/iaso/tasks/process_mobile_bulk_upload.py +++ b/iaso/tasks/process_mobile_bulk_upload.py @@ -99,8 +99,6 @@ def process_mobile_bulk_upload(api_import_id, project_id, task=None): import_instances(instances_data, user, project.app_id, api_import=api_import, cache=cache) new_instance_files = [] dirs = get_directory_handlers(zip_ref) - # The batch's instances share a few (form, version) pairs: look each up once. - form_versions_cache = {} # One INSERT per chunk of audit rows instead of one per instance modifications = [] @@ -108,7 +106,7 @@ def process_mobile_bulk_upload(api_import_id, project_id, task=None): uuid = instance_data["id"] instance = cache.instance(uuid) or Instance.objects.get(uuid=uuid) original = copy(instance) - instance = process_instance_xml(instance, instance_data, zip_ref, user, form_versions_cache) + instance = process_instance_xml(instance, instance_data, zip_ref, user, cache.form_versions) if instance is None: continue stats["new_instances"] += 1 diff --git a/iaso/tests/api/instances/test_import_cache.py b/iaso/tests/api/instances/test_import_cache.py index 3aaca75ad8..e0e1846a7b 100644 --- a/iaso/tests/api/instances/test_import_cache.py +++ b/iaso/tests/api/instances/test_import_cache.py @@ -99,3 +99,32 @@ def test_add_entity(self): cache.add_entity(entity) self.assertEqual(cache.entities(new_uuid, str(entity_type.id)), [entity]) + + def test_forms_in_a_single_query(self): + registration = m.Form.objects.create(name="Registration") + follow_up = m.Form.objects.create(name="Follow-up") + cache = InstanceImportCache([]) + + with self.assertNumQueries(1): + cache.prefetch_forms([registration.id, follow_up.id, registration.id]) + + with self.assertNumQueries(0): + self.assertEqual(cache.form(registration.id), registration) + self.assertEqual(cache.form(follow_up.id), follow_up) + self.assertIsNone(cache.form(0)) + self.assertIsNone(cache.form(None)) + + def test_reference_form_looked_up_once_per_entity_type(self): + registration = m.Form.objects.create(name="Registration") + entity_type = m.EntityType.objects.create(name="Patient", account=self.account, reference_form=registration) + existing_patient = m.Entity.objects.create(uuid=uuid.uuid4(), entity_type=entity_type, account=self.account) + existing_patient = m.Entity.objects.get(id=existing_patient.id) # without its entity type loaded + # Created as `find_entity()` does, with the payload's string entity type id + new_patient = m.Entity.objects.create( + uuid=uuid.uuid4(), entity_type_id=str(entity_type.id), account=self.account + ) + cache = InstanceImportCache([]) + + with self.assertNumQueries(1): + self.assertEqual(cache.reference_form_id(existing_patient), registration.id) + self.assertEqual(cache.reference_form_id(new_patient), registration.id)