Improve partition support for Bulk Patch operation

hapi-fhir-base/src/main/java/ca/uhn/fhir/interceptor/model/RequestPartitionId.java

@@ -45,7 +45,7 @@
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
-import static org.apache.commons.lang3.ObjectUtils.defaultIfNull;
+import static org.apache.commons.lang3.ObjectUtils.getIfNull;
 
 /**
  * @since 5.0.0
@@ -72,10 +72,7 @@ public class RequestPartitionId implements IModelJson {
 	 */
 	private RequestPartitionId(
 			@Nullable String thePartitionName, @Nullable Integer thePartitionId, @Nullable LocalDate thePartitionDate) {
-		myPartitionIds = toListOrNull(thePartitionId);
-		myPartitionNames = toListOrNull(thePartitionName);
-		myPartitionDate = thePartitionDate;
-		myAllPartitions = false;
+		this(toListOrNull(thePartitionName), toListOrNull(thePartitionId), thePartitionDate);
 	}
 
 	/**
@@ -85,10 +82,21 @@ private RequestPartitionId(
 			@Nullable List<String> thePartitionName,
 			@Nullable List<Integer> thePartitionId,
 			@Nullable LocalDate thePartitionDate) {
+		this(thePartitionName, thePartitionId, thePartitionDate, false);
+	}
+
+	/**
+	 * Constructor for a multiple partitions with explicit "all partitions" flag
+	 */
+	private RequestPartitionId(
+			@Nullable List<String> thePartitionName,
+			@Nullable List<Integer> thePartitionId,
+			@Nullable LocalDate thePartitionDate,
+			boolean theAllPartitions) {
 		myPartitionIds = toListOrNull(thePartitionId);
 		myPartitionNames = toListOrNull(thePartitionName);
 		myPartitionDate = thePartitionDate;
-		myAllPartitions = false;
+		myAllPartitions = theAllPartitions;
 	}
 
 	/**
@@ -117,7 +125,7 @@ public static Optional<RequestPartitionId> getPartitionIfAssigned(IBaseResource
 	 * @since 7.4.0
 	 */
 	public RequestPartitionId mergeIds(RequestPartitionId theOther) {
-		if (isAllPartitions() || theOther.isAllPartitions()) {
+		if ((isAllPartitions() && !hasPartitionIds()) || (theOther.isAllPartitions() && !theOther.hasPartitionIds())) {
 			return RequestPartitionId.allPartitions();
 		}
 
@@ -131,7 +139,12 @@ public RequestPartitionId mergeIds(RequestPartitionId theOther) {
 		List<Integer> newPartitionIds = Stream.concat(thisPartitionIds.stream(), otherPartitionIds.stream())
 				.distinct()
 				.collect(Collectors.toList());
-		return RequestPartitionId.fromPartitionIds(newPartitionIds);
+		boolean newAllPartitions = isAllPartitions() || theOther.isAllPartitions();
+		if (newAllPartitions) {
+			return RequestPartitionId.allPartitionsWithPartitionIds(newPartitionIds);
+		} else {
+			return RequestPartitionId.fromPartitionIds(newPartitionIds);
+		}
 	}
 
 	public static RequestPartitionId fromJson(String theJson) throws JsonProcessingException {
@@ -260,7 +273,7 @@ public boolean isDefaultPartition() {
 	 *         <code>thePartitionId</code>.
 	 */
 	public boolean isPartition(@Nullable Integer thePartitionId) {
-		if (isAllPartitions()) {
+		if (isAllPartitions() && !hasPartitionIds()) {
 			return false;
 		}
 		return hasPartitionIds()
@@ -348,6 +361,26 @@ public static RequestPartitionId allPartitions() {
 		return ALL_PARTITIONS;
 	}
 
+	/**
+	 * Creates a new RequestPartitionId which indicates "all partitions" and explicitly lists them
+	 *
+	 * @since 8.8.0
+	 */
+	@Nonnull
+	public static RequestPartitionId allPartitionsWithPartitionIds(Integer... thePartitionIds) {
+		return allPartitionsWithPartitionIds(toListOrNull(thePartitionIds));
+	}
+
+	/**
+	 * Creates a new RequestPartitionId which indicates "all partitions" and explicitly lists them
+	 *
+	 * @since 8.8.0
+	 */
+	@Nonnull
+	public static RequestPartitionId allPartitionsWithPartitionIds(List<Integer> thePartitionIds) {
+		return new RequestPartitionId(null, thePartitionIds, null, true);
+	}
+
 	/**
 	 * @deprecated use {@link RequestPartitionId#defaultPartition(IDefaultPartitionSettings)} instead
 	 */
@@ -443,18 +476,32 @@ public static RequestPartitionId forPartitionIdsAndNames(
 		return new RequestPartitionId(thePartitionNames, thePartitionIds, thePartitionDate);
 	}
 
+	@Nonnull
+	public static RequestPartitionId forPartitionIdsAndNames(
+			List<String> thePartitionNames,
+			List<Integer> thePartitionIds,
+			LocalDate thePartitionDate,
+			boolean theAllPartitions) {
+		return new RequestPartitionId(thePartitionNames, thePartitionIds, thePartitionDate, theAllPartitions);
+	}
+
 	/**
 	 * Create a string representation suitable for use as a cache key. Null aware.
 	 * <p>
 	 * Returns the partition IDs (numeric) as a joined string with a space between, using the string "null" for any null values
 	 */
 	public static String stringifyForKey(@Nonnull RequestPartitionId theRequestPartitionId) {
-		String retVal = "(all)";
-		if (!theRequestPartitionId.isAllPartitions()) {
+		String retVal;
+		if (theRequestPartitionId.hasPartitionIds()) {
 			assert theRequestPartitionId.hasPartitionIds();
 			retVal = theRequestPartitionId.getPartitionIds().stream()
-					.map(t -> defaultIfNull(t, "null").toString())
+					.map(t -> getIfNull(t, "null").toString())
 					.collect(Collectors.joining(" "));
+			if (theRequestPartitionId.isAllPartitions()) {
+				retVal = "(all) " + retVal;
+			}
+		} else {
+			retVal = "(all)";
 		}
 		return retVal;
 	}

hapi-fhir-base/src/test/java/ca/uhn/fhir/interceptor/model/RequestPartitionIdTest.java

@@ -9,6 +9,8 @@
 import org.slf4j.LoggerFactory;
 
 import java.time.LocalDate;
+import java.util.List;
+import java.util.stream.Stream;
 
 import static ca.uhn.fhir.interceptor.model.RequestPartitionId.allPartitions;
 import static ca.uhn.fhir.interceptor.model.RequestPartitionId.defaultPartition;
@@ -118,6 +120,24 @@ public void testMergeIds_IncludesDefault() {
 
 	}
 
+	@ParameterizedTest
+	@MethodSource("testStringifyForKeyTestCases")
+	public void testStringifyForKey(RequestPartitionId theRequestPartitionId, String theExpectedString) {
+		String actual = RequestPartitionId.stringifyForKey(theRequestPartitionId);
+		assertEquals(theExpectedString, actual);
+	}
+
+
+	static Stream<Object[]> testStringifyForKeyTestCases() {
+		return List.of(
+			new Object[]{RequestPartitionId.allPartitions(), "(all)"},
+			new Object[]{RequestPartitionId.defaultPartition(), "null"},
+			new Object[]{RequestPartitionId.fromPartitionIds(1, 2, 3), "1 2 3"},
+			new Object[]{RequestPartitionId.fromPartitionIds(null, 2, 3), "null 2 3"},
+			new Object[]{RequestPartitionId.allPartitionsWithPartitionIds(1, 2, 3), "(all) 1 2 3"}
+		).stream();
+	}
+
 	record ContainsTestCase(String description, RequestPartitionId left, RequestPartitionId right, Comparison comparison) {
 		enum Comparison {
 			LEFT_CONTAINS_RIGHT,

hapi-fhir-docs/src/main/resources/ca/uhn/hapi/fhir/changelog/8_8_0/7406-add-partition-param-to-bulk-patch.yaml

@@ -0,0 +1,8 @@
+---
+type: add
+issue: 7406
+title: "A new parameter has been added to the `$hapi.fhir.bulk-patch` and
+   `$hapi.fhir.bulk-patch-rewrite-history` operations which can be used to
+   explicitly specify the partition(s) to use when applying these
+   operations. This change also generally improves support for using these
+   jobs in a partitioned environment."

hapi-fhir-jpaserver-base/src/main/java/ca/uhn/fhir/jpa/batch2/Batch2DaoSvcImpl.java

@@ -171,12 +171,13 @@ private Stream<TypedResourcePid> searchForResourceIdsAndType(
 		ISearchBuilder<JpaPid> builder = mySearchBuilderFactory.newSearchBuilder(null, null);
 		return myTransactionService
 				.withRequest(theRequestDetails)
-				.search(() -> builder.createQueryStream(
+				.withRequestPartitionId(theRequestPartitionId)
+				.search(partition -> builder.createQueryStream(
 						theSearchParams,
 						new SearchRuntimeDetails(
 								theRequestDetails, UUID.randomUUID().toString()),
 						theRequestDetails,
-						theRequestPartitionId))
+						partition))
 				.map(pid -> new TypedResourcePid(pid.getResourceType(), pid));
 	}
 

hapi-fhir-jpaserver-base/src/main/java/ca/uhn/fhir/jpa/cache/ResourceVersionSvcDaoImpl.java

@@ -27,6 +27,7 @@
 import ca.uhn.fhir.jpa.dao.data.IResourceTableDao;
 import ca.uhn.fhir.jpa.model.cross.IResourceLookup;
 import ca.uhn.fhir.jpa.model.dao.JpaPid;
+import ca.uhn.fhir.jpa.partition.IRequestPartitionHelperSvc;
 import ca.uhn.fhir.jpa.searchparam.SearchParameterMap;
 import ca.uhn.fhir.rest.api.server.SystemRequestDetails;
 import jakarta.annotation.Nonnull;
@@ -60,6 +61,9 @@ public class ResourceVersionSvcDaoImpl implements IResourceVersionSvc {
 	@Autowired
 	IIdHelperService<JpaPid> myIdHelperService;
 
+	@Autowired
+	IRequestPartitionHelperSvc myRequestPartitionHelperSvc;
+
 	@Override
 	@Nonnull
 	public ResourceVersionMap getVersionMap(
@@ -76,6 +80,13 @@ public ResourceVersionMap getVersionMap(
 		return ResourceVersionMap.fromIdsWithVersions(fhirIds);
 	}
 
+	@Override
+	public ResourceVersionMap getVersionMap(String theResourceName, SearchParameterMap theSearchParamMap) {
+		RequestPartitionId requestPartition = myRequestPartitionHelperSvc.determineReadPartitionForRequestForSearchType(
+				null, theResourceName, theSearchParamMap);
+		return getVersionMap(requestPartition, theResourceName, theSearchParamMap);
+	}
+
 	/**
 	 * Retrieves the latest versions for any resourceid that are found.
 	 * If they are not found, they will not be contained in the returned map.

hapi-fhir-jpaserver-base/src/main/java/ca/uhn/fhir/jpa/dao/BaseHapiFhirResourceDao.java

@@ -934,23 +934,32 @@ public DeleteMethodOutcome deleteByUrl(
 			@Nonnull TransactionDetails theTransactionDetails) {
 		validateDeleteEnabled();
 
+		ResourceSearch resourceSearch = myMatchUrlService.getResourceSearch(theUrl);
+		SearchParameterMap paramMap = resourceSearch.getSearchParameterMap();
+
+		RequestPartitionId requestPartitionId =
+				myRequestPartitionHelperService.determineReadPartitionForRequestForSearchType(
+						theRequestDetails, myResourceName, paramMap);
+
 		return myTransactionService
 				.withRequest(theRequestDetails)
+				.withRequestPartitionId(requestPartitionId)
 				.withTransactionDetails(theTransactionDetails)
-				.execute(tx -> doDeleteByUrl(theUrl, deleteConflicts, theTransactionDetails, theRequestDetails));
+				.execute(tx ->
+						doDeleteByUrl(theUrl, paramMap, deleteConflicts, theTransactionDetails, theRequestDetails));
 	}
 
 	@Nonnull
 	private DeleteMethodOutcome doDeleteByUrl(
 			String theUrl,
+			SearchParameterMap theParamMap,
 			DeleteConflictList deleteConflicts,
 			TransactionDetails theTransactionDetails,
 			RequestDetails theRequestDetails) {
-		ResourceSearch resourceSearch = myMatchUrlService.getResourceSearch(theUrl);
-		SearchParameterMap paramMap = resourceSearch.getSearchParameterMap();
-		paramMap.setLoadSynchronous(true);
 
-		Set<JpaPid> resourceIds = myMatchResourceUrlService.search(paramMap, myResourceType, theRequestDetails, null);
+		theParamMap.setLoadSynchronous(true);
+		Set<JpaPid> resourceIds =
+				myMatchResourceUrlService.search(theParamMap, myResourceType, theRequestDetails, null);
 
 		if (resourceIds.size() > 1) {
 			if (!getStorageSettings().isAllowMultipleDelete()) {
@@ -1663,7 +1672,7 @@ public T read(IIdType theId, RequestDetails theRequest, boolean theDeletedOk) {
 				.withRequest(theRequest)
 				.withTransactionDetails(transactionDetails)
 				.withRequestPartitionId(requestPartitionId)
-				.read(() -> doReadInTransaction(theId, theRequest, theDeletedOk, requestPartitionId));
+				.read(partition -> doReadInTransaction(theId, theRequest, theDeletedOk, partition));
 	}
 
 	private T doReadInTransaction(
@@ -2278,7 +2287,7 @@ public List<JpaPid> searchForIds(
 				.withRequest(theRequest)
 				.withTransactionDetails(transactionDetails)
 				.withRequestPartitionId(requestPartitionId)
-				.searchList(() -> {
+				.searchList(partition -> {
 					if (isNull(theParams.getLoadSynchronousUpTo())) {
 						theParams.setLoadSynchronousUpTo(myStorageSettings.getInternalSynchronousSearchSize());
 					}
@@ -2292,7 +2301,7 @@ public List<JpaPid> searchForIds(
 
 					SearchRuntimeDetails searchRuntimeDetails = new SearchRuntimeDetails(theRequest, uuid);
 					try (IResultIterator<JpaPid> iter =
-							builder.createQuery(theParams, searchRuntimeDetails, theRequest, requestPartitionId)) {
+							builder.createQuery(theParams, searchRuntimeDetails, theRequest, partition)) {
 						while (iter.hasNext()) {
 							ids.add(iter.next());
 						}
@@ -2325,8 +2334,8 @@ public <PID extends IResourcePersistentId<?>> Stream<PID> searchForIdStream(
 		//noinspection unchecked
 		return (Stream<PID>) myTransactionService
 				.withRequest(theRequest)
-				.search(() ->
-						builder.createQueryStream(theParams, searchRuntimeDetails, theRequest, requestPartitionId));
+				.withRequestPartitionId(requestPartitionId)
+				.search(partition -> builder.createQueryStream(theParams, searchRuntimeDetails, theRequest, partition));
 	}
 
 	@Override
@@ -2353,13 +2362,14 @@ private <V> List<V> searchForTransformedIds(
 		return myTransactionService
 				.withRequest(theRequest)
 				.withPropagation(Propagation.REQUIRED)
-				.searchList(() -> {
+				.withRequestPartitionId(requestPartitionId)
+				.searchList(partition -> {
 					ISearchBuilder<JpaPid> builder =
 							mySearchBuilderFactory.newSearchBuilder(getResourceName(), getResourceType());
 					Stream<JpaPid> pidStream =
-							builder.createQueryStream(theParams, searchRuntimeDetails, theRequest, requestPartitionId);
+							builder.createQueryStream(theParams, searchRuntimeDetails, theRequest, partition);
 
-					Stream<V> transformedStream = transform.apply(theRequest, pidStream, requestPartitionId);
+					Stream<V> transformedStream = transform.apply(theRequest, pidStream, partition);
 
 					return transformedStream.collect(Collectors.toList());
 				});

hapi-fhir-jpaserver-base/src/main/java/ca/uhn/fhir/jpa/partition/RequestPartitionHelperSvc.java

@@ -42,15 +42,16 @@ public RequestPartitionHelperSvc() {}
 	public RequestPartitionId validateAndNormalizePartitionIds(RequestPartitionId theRequestPartitionId) {
 		List<String> names = null;
 		List<Integer> partitionIds = null;
-		for (int i = 0; i < theRequestPartitionId.getPartitionIds().size(); i++) {
+		List<Integer> originalPartitionIds = theRequestPartitionId.getPartitionIds();
+		for (int i = 0; i < originalPartitionIds.size(); i++) {
 
 			PartitionEntity partition;
-			Integer id = theRequestPartitionId.getPartitionIds().get(i);
+			Integer id = originalPartitionIds.get(i);
 			if (id == null) {
 				partition = null;
 				if (myPartitionSettings.getDefaultPartitionId() != null) {
 					if (partitionIds == null) {
-						partitionIds = new ArrayList<>(theRequestPartitionId.getPartitionIds());
+						partitionIds = new ArrayList<>(originalPartitionIds);
 					}
 					partitionIds.set(i, myPartitionSettings.getDefaultPartitionId());
 				}
@@ -63,15 +64,15 @@ public RequestPartitionId validateAndNormalizePartitionIds(RequestPartitionId th
 							.getMessage(
 									BaseRequestPartitionHelperSvc.class,
 									"unknownPartitionId",
-									theRequestPartitionId.getPartitionIds().get(i));
+									originalPartitionIds.get(i));
 					throw new ResourceNotFoundException(Msg.code(1316) + msg);
 				}
 			}
 
 			if (theRequestPartitionId.hasPartitionNames()) {
 				if (partition == null) {
 					Validate.isTrue(
-							theRequestPartitionId.getPartitionIds().get(i) == null,
+							originalPartitionIds.get(i) == null,
 							"Partition %s must not have an ID",
 							JpaConstants.DEFAULT_PARTITION_NAME);
 				} else {
@@ -80,7 +81,7 @@ public RequestPartitionId validateAndNormalizePartitionIds(RequestPartitionId th
 									theRequestPartitionId.getPartitionNames().get(i), partition.getName()),
 							"Partition name %s does not match ID %s",
 							theRequestPartitionId.getPartitionNames().get(i),
-							theRequestPartitionId.getPartitionIds().get(i));
+							originalPartitionIds.get(i));
 				}
 			} else {
 				if (names == null) {
@@ -95,12 +96,15 @@ public RequestPartitionId validateAndNormalizePartitionIds(RequestPartitionId th
 		}
 
 		if (names != null) {
-			List<Integer> partitionIdsToUse = theRequestPartitionId.getPartitionIds();
+			List<Integer> partitionIdsToUse = originalPartitionIds;
 			if (partitionIds != null) {
 				partitionIdsToUse = partitionIds;
 			}
 			return RequestPartitionId.forPartitionIdsAndNames(
-					names, partitionIdsToUse, theRequestPartitionId.getPartitionDate());
+					names,
+					partitionIdsToUse,
+					theRequestPartitionId.getPartitionDate(),
+					theRequestPartitionId.isAllPartitions());
 		}
 
 		return theRequestPartitionId;

hapi-fhir-jpaserver-model/src/main/java/ca/uhn/fhir/jpa/model/config/PartitionSettings.java

@@ -138,7 +138,7 @@ public PartitionSettings setPartitioningEnabled(boolean thePartitioningEnabled)
 	}
 
 	/**
-	 * Should resources references be permitted to cross partition boundaries. Default is {@link CrossPartitionReferenceMode#NOT_ALLOWED}.
+	 * Should resource references be permitted to cross partition boundaries? Default is {@link CrossPartitionReferenceMode#NOT_ALLOWED}.
 	 *
 	 * @since 5.0.0
 	 */