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..39848065d3 --- /dev/null +++ b/iaso/api/instances/import_cache.py @@ -0,0 +1,113 @@ +import uuid as uuid_lib + +from collections import defaultdict +from typing import Dict, Iterable, List, Optional, Tuple + +from iaso.models import Account, Entity, Form, Instance, OrgUnit + + +EntityKey = Tuple[str, int] + + +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). + - 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]]): + # 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) + 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]] = {} + 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.""" + 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] + + 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 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.""" + 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 d49effaec9..88ebf3ce0d 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 @@ -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 ( @@ -1092,7 +1093,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: @@ -1106,7 +1116,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: @@ -1128,22 +1139,17 @@ 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 = [] - # 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( @@ -1152,10 +1158,26 @@ 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]) + # 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, + [ + (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) - 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 @@ -1163,7 +1185,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: @@ -1180,10 +1202,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) @@ -1200,23 +1223,26 @@ 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 # 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) 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( @@ -1242,11 +1268,7 @@ def import_data(instances, user, app_id, api_import=None): 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() @@ -1292,6 +1314,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..46add5d1e6 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,38 +91,31 @@ 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) + instance = process_instance_xml(instance, instance_data, zip_ref, user, cache.form_versions) if instance is None: continue 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..e0e1846a7b --- /dev/null +++ b/iaso/tests/api/instances/test_import_cache.py @@ -0,0 +1,130 @@ +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 + + +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) + + 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]) + + 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) diff --git a/iaso/tests/tasks/test_process_mobile_bulk_upload.py b/iaso/tests/tasks/test_process_mobile_bulk_upload.py index 578e6410e0..7326c5ff8e 100644 --- a/iaso/tests/tasks/test_process_mobile_bulk_upload.py +++ b/iaso/tests/tasks/test_process_mobile_bulk_upload.py @@ -637,21 +637,22 @@ 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. `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`: 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 +665,12 @@ def test_form_version_query_count_baseline(self): ): profiler.assertLessEqualQueryCount( { - "iaso_orgunit": 9, + "iaso_orgunit": 6, "iaso_form": 1, - "iaso_entity": 8, + "iaso_entity": 5, "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 +684,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"]), 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 @@ -733,10 +734,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 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`: 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 +747,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_entity": 51, "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 +765,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"]), 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. @@ -771,13 +773,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 +793,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 +824,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 +858,14 @@ 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) + # 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) @@ -2105,9 +2111,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 +2142,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 +2175,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 +2208,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 +2342,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 +2378,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):