Add partition selector to id parameter (#2208)

* Add partition selector to id parameter

* Add changelog

hapi-fhir-docs/src/main/resources/ca/uhn/hapi/fhir/changelog/5_3_0/2208-add-partition-selector-to-id-param.yaml

@@ -0,0 +1,6 @@
+---
+type: bug
+issue: 2208
+title: "When performing a JPA server search on a partitioned server, searches with only the `_id` parameter and no other
+   parameters did not include the partition selector in the generated SQL, resulting in leakage across partitions. Thanks to
+   GitHub user @jtheory for reporting!"

hapi-fhir-jpaserver-base/src/main/java/ca/uhn/fhir/jpa/search/builder/QueryStack.java

@@ -739,7 +739,7 @@ public class QueryStack {
 			codePredicates.add(singleCode);
 		}
 
-		return join.combineWithRequestPartitionIdPredicate(theRequestPartitionId, ComboCondition.or(codePredicates.toArray(new Condition[0])));
+		return join.combineWithRequestPartitionIdPredicate(theRequestPartitionId, toOrPredicate(codePredicates));
 	}
 
 	public Condition createPredicateTag(@Nullable DbColumn theSourceJoinColumn, List<List<IQueryParameterType>> theList, String theParamName, RequestPartitionId theRequestPartitionId) {

hapi-fhir-jpaserver-base/src/main/java/ca/uhn/fhir/jpa/search/builder/predicate/ResourceIdPredicateBuilder.java

@@ -115,12 +115,15 @@ public class ResourceIdPredicateBuilder extends BasePredicateBuilder {
 			List<Long> resourceIds = ResourcePersistentId.toLongList(allOrPids);
 			if (theSourceJoinColumn == null) {
 				BaseJoiningPredicateBuilder queryRootTable = super.getOrCreateQueryRootTable();
+				Condition predicate;
 				switch (operation) {
 					default:
 					case eq:
-						return queryRootTable.createPredicateResourceIds(false, resourceIds);
+						predicate = queryRootTable.createPredicateResourceIds(false, resourceIds);
+						return queryRootTable.combineWithRequestPartitionIdPredicate(theRequestPartitionId, predicate);
 					case ne:
-						return queryRootTable.createPredicateResourceIds(true, resourceIds);
+						predicate = queryRootTable.createPredicateResourceIds(true, resourceIds);
+						return queryRootTable.combineWithRequestPartitionIdPredicate(theRequestPartitionId, predicate);
 				}
 			} else {
 				return QueryStack.toEqualToOrInPredicate(theSourceJoinColumn, generatePlaceholders(resourceIds), operation == SearchFilterParser.CompareOperation.ne);

hapi-fhir-jpaserver-base/src/test/java/ca/uhn/fhir/jpa/dao/r4/PartitioningSqlR4Test.java

@@ -1136,6 +1136,231 @@ public class PartitioningSqlR4Test extends BaseJpaR4SystemTest {
 		}
 	}
 
+	@Test
+	public void testSearch_IdParamOnly_PidId_SpecificPartition() {
+		IIdType patientIdNull = createPatient(withPartition(null), withActiveTrue());
+		IIdType patientId1 = createPatient(withPartition(1), withActiveTrue());
+		IIdType patientId2 = createPatient(withPartition(2), withActiveTrue());
+
+		/* *******************************
+		 * _id param is only parameter
+		 * *******************************/
+
+		// Read in correct Partition
+		{
+			myCaptureQueriesListener.clear();
+			addReadPartition(1);
+
+			SearchParameterMap map = SearchParameterMap.newSynchronous(Patient.SP_RES_ID, new TokenParam(patientId1.toUnqualifiedVersionless().getValue()));
+			IBundleProvider searchOutcome = myPatientDao.search(map);
+			assertEquals(1, searchOutcome.size());
+			IIdType gotId1 = searchOutcome.getResources(0,1).get(0).getIdElement().toUnqualifiedVersionless();
+			assertEquals(patientId1, gotId1);
+
+			String searchSql = myCaptureQueriesListener.getSelectQueriesForCurrentThread().get(0).getSql(true, false);
+			ourLog.info("Search SQL:\n{}", searchSql);
+
+			// Only the read columns should be used, no criteria use partition
+			assertThat(searchSql, searchSql, containsString("PARTITION_ID IN ('1')"));
+			assertEquals(1, StringUtils.countMatches(searchSql, "PARTITION_ID"), searchSql);
+		}
+
+		// Read in null Partition
+		{
+			addReadPartition(1);
+
+			SearchParameterMap map = SearchParameterMap.newSynchronous(Patient.SP_RES_ID, new TokenParam(patientIdNull.toUnqualifiedVersionless().getValue()));
+			IBundleProvider searchOutcome = myPatientDao.search(map);
+			assertEquals(0, searchOutcome.size());
+		}
+
+		// Read in wrong Partition
+		{
+			addReadPartition(1);
+
+			SearchParameterMap map = SearchParameterMap.newSynchronous(Patient.SP_RES_ID, new TokenParam(patientId2.toUnqualifiedVersionless().getValue()));
+			IBundleProvider searchOutcome = myPatientDao.search(map);
+			assertEquals(0, searchOutcome.size());
+		}
+
+	}
+
+
+	@Test
+	public void testSearch_IdParamSecond_PidId_SpecificPartition() {
+		IIdType patientIdNull = createPatient(withPartition(null), withActiveTrue());
+		IIdType patientId1 = createPatient(withPartition(1), withActiveTrue());
+		IIdType patientId2 = createPatient(withPartition(2), withActiveTrue());
+
+		/* *******************************
+		 * _id param is second parameter
+		 * *******************************/
+
+		// Read in correct Partition
+		{
+			myCaptureQueriesListener.clear();
+			addReadPartition(1);
+
+			SearchParameterMap map = SearchParameterMap.newSynchronous()
+				.add(Patient.SP_ACTIVE, new TokenParam("true"))
+				.add(Patient.SP_RES_ID, new TokenParam(patientId1.toUnqualifiedVersionless().getValue()));
+			IBundleProvider searchOutcome = myPatientDao.search(map);
+			assertEquals(1, searchOutcome.size());
+			IIdType gotId1 = searchOutcome.getResources(0,1).get(0).getIdElement().toUnqualifiedVersionless();
+			assertEquals(patientId1, gotId1);
+
+			String searchSql = myCaptureQueriesListener.getSelectQueriesForCurrentThread().get(0).getSql(true, false);
+			ourLog.info("Search SQL:\n{}", searchSql);
+
+			// Only the read columns should be used, no criteria use partition
+			assertThat(searchSql, searchSql, containsString("PARTITION_ID IN ('1')"));
+			assertEquals(2, StringUtils.countMatches(searchSql, "PARTITION_ID"), searchSql); // If this switches to 1 that would be fine
+		}
+
+		// Read in null Partition
+		{
+			addReadPartition(1);
+
+			SearchParameterMap map = SearchParameterMap.newSynchronous()
+				.add(Patient.SP_ACTIVE, new TokenParam("true"))
+				.add(Patient.SP_RES_ID, new TokenParam(patientIdNull.toUnqualifiedVersionless().getValue()));
+			IBundleProvider searchOutcome = myPatientDao.search(map);
+			assertEquals(0, searchOutcome.size());
+		}
+
+		// Read in wrong Partition
+		{
+			addReadPartition(1);
+
+			SearchParameterMap map = SearchParameterMap.newSynchronous()
+				.add(Patient.SP_ACTIVE, new TokenParam("true"))
+				.add(Patient.SP_RES_ID, new TokenParam(patientId2.toUnqualifiedVersionless().getValue()));
+			IBundleProvider searchOutcome = myPatientDao.search(map);
+			assertEquals(0, searchOutcome.size());
+		}
+
+	}
+
+
+	@Test
+	public void testSearch_IdParamOnly_ForcedId_SpecificPartition() {
+		addReadPartition(new Integer[]{null});
+		IIdType patientIdNull = createPatient(withPartition(null), withId("PT-NULL"), withActiveTrue());
+		addReadPartition(1);
+		IIdType patientId1 = createPatient(withPartition(1), withId("PT-1"), withActiveTrue());
+		addReadPartition(2);
+		IIdType patientId2 = createPatient(withPartition(2), withId("PT-2"), withActiveTrue());
+
+		/* *******************************
+		 * _id param is only parameter
+		 * *******************************/
+
+		// Read in correct Partition
+		{
+			myCaptureQueriesListener.clear();
+			addReadPartition(1);
+
+			SearchParameterMap map = SearchParameterMap.newSynchronous(Patient.SP_RES_ID, new TokenParam(patientId1.toUnqualifiedVersionless().getValue()));
+			IBundleProvider searchOutcome = myPatientDao.search(map);
+			assertEquals(1, searchOutcome.size());
+			IIdType gotId1 = searchOutcome.getResources(0,1).get(0).getIdElement().toUnqualifiedVersionless();
+			assertEquals(patientId1, gotId1);
+
+			String searchSql = myCaptureQueriesListener.getSelectQueriesForCurrentThread().get(0).getSql(true, false).toUpperCase();
+			ourLog.info("Search SQL:\n{}", searchSql);
+
+			// Only the read columns should be used, no criteria use partition
+			assertThat(searchSql, searchSql, containsString("PARTITION_ID IN ('1')"));
+			assertEquals(1, StringUtils.countMatches(searchSql, "PARTITION_ID"), searchSql);
+		}
+
+		// Read in null Partition
+		{
+			addReadPartition(1);
+
+			SearchParameterMap map = SearchParameterMap.newSynchronous(Patient.SP_RES_ID, new TokenParam(patientIdNull.toUnqualifiedVersionless().getValue()));
+			IBundleProvider searchOutcome = myPatientDao.search(map);
+			assertEquals(0, searchOutcome.size());
+		}
+
+		// Read in wrong Partition
+		{
+			addReadPartition(1);
+
+			SearchParameterMap map = SearchParameterMap.newSynchronous(Patient.SP_RES_ID, new TokenParam(patientId2.toUnqualifiedVersionless().getValue()));
+			IBundleProvider searchOutcome = myPatientDao.search(map);
+			assertEquals(0, searchOutcome.size());
+		}
+
+	}
+
+
+	@Test
+	public void testSearch_IdParamSecond_ForcedId_SpecificPartition() {
+		addReadPartition(new Integer[]{null});
+		IIdType patientIdNull = createPatient(withPartition(null), withId("PT-NULL"), withActiveTrue());
+		addReadPartition(1);
+		IIdType patientId1 = createPatient(withPartition(1), withId("PT-1"), withActiveTrue());
+		addReadPartition(2);
+		IIdType patientId2 = createPatient(withPartition(2), withId("PT-2"), withActiveTrue());
+
+		/* *******************************
+		 * _id param is second parameter
+		 * *******************************/
+
+		// Read in correct Partition
+		{
+			myCaptureQueriesListener.clear();
+			addReadPartition(1);
+
+			SearchParameterMap map = SearchParameterMap.newSynchronous()
+				.add(Patient.SP_ACTIVE, new TokenParam("true"))
+				.add(Patient.SP_RES_ID, new TokenParam(patientId1.toUnqualifiedVersionless().getValue()));
+			IBundleProvider searchOutcome = myPatientDao.search(map);
+			assertEquals(1, searchOutcome.size());
+			IIdType gotId1 = searchOutcome.getResources(0,1).get(0).getIdElement().toUnqualifiedVersionless();
+			assertEquals(patientId1, gotId1);
+
+			// First SQL resolves the forced ID
+			String searchSql = myCaptureQueriesListener.getSelectQueriesForCurrentThread().get(0).getSql(true, false).toUpperCase();
+			ourLog.info("Search SQL:\n{}", searchSql);
+			assertThat(searchSql, searchSql, containsString("PARTITION_ID IN ('1')"));
+			assertEquals(1, StringUtils.countMatches(searchSql, "PARTITION_ID"), searchSql);
+
+			// Second SQL performs the search
+			searchSql = myCaptureQueriesListener.getSelectQueriesForCurrentThread().get(1).getSql(true, false).toUpperCase();
+			ourLog.info("Search SQL:\n{}", searchSql);
+			assertThat(searchSql, searchSql, containsString("PARTITION_ID IN ('1')"));
+			assertEquals(2, StringUtils.countMatches(searchSql, "PARTITION_ID"), searchSql); // If this switches to 1 that would be fine
+		}
+
+		// Read in null Partition
+		{
+			addReadPartition(1);
+
+			SearchParameterMap map = SearchParameterMap.newSynchronous()
+				.add(Patient.SP_ACTIVE, new TokenParam("true"))
+				.add(Patient.SP_RES_ID, new TokenParam(patientIdNull.toUnqualifiedVersionless().getValue()));
+			IBundleProvider searchOutcome = myPatientDao.search(map);
+			assertEquals(0, searchOutcome.size());
+		}
+
+		// Read in wrong Partition
+		{
+			addReadPartition(1);
+
+			SearchParameterMap map = SearchParameterMap.newSynchronous()
+				.add(Patient.SP_ACTIVE, new TokenParam("true"))
+				.add(Patient.SP_RES_ID, new TokenParam(patientId2.toUnqualifiedVersionless().getValue()));
+			IBundleProvider searchOutcome = myPatientDao.search(map);
+			assertEquals(0, searchOutcome.size());
+		}
+
+	}
+
+
+
+
 	@Test
 	public void testSearch_MissingParamString_SearchAllPartitions() {
 		myPartitionSettings.setIncludePartitionInSearchHashes(false);

hapi-fhir-test-utilities/src/main/java/ca/uhn/fhir/test/utilities/ITestDataBuilder.java

@@ -115,7 +115,7 @@ public interface ITestDataBuilder {
 
 	default Consumer<IBaseResource> withId(String theId) {
 		return t -> {
-			assertThat(theId, matchesPattern("[a-zA-Z0-9]+"));
+			assertThat(theId, matchesPattern("[a-zA-Z0-9-]+"));
 			t.setId(theId);
 		};
 	}