MDM expansion support for Patient instance bulk export operations

hapi-fhir-docs/src/main/resources/ca/uhn/hapi/fhir/changelog/8_8_0/7421-patient-bulk-export-mdm-expansion.yaml

@@ -0,0 +1,11 @@
+---
+type: add
+issue: 7421
+title: "Patient instance bulk export operation now supports MDM expansion. When the `_mdm` parameter is enabled on
+a Patient instance bulk export request (POST /Patient/123/$export?_mdm=true), the exported set of resources will include:
+<ul>
+<li>Patient/123 (the requested patient)</li>
+<li>Patient/123's golden resource (if MDM-linked)</li>
+<li>All other patients linked to the same golden resource</li>
+<li>All resources in the patient compartment referring to ANY of the above patients</li>
+</ul>"

hapi-fhir-jpaserver-base/src/main/java/ca/uhn/fhir/jpa/bulk/export/svc/JpaBulkExportProcessor.java

@@ -124,7 +124,7 @@ public Iterator<JpaPid> getResourcePidIterator(ExportPIDIteratorParameters thePa
 					String chunkId = theParams.getChunkId();
 					RuntimeResourceDefinition def = myContext.getResourceDefinition(resourceType);
 
-					LinkedHashSet<JpaPid> pids;
+					HashSet<JpaPid> pids;
 					if (theParams.getExportStyle() == BulkExportJobParameters.ExportStyle.PATIENT) {
 						pids = getPidsForPatientStyleExport(theParams, resourceType, jobId, chunkId, def);
 					} else if (theParams.getExportStyle() == BulkExportJobParameters.ExportStyle.GROUP) {
@@ -139,7 +139,7 @@ public Iterator<JpaPid> getResourcePidIterator(ExportPIDIteratorParameters thePa
 	}
 
 	@SuppressWarnings("unchecked")
-	private LinkedHashSet<JpaPid> getPidsForPatientStyleExport(
+	protected LinkedHashSet<JpaPid> getPidsForPatientStyleExport(
 			ExportPIDIteratorParameters theParams,
 			String resourceType,
 			String theJobId,
@@ -155,6 +155,8 @@ private LinkedHashSet<JpaPid> getPidsForPatientStyleExport(
 			throw new IllegalStateException(Msg.code(797) + errorMessage);
 		}
 
+		Set<String> expandedPatientIds = getExpandedPatientSetForPatientExport(theParams);
+
 		Set<String> patientSearchParams = getPatientActiveSearchParamsForResourceType(theParams.getResourceType());
 		for (String patientSearchParam : patientSearchParams) {
 			List<SearchParameterMap> maps =
@@ -165,28 +167,30 @@ private LinkedHashSet<JpaPid> getPidsForPatientStyleExport(
 
 				ISearchBuilder<JpaPid> searchBuilder = getSearchBuilderForResourceType(theParams.getResourceType());
 
-				filterBySpecificPatient(theParams, resourceType, patientSearchParam, map);
+				filterBySpecificPatient(expandedPatientIds, resourceType, patientSearchParam, map);
 
 				SearchRuntimeDetails searchRuntime = new SearchRuntimeDetails(null, theJobId);
 
 				Logs.getBatchTroubleshootingLog()
-						.debug(
-								"Executing query for bulk export job[{}] chunk[{}]: {}",
-								theJobId,
-								theChunkId,
-								map.toNormalizedQueryString(myContext));
+						.atDebug()
+						.setMessage("Executing query for bulk export job[{}] chunk[{}]: {}")
+						.addArgument(theJobId)
+						.addArgument(theChunkId)
+						.addArgument(map.toNormalizedQueryString())
+						.log();
 
 				try (IResultIterator<JpaPid> resultIterator = searchBuilder.createQuery(
 						map, searchRuntime, new SystemRequestDetails(), theParams.getPartitionIdOrAllPartitions())) {
 					int pidCount = 0;
 					while (resultIterator.hasNext()) {
 						if (pidCount % 10000 == 0) {
 							Logs.getBatchTroubleshootingLog()
-									.debug(
-											"Bulk export job[{}] chunk[{}] has loaded {} pids",
-											theJobId,
-											theChunkId,
-											pidCount);
+									.atDebug()
+									.setMessage("Bulk export job[{}] chunk[{}] has loaded {} pids")
+									.addArgument(theJobId)
+									.addArgument(theChunkId)
+									.addArgument(pidCount)
+									.log();
 						}
 						pidCount++;
 						pids.add(resultIterator.next());
@@ -198,18 +202,17 @@ map, searchRuntime, new SystemRequestDetails(), theParams.getPartitionIdOrAllPar
 	}
 
 	private void filterBySpecificPatient(
-			ExportPIDIteratorParameters theParams,
-			String resourceType,
-			String patientSearchParam,
-			SearchParameterMap map) {
+			Set<String> theExpandedPatientIds, String resourceType, String patientSearchParam, SearchParameterMap map) {
 		if (resourceType.equalsIgnoreCase("Patient")) {
-			if (theParams.getPatientIds() != null) {
-				ReferenceOrListParam referenceOrListParam = makeReferenceOrListParam(theParams.getPatientIds());
+			if (theExpandedPatientIds != null) {
+				ReferenceOrListParam referenceOrListParam =
+						makeReferenceOrListParam(new ArrayList<>(theExpandedPatientIds));
 				map.add(PARAM_ID, referenceOrListParam);
 			}
 		} else {
-			if (theParams.getPatientIds() != null) {
-				ReferenceOrListParam referenceOrListParam = makeReferenceOrListParam(theParams.getPatientIds());
+			if (theExpandedPatientIds != null) {
+				ReferenceOrListParam referenceOrListParam =
+						makeReferenceOrListParam(new ArrayList<>(theExpandedPatientIds));
 				map.add(patientSearchParam, referenceOrListParam);
 			} else {
 				map.add(patientSearchParam, new ReferenceParam().setMissing(false));
@@ -236,11 +239,11 @@ private LinkedHashSet<JpaPid> getPidsForSystemStyleExport(
 
 		for (SearchParameterMap map : maps) {
 			Logs.getBatchTroubleshootingLog()
-					.debug(
-							"Executing query for bulk export job[{}] chunk[{}]: {}",
-							theJobId,
-							theChunkId,
-							map.toNormalizedQueryString(myContext));
+					.atDebug()
+					.setMessage("Executing query for bulk export job[{}] chunk[{}]: {}")
+					.addArgument(theJobId)
+					.addArgument(theChunkId)
+					.addArgument(map.toNormalizedQueryString());
 
 			// requires a transaction
 			try (IResultIterator<JpaPid> resultIterator = searchBuilder.createQuery(
@@ -263,14 +266,14 @@ map, new SearchRuntimeDetails(null, theJobId), null, theParams.getPartitionIdOrA
 		return pids;
 	}
 
-	private LinkedHashSet<JpaPid> getPidsForGroupStyleExport(
+	private HashSet<JpaPid> getPidsForGroupStyleExport(
 			ExportPIDIteratorParameters theParams, String theResourceType, RuntimeResourceDefinition theDef)
 			throws IOException {
-		LinkedHashSet<JpaPid> pids;
+		HashSet<JpaPid> pids;
 
 		if (theResourceType.equalsIgnoreCase("Patient")) {
 			ourLog.info("Expanding Patients of a Group Bulk Export.");
-			pids = getExpandedPatientList(theParams, true);
+			pids = getExpandedPatientSetForGroupExport(theParams, true);
 			ourLog.info("Obtained {} PIDs", pids.size());
 		} else if (theResourceType.equalsIgnoreCase("Group")) {
 			pids = getSingletonGroupList(theParams);
@@ -288,15 +291,15 @@ private LinkedHashSet<JpaPid> getRelatedResourceTypePids(
 				getActivePatientSearchParamForCurrentResourceType(theParams.getResourceType());
 		if (activeSearchParam != null) {
 			// expand the group pid -> list of patients in that group (list of patient pids)
-			Set<JpaPid> expandedMemberResourceIds = getExpandedPatientList(theParams, false);
+			HashSet<JpaPid> expandedMemberResourceIds = getExpandedPatientSetForGroupExport(theParams, false);
 			assert !expandedMemberResourceIds.isEmpty();
 			Logs.getBatchTroubleshootingLog()
 					.debug("{} has been expanded to members:[{}]", theParams.getGroupId(), expandedMemberResourceIds);
 
 			// for each patient pid ->
 			//	search for the target resources, with their correct patient references, chunked.
 			// The results will be jammed into myReadPids
-			TaskChunker.chunk(expandedMemberResourceIds, QUERY_CHUNK_SIZE, (idChunk) -> {
+			TaskChunker.chunk(expandedMemberResourceIds, QUERY_CHUNK_SIZE, idChunk -> {
 				try {
 					queryResourceTypeWithReferencesToPatients(pids, idChunk, theParams, theDef);
 				} catch (IOException ex) {
@@ -309,8 +312,9 @@ private LinkedHashSet<JpaPid> getRelatedResourceTypePids(
 				}
 			});
 		} else {
-			ourLog.warn("No active patient compartment search parameter(s) for resource type "
-					+ theParams.getResourceType());
+			ourLog.warn(
+					"No active patient compartment search parameter(s) for resource type {}",
+					theParams.getResourceType());
 		}
 		return pids;
 	}
@@ -388,30 +392,95 @@ private void validateSearchParametersForGroup(SearchParameterMap expandedSpMap,
 
 	/**
 	 * Given the local myGroupId, perform an expansion to retrieve all resource IDs of member patients.
-	 * if myMdmEnabled is set to true, we also attempt to also expand it into matched
-	 * patients.
+	 * If myMdmEnabled is set to true, we also expand into MDM-matched patients.
 	 *
-	 * @return a Set of Strings representing the resource IDs of all members of a group.
+	 * CACHING: Results are cached in theParameters.myExpandedPatientIdsForGroupExport to avoid redundant expansion
+	 * across multiple resource type iterations.
+	 *
+	 * @param theParameters - export parameters containing group ID and MDM flag
+	 * @param theConsiderDateRange - whether to apply date range filters
+	 * @return a LinkedHashSet of JpaPids representing all member patients (with MDM expansion if enabled)
 	 */
-	private LinkedHashSet<JpaPid> getExpandedPatientList(
+	private HashSet<JpaPid> getExpandedPatientSetForGroupExport(
 			ExportPIDIteratorParameters theParameters, boolean theConsiderDateRange) throws IOException {
+
 		List<JpaPid> members = getMembersFromGroupWithFilter(theParameters, theConsiderDateRange);
-		ourLog.info(
-				"Group with ID [{}] has been expanded to {} members, member JpaIds: {}",
+		ourLog.debug(
+				"Group with ID [{}] has {} members, member JpaIds: {}",
 				theParameters.getGroupId(),
 				members.size(),
 				members);
 		LinkedHashSet<JpaPid> patientPidsToExport = new LinkedHashSet<>(members);
 
 		if (theParameters.isExpandMdm()) {
 			RequestPartitionId partitionId = theParameters.getPartitionIdOrAllPartitions();
-			patientPidsToExport.addAll(myMdmExpandersHolder
+
+			Set<JpaPid> singlePatientExpandedSet = myMdmExpandersHolder
 					.getBulkExportMDMResourceExpanderInstance()
-					.expandGroup(theParameters.getGroupId(), partitionId));
+					.expandGroup(theParameters.getGroupId(), partitionId);
+
+			patientPidsToExport.addAll(singlePatientExpandedSet);
+
+			ourLog.debug(
+					"Group with ID [{}] has been expanded to {} members, member JpaIds: {}",
+					theParameters.getGroupId(),
+					singlePatientExpandedSet.size(),
+					singlePatientExpandedSet);
 		}
+
 		return patientPidsToExport;
 	}
 
+	/**
+	 * Expands patient IDs for Patient-style bulk export.
+	 * If MDM expansion is enabled, expands each patient to include their MDM-linked patients.
+	 *
+	 * CACHING: Results are cached in theParams.myExpandedPatientIdsForPatientExport to avoid redundant expansion
+	 * across multiple resource type iterations.
+	 *
+	 * @param theParams - export parameters containing patient IDs and MDM flag
+	 * @return HashSet of String patient IDs for all patients (original + MDM-expanded)
+	 *
+	 * Created by Claude 4.5 Sonnet
+	 */
+	Set<String> getExpandedPatientSetForPatientExport(ExportPIDIteratorParameters theParams) {
+		if (theParams.hasExpandedPatientIdsForPatientExport()) {
+			ourLog.debug(
+					"Using cached expanded patient ID set with {} patients",
+					theParams.getExpandedPatientIdsForPatientExport().size());
+			return theParams.getExpandedPatientIdsForPatientExport();
+		}
+
+		HashSet<String> expandedPatientIds = new HashSet<>();
+
+		List<String> patientIds = theParams.getPatientIds();
+
+		if (patientIds == null || patientIds.isEmpty()) {
+			return expandedPatientIds;
+		}
+
+		expandedPatientIds.addAll(patientIds);
+
+		RequestPartitionId partitionId = theParams.getPartitionIdOrAllPartitions();
+
+		if (theParams.isExpandMdm()) {
+			ourLog.debug("MDM expansion enabled - expanding {} patients", patientIds.size());
+
+			for (String patientId : patientIds) {
+				Set<String> mdmExpandedIds = myMdmExpandersHolder
+						.getBulkExportMDMResourceExpanderInstance()
+						.expandPatient(patientId, partitionId);
+				expandedPatientIds.addAll(mdmExpandedIds);
+			}
+		}
+
+		ourLog.debug("Patient expansion resulted in {} total patient IDs", expandedPatientIds.size());
+
+		theParams.setExpandedPatientIdsForPatientExport(expandedPatientIds);
+
+		return expandedPatientIds;
+	}
+
 	/**
 	 * Given the parameters, find all members' patient references in the group with the typeFilter applied.
 	 *
@@ -533,7 +602,7 @@ expandedSpMap, new SearchRuntimeDetails(null, theParams.getInstanceId()), null,
 
 			// gets rid of the Patient duplicates
 			theReadPids.addAll(includeIds.stream()
-					.filter((id) -> !id.getResourceType().equals("Patient"))
+					.filter(id -> !id.getResourceType().equals("Patient"))
 					.collect(Collectors.toSet()));
 		}
 	}

hapi-fhir-jpaserver-base/src/test/java/ca/uhn/fhir/jpa/bulk/export/svc/JpaBulkExportProcessorTest.java

@@ -65,6 +65,7 @@
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.ArgumentMatchers.nullable;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
@@ -405,8 +406,9 @@ public void getResourcePidIterator_groupExportStyleWithNonPatientResource_return
 		Iterator<JpaPid> pidIterator = myProcessor.getResourcePidIterator(parameters);
 
 		// verify
-		assertThat(pidIterator).as("PID iterator null for mdm = " + theMdm).isNotNull();
-		assertThat(pidIterator).toIterable().containsExactly(observationPid, observationPid2);
+		assertThat(pidIterator)
+			.as("PID iterator null for mdm = " + theMdm).isNotNull()
+			.toIterable().containsExactly(observationPid, observationPid2);
 
 		ArgumentCaptor<SearchBuilderLoadIncludesParameters> searchBuilderLoadIncludesRequestDetailsCaptor = ArgumentCaptor.forClass(SearchBuilderLoadIncludesParameters.class);
 		verify(observationSearchBuilder).loadIncludes(searchBuilderLoadIncludesRequestDetailsCaptor.capture());
@@ -506,4 +508,90 @@ public void getResourcePidIterator_groupExportStyleWithGroupResource_returnsAnIt
 		validatePartitionId(thePartitioned, resourceDaoServletRequestDetailsCaptor.getValue().getRequestPartitionId());
 	}
 
+	@Test
+	void testGetExpandedPatientSetForPatientExport_withoutMdm_returnsOriginalPatients() {
+		// Given - Parameters with MDM disabled
+		ExportPIDIteratorParameters params = new ExportPIDIteratorParameters();
+		params.setPatientIds(List.of("Patient/1"));
+		params.setExpandMdm(false); // MDM disabled
+
+		// When - Call method
+		Set<String> result = myProcessor.getExpandedPatientSetForPatientExport(params);
+
+		// Then - Returns only original patient ID
+		assertThat(result).containsExactly("Patient/1");
+
+		// Verify MDM expander was NOT called
+		verify(myMdmExpandersHolder, never()).getBulkExportMDMResourceExpanderInstance();
+		verify(myBulkExportMDMResourceExpander, never()).expandPatient(any(), any());
+
+		// Verify results were cached
+		assertThat(params.getExpandedPatientIdsForPatientExport()).isEqualTo(result);
+		assertThat(params.hasExpandedPatientIdsForPatientExport()).isTrue();
+	}
+
+	@Test
+	void testGetExpandedPatientSetForPatientExport_withMdm_expandsAndCachesPatients() {
+		// Given - Parameters with MDM enabled
+		ExportPIDIteratorParameters params = new ExportPIDIteratorParameters();
+		params.setPatientIds(List.of("Patient/1"));
+		params.setExpandMdm(true); // MDM enabled
+		params.setPartitionId(RequestPartitionId.allPartitions());
+
+		// Mock MDM expander behavior
+		when(myMdmExpandersHolder.getBulkExportMDMResourceExpanderInstance())
+				.thenReturn(myBulkExportMDMResourceExpander);
+
+		// Patient/1 expands to Patient/1 + Patient/golden (linked via MDM)
+		when(myBulkExportMDMResourceExpander.expandPatient(eq("Patient/1"), any()))
+				.thenReturn(Set.of("Patient/1", "Patient/golden"));
+
+		// When - Call method
+		Set<String> result = myProcessor.getExpandedPatientSetForPatientExport(params);
+
+		// Then - Returns original + expanded patient IDs
+		assertThat(result).containsExactlyInAnyOrder("Patient/1", "Patient/golden");
+
+		// Verify MDM expander was called for the patient
+		verify(myMdmExpandersHolder, times(1)).getBulkExportMDMResourceExpanderInstance();
+		verify(myBulkExportMDMResourceExpander, times(1))
+				.expandPatient(eq("Patient/1"), eq(RequestPartitionId.allPartitions()));
+
+		// Verify results were cached
+		assertThat(params.getExpandedPatientIdsForPatientExport()).isEqualTo(result);
+		assertThat(params.hasExpandedPatientIdsForPatientExport()).isTrue();
+	}
+
+	@Test
+	void testGetExpandedPatientSetForPatientExport_cachePreventsRedundantExpansion() {
+		// Given - Parameters with MDM enabled
+		ExportPIDIteratorParameters params = new ExportPIDIteratorParameters();
+		params.setPatientIds(List.of("Patient/1"));
+		params.setExpandMdm(true);
+		params.setPartitionId(RequestPartitionId.allPartitions());
+
+		// Mock MDM expander to return expanded set
+		when(myMdmExpandersHolder.getBulkExportMDMResourceExpanderInstance())
+				.thenReturn(myBulkExportMDMResourceExpander);
+		when(myBulkExportMDMResourceExpander.expandPatient(eq("Patient/1"), any()))
+				.thenReturn(Set.of("Patient/1", "Patient/golden"));
+
+		// When - Call method FIRST time
+		Set<String> result1 = myProcessor.getExpandedPatientSetForPatientExport(params);
+
+		// Then - First call should expand and cache
+		assertThat(result1).containsExactlyInAnyOrder("Patient/1", "Patient/golden");
+		verify(myBulkExportMDMResourceExpander, times(1)).expandPatient(eq("Patient/1"), any());
+
+		// When - Call method SECOND time with SAME params
+		Set<String> result2 = myProcessor.getExpandedPatientSetForPatientExport(params);
+
+		// Then - Second call should use cache (no additional MDM calls)
+		assertThat(result2).containsExactlyInAnyOrder("Patient/1", "Patient/golden");
+
+		// Verify MDM expander was called only ONCE (on first call)
+		verify(myBulkExportMDMResourceExpander, times(1)).expandPatient(eq("Patient/1"), any());
+		verify(myMdmExpandersHolder, times(1)).getBulkExportMDMResourceExpanderInstance();
+	}
+
 }