rebuilding golden resources with survivorship in partitions (#5957)

* fixing a regression that prevented rebuilding golden resources with survivorship in a partitioned system

* spotless

* update var usage

---------

Co-authored-by: leif stawnyczy <leifstawnyczy@leifs-mbp.home>

hapi-fhir-docs/src/main/resources/ca/uhn/hapi/fhir/changelog/7_4_0/5956-rebuild-golden-resources-with-survivorship-on-partitions.yaml

@@ -0,0 +1,7 @@
+---
+type: fix
+issue: 5956
+title: "Fixed a regression that prevented rebuilding golden resources with mdm
+        survivorship rules when resources have non-numeric fhir ids in
+        a partitioned system.
+"

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

@@ -155,25 +155,26 @@ public class IdHelperService implements IIdHelperService<JpaPid> {
 	public IResourceLookup<JpaPid> resolveResourceIdentity(
 			@Nonnull RequestPartitionId theRequestPartitionId,
 			String theResourceType,
-			String theResourceId,
+			final String theResourceId,
 			boolean theExcludeDeleted)
 			throws ResourceNotFoundException {
 		assert myDontCheckActiveTransactionForUnitTest || TransactionSynchronizationManager.isSynchronizationActive()
 				: "no transaction active";
 
-		if (theResourceId.contains("/")) {
+		String resourceIdToUse = theResourceId;
-			theResourceId = theResourceId.substring(theResourceId.indexOf("/") + 1);
+		if (resourceIdToUse.contains("/")) {
+			resourceIdToUse = theResourceId.substring(resourceIdToUse.indexOf("/") + 1);
 		}
-		IdDt id = new IdDt(theResourceType, theResourceId);
+		IdDt id = new IdDt(theResourceType, resourceIdToUse);
 		Map<String, List<IResourceLookup<JpaPid>>> matches =
 				translateForcedIdToPids(theRequestPartitionId, Collections.singletonList(id), theExcludeDeleted);
 
 		// We only pass 1 input in so only 0..1 will come back
-		if (matches.isEmpty() || !matches.containsKey(theResourceId)) {
+		if (matches.isEmpty() || !matches.containsKey(resourceIdToUse)) {
 			throw new ResourceNotFoundException(Msg.code(2001) + "Resource " + id + " is not known");
 		}
 
-		if (matches.size() > 1 || matches.get(theResourceId).size() > 1) {
+		if (matches.size() > 1 || matches.get(resourceIdToUse).size() > 1) {
 			/*
 			 *  This means that:
 			 *  1. There are two resources with the exact same resource type and forced id
@@ -183,7 +184,7 @@ public class IdHelperService implements IIdHelperService<JpaPid> {
 			throw new PreconditionFailedException(Msg.code(1099) + msg);
 		}
 
-		return matches.get(theResourceId).get(0);
+		return matches.get(resourceIdToUse).get(0);
 	}
 
 	/**

hapi-fhir-jpaserver-mdm/src/main/java/ca/uhn/fhir/jpa/mdm/svc/MdmModelConverterSvcImpl.java

@@ -31,9 +31,11 @@ import java.math.RoundingMode;
 
 public class MdmModelConverterSvcImpl implements IMdmModelConverterSvc {
 
+	@SuppressWarnings("rawtypes")
 	@Autowired
 	IIdHelperService myIdHelperService;
 
+	@SuppressWarnings({"rawtypes", "unchecked"})
 	@Override
 	public MdmLinkJson toJson(IMdmLink theLink) {
 		MdmLinkJson retVal = new MdmLinkJson();
@@ -42,11 +44,13 @@ public class MdmModelConverterSvcImpl implements IMdmModelConverterSvc {
 				.toVersionless()
 				.getValue();
 		retVal.setSourceId(sourceId);
+		retVal.setSourcePid(theLink.getSourcePersistenceId());
 		String goldenResourceId = myIdHelperService
 				.resourceIdFromPidOrThrowException(theLink.getGoldenResourcePersistenceId(), theLink.getMdmSourceType())
 				.toVersionless()
 				.getValue();
 		retVal.setGoldenResourceId(goldenResourceId);
+		retVal.setGoldenPid(theLink.getGoldenResourcePersistenceId());
 		retVal.setCreated(theLink.getCreated());
 		retVal.setEidMatch(theLink.getEidMatch());
 		retVal.setLinkSource(theLink.getLinkSource());

hapi-fhir-jpaserver-mdm/src/test/java/ca/uhn/fhir/jpa/mdm/svc/MdmModelConverterSvcImplTest.java

@@ -2,13 +2,14 @@ package ca.uhn.fhir.jpa.mdm.svc;
 
 import ca.uhn.fhir.jpa.entity.MdmLink;
 import ca.uhn.fhir.jpa.mdm.BaseMdmR4Test;
+import ca.uhn.fhir.jpa.model.dao.JpaPid;
 import ca.uhn.fhir.jpa.model.entity.EnversRevision;
 import ca.uhn.fhir.mdm.api.IMdmLink;
-import ca.uhn.fhir.mdm.model.mdmevents.MdmLinkJson;
-import ca.uhn.fhir.mdm.model.mdmevents.MdmLinkWithRevisionJson;
 import ca.uhn.fhir.mdm.api.MdmLinkSourceEnum;
 import ca.uhn.fhir.mdm.api.MdmLinkWithRevision;
 import ca.uhn.fhir.mdm.api.MdmMatchResultEnum;
+import ca.uhn.fhir.mdm.model.mdmevents.MdmLinkJson;
+import ca.uhn.fhir.mdm.model.mdmevents.MdmLinkWithRevisionJson;
 import ca.uhn.fhir.rest.api.server.storage.IResourcePersistentId;
 import org.hibernate.envers.RevisionType;
 import org.junit.jupiter.api.Test;
@@ -16,8 +17,6 @@ import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.springframework.beans.factory.annotation.Autowired;
 
-import static org.junit.jupiter.api.Assertions.assertEquals;
-
 import java.math.BigDecimal;
 import java.math.RoundingMode;
 import java.time.LocalDateTime;
@@ -25,6 +24,8 @@ import java.time.Month;
 import java.time.ZoneId;
 import java.util.Date;
 
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
 public class MdmModelConverterSvcImplTest extends BaseMdmR4Test {
 	private static final Logger ourLog = LoggerFactory.getLogger(MdmModelConverterSvcImplTest.class);
 
@@ -49,7 +50,8 @@ public class MdmModelConverterSvcImplTest extends BaseMdmR4Test {
 
 		ourLog.info("actualMdmLinkJson: {}", actualMdmLinkJson);
 
-		assertEquals(getExepctedMdmLinkJson(mdmLink.getGoldenResourcePersistenceId().getId(), mdmLink.getSourcePersistenceId().getId(), MdmMatchResultEnum.MATCH, MdmLinkSourceEnum.MANUAL, version, createTime, updateTime, isLinkCreatedResource, scoreRounded), actualMdmLinkJson);
+		MdmLinkJson exepctedMdmLinkJson = getExepctedMdmLinkJson(mdmLink.getGoldenResourcePersistenceId().getId(), mdmLink.getSourcePersistenceId().getId(), MdmMatchResultEnum.MATCH, MdmLinkSourceEnum.MANUAL, version, createTime, updateTime, isLinkCreatedResource, scoreRounded);
+		assertEquals(exepctedMdmLinkJson, actualMdmLinkJson);
 	}
 
 	@Test
@@ -99,7 +101,9 @@ public class MdmModelConverterSvcImplTest extends BaseMdmR4Test {
 		final MdmLinkJson mdmLinkJson = new MdmLinkJson();
 
 		mdmLinkJson.setGoldenResourceId("Patient/" + theGoldenPatientId);
+		mdmLinkJson.setGoldenPid(JpaPid.fromId(theGoldenPatientId));
 		mdmLinkJson.setSourceId("Patient/" + theSourceId);
+		mdmLinkJson.setSourcePid(JpaPid.fromId(theSourceId));
 		mdmLinkJson.setMatchResult(theMdmMatchResultEnum);
 		mdmLinkJson.setLinkSource(theMdmLinkSourceEnum);
 		mdmLinkJson.setVersion(version);

hapi-fhir-jpaserver-mdm/src/test/java/ca/uhn/fhir/jpa/mdm/svc/MdmSurvivorshipSvcImplIT.java

@@ -1,17 +1,22 @@
 package ca.uhn.fhir.jpa.mdm.svc;
 
+import ca.uhn.fhir.interceptor.model.RequestPartitionId;
+import ca.uhn.fhir.jpa.entity.PartitionEntity;
 import ca.uhn.fhir.jpa.mdm.BaseMdmR4Test;
+import ca.uhn.fhir.jpa.model.entity.StorageSettings;
+import ca.uhn.fhir.jpa.searchparam.SearchParameterMap;
 import ca.uhn.fhir.mdm.api.IMdmSurvivorshipService;
 import ca.uhn.fhir.mdm.api.MdmLinkSourceEnum;
 import ca.uhn.fhir.mdm.api.MdmMatchOutcome;
 import ca.uhn.fhir.mdm.model.MdmTransactionContext;
+import ca.uhn.fhir.rest.api.server.IBundleProvider;
 import ca.uhn.fhir.rest.api.server.SystemRequestDetails;
+import ca.uhn.fhir.rest.api.server.storage.TransactionDetails;
+import ca.uhn.fhir.rest.server.interceptor.partition.RequestTenantPartitionInterceptor;
 import org.hl7.fhir.r4.model.Patient;
 import org.junit.jupiter.api.Test;
 import org.springframework.beans.factory.annotation.Autowired;
 
-import java.util.List;
-
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertNull;
@@ -77,4 +82,104 @@ class MdmSurvivorshipSvcImplIT extends BaseMdmR4Test {
 
 		myMdmSurvivorshipService.rebuildGoldenResourceWithSurvivorshipRules(goldenPatient, new MdmTransactionContext(MdmTransactionContext.OperationType.UPDATE_LINK));
 	}
+
+	@Test
+	public void rebuildGoldenResourceWithSurvivorshipRules_usingPatchedResourcesAndPartitions_doesNotThrow() {
+		// setup
+		boolean isPartitioningEnabled = myPartitionSettings.isPartitioningEnabled();
+		boolean isUnnamedPartitionMode = myPartitionSettings.isUnnamedPartitionMode();
+		StorageSettings.IndexEnabledEnum indexMissingFields = myStorageSettings.getIndexMissingFields();
+		boolean isSearchAllPartitionForMatch = myMdmSettings.getSearchAllPartitionForMatch();
+		String goldenResourcePartitionName = myMdmSettings.getGoldenResourcePartitionName();
+
+		myPartitionSettings.setPartitioningEnabled(true);
+		myPartitionSettings.setUnnamedPartitionMode(false);
+		myStorageSettings.setIndexMissingFields(StorageSettings.IndexEnabledEnum.ENABLED);
+		myMdmSettings.setSearchAllPartitionForMatch(true);
+		myMdmSettings.setGoldenResourcePartitionName(PARTITION_2);
+
+		MdmTransactionContext transactionContext = new MdmTransactionContext(MdmTransactionContext.OperationType.UPDATE_LINK);
+
+		// register the tenant partition interceptor
+		// req'd so partition values will be filled in
+		RequestTenantPartitionInterceptor tenantPartitionInterceptor = new RequestTenantPartitionInterceptor();
+		myInterceptorRegistry.registerInterceptor(tenantPartitionInterceptor);
+		try {
+			// create the partitions we need
+			PartitionEntity partition1 = new PartitionEntity();
+			partition1.setName(PARTITION_1);
+			partition1.setId(1);
+			partition1.setDescription("first");
+			myPartitionLookupSvc.createPartition(partition1, new SystemRequestDetails());
+			PartitionEntity partition2 = new PartitionEntity();
+			partition2.setName(PARTITION_2);
+			partition2.setId(2);
+			partition2.setDescription("second");
+			myPartitionLookupSvc.createPartition(partition2, new SystemRequestDetails());
+			myPartitionLookupSvc.invalidateCaches();
+
+			// create the patients - 2 in part-a, golden in partition golden
+			SystemRequestDetails partitionARequestDetails = new SystemRequestDetails();
+			partitionARequestDetails.setRequestPartitionId(RequestPartitionId.fromPartitionId(1));
+			Patient patient1 = buildJanePatient();
+			patient1.setId("Patient/pat-a");
+			myPatientDao.update(patient1,
+				null,
+				true,
+				false,
+				partitionARequestDetails,
+				new TransactionDetails());
+
+			Patient patient2 = buildJanePatient();
+			patient2.setId("Patient/pat-b");
+			myPatientDao.update(
+				patient2,
+				null,
+				true,
+				false,
+				partitionARequestDetails,
+				new TransactionDetails()
+			);
+
+			// manually create our golden resource
+			SystemRequestDetails goldenRequestDetails = new SystemRequestDetails();
+			goldenRequestDetails.setRequestPartitionId(RequestPartitionId.fromPartitionId(2));
+			final Patient goldenPatient = buildJanePatient();
+			goldenPatient.setId("pat-gold");
+			myPatientDao.update(goldenPatient,
+				null,
+				true,
+				false,
+				goldenRequestDetails,
+				new TransactionDetails());
+
+			// save our links
+			{
+				myMdmLinkDaoSvc.createOrUpdateLinkEntity(goldenPatient, patient1, MdmMatchOutcome.NEW_GOLDEN_RESOURCE_MATCH, MdmLinkSourceEnum.AUTO, createContextForCreate("Patient"));
+				myMdmLinkDaoSvc.createOrUpdateLinkEntity(goldenPatient, patient2, MdmMatchOutcome.NEW_GOLDEN_RESOURCE_MATCH, MdmLinkSourceEnum.AUTO, createContextForCreate("Patient"));
+			}
+
+			// remove link
+			{
+				Patient patb = new Patient();
+				patb.setId(patient2.getId());
+				myMdmLinkDaoSvc.createOrUpdateLinkEntity(goldenPatient, patb, MdmMatchOutcome.NO_MATCH, MdmLinkSourceEnum.MANUAL, createContextForUpdate("Patient"));
+			}
+
+			// test
+			myMdmSurvivorshipService.rebuildGoldenResourceWithSurvivorshipRules(goldenPatient, transactionContext);
+
+			IBundleProvider provider = myPatientDao.search(new SearchParameterMap().setLoadSynchronous(true),
+				new SystemRequestDetails().setRequestPartitionId(RequestPartitionId.allPartitions()));
+		} finally {
+			// reset
+			myInterceptorRegistry.unregisterInterceptor(tenantPartitionInterceptor);
+
+			myPartitionSettings.setPartitioningEnabled(isPartitioningEnabled);
+			myPartitionSettings.setUnnamedPartitionMode(isUnnamedPartitionMode);
+			myStorageSettings.setIndexMissingFields(indexMissingFields);
+			myMdmSettings.setSearchAllPartitionForMatch(isSearchAllPartitionForMatch);
+			myMdmSettings.setGoldenResourcePartitionName(goldenResourcePartitionName);
+		}
+	}
 }

hapi-fhir-jpaserver-mdm/src/test/java/ca/uhn/fhir/jpa/mdm/svc/MdmSurvivorshipSvcImplTest.java

@@ -14,7 +14,6 @@ import ca.uhn.fhir.mdm.api.MdmMatchResultEnum;
 import ca.uhn.fhir.mdm.model.CanonicalEID;
 import ca.uhn.fhir.mdm.model.MdmTransactionContext;
 import ca.uhn.fhir.mdm.model.mdmevents.MdmLinkJson;
-import ca.uhn.fhir.mdm.model.mdmevents.MdmLinkWithRevisionJson;
 import ca.uhn.fhir.mdm.rules.json.MdmRulesJson;
 import ca.uhn.fhir.mdm.svc.MdmSurvivorshipSvcImpl;
 import ca.uhn.fhir.mdm.util.EIDHelper;
@@ -102,6 +101,7 @@ public class MdmSurvivorshipSvcImplTest {
 	public void rebuildGoldenResourceCurrentLinksUsingSurvivorshipRules_withManyLinks_rebuildsInUpdateOrder(boolean theIsUseNonNumericId) {
 		// setup
 		// create resources
+		int goldenId = 777;
 		Patient goldenPatient = new Patient();
 		goldenPatient.addAddress()
 			.setCity("Toronto")
@@ -109,13 +109,13 @@ public class MdmSurvivorshipSvcImplTest {
 		goldenPatient.addName()
 			.setFamily("Doe")
 			.addGiven("Jane");
-		goldenPatient.setId("Patient/777");
+		goldenPatient.setId("Patient/" + goldenId);
 		MdmResourceUtil.setMdmManaged(goldenPatient);
 		MdmResourceUtil.setGoldenResource(goldenPatient);
 
 		List<IBaseResource> resources = new ArrayList<>();
 		List<MdmLinkJson> links = new ArrayList<>();
-		List<MdmLinkWithRevisionJson> linksWithRevisions = new ArrayList<>();
+
 		for (int i = 0; i < 10; i++) {
 			// we want our resources to be slightly different
 			Patient patient = new Patient();
@@ -128,21 +128,16 @@ public class MdmSurvivorshipSvcImplTest {
 			patient.addIdentifier()
 				.setSystem("http://example.com")
 				.setValue("Value" + i);
-			patient.setId("Patient/" + (theIsUseNonNumericId ? "pat"+i : Integer.toString(i)));
+			patient.setId("Patient/" + (theIsUseNonNumericId ? "pat" + i : Integer.toString(i)));
 			resources.add(patient);
 
 			MdmLinkJson link = createLinkJson(
 				patient,
 				goldenPatient
 			);
+			link.setSourcePid(JpaPid.fromId((long)i));
+			link.setGoldenPid(JpaPid.fromId((long)goldenId));
 			links.add(link);
-			linksWithRevisions.add(new MdmLinkWithRevisionJson(
-				link,
-				1L,
-				Date.from(
-					Instant.now().minus(i, ChronoUnit.HOURS)
-				)
-			));
 		}
 
 		IFhirResourceDao resourceDao = mock(IFhirResourceDao.class);
@@ -151,13 +146,8 @@ public class MdmSurvivorshipSvcImplTest {
 		when(myDaoRegistry.getResourceDao(eq("Patient")))
 			.thenReturn(resourceDao);
 		AtomicInteger counter = new AtomicInteger();
-		if (theIsUseNonNumericId) {
+		when(resourceDao.readByPid(any()))
-			when(resourceDao.read(any(), any()))
+			.thenAnswer(params -> resources.get(counter.getAndIncrement()));
-				.thenAnswer(params -> resources.get(counter.getAndIncrement()));
-		} else {
-			when(resourceDao.readByPid(any()))
-				.thenAnswer(params -> resources.get(counter.getAndIncrement()));
-		}
 		Page<MdmLinkJson> linkPage = mock(Page.class);
 		when(myMdmLinkQuerySvc.queryLinks(any(), any()))
 			.thenReturn(linkPage);

hapi-fhir-server-mdm/src/main/java/ca/uhn/fhir/mdm/model/mdmevents/MdmLinkJson.java

@@ -22,6 +22,7 @@ package ca.uhn.fhir.mdm.model.mdmevents;
 import ca.uhn.fhir.mdm.api.MdmLinkSourceEnum;
 import ca.uhn.fhir.mdm.api.MdmMatchResultEnum;
 import ca.uhn.fhir.model.api.IModelJson;
+import ca.uhn.fhir.rest.api.server.storage.IResourcePersistentId;
 import com.fasterxml.jackson.annotation.JsonProperty;
 
 import java.util.Date;
@@ -29,15 +30,39 @@ import java.util.Objects;
 
 public class MdmLinkJson implements IModelJson {
 
+	/**
+	 * Golden resource FhirId
+	 */
 	@JsonProperty("goldenResourceId")
 	private String myGoldenResourceId;
 
+	/**
+	 * Source resource FhirId
+	 */
 	@JsonProperty("sourceId")
 	private String mySourceId;
 
+	/**
+	 * Golden resource PID
+	 */
+	@JsonProperty("goldenPid")
+	private IResourcePersistentId<?> myGoldenPid;
+
+	/**
+	 * Source PID
+	 */
+	@JsonProperty("sourcePid")
+	private IResourcePersistentId<?> mySourcePid;
+
+	/**
+	 * Kind of link (MATCH, etc)
+	 */
 	@JsonProperty("matchResult")
 	private MdmMatchResultEnum myMatchResult;
 
+	/**
+	 * How the link was constructed (AUTO - by the system, MANUAL - by a user)
+	 */
 	@JsonProperty("linkSource")
 	private MdmLinkSourceEnum myLinkSource;
 
@@ -178,6 +203,22 @@ public class MdmLinkJson implements IModelJson {
 		myRuleCount = theRuleCount;
 	}
 
+	public IResourcePersistentId<?> getGoldenPid() {
+		return myGoldenPid;
+	}
+
+	public void setGoldenPid(IResourcePersistentId<?> theGoldenPid) {
+		myGoldenPid = theGoldenPid;
+	}
+
+	public IResourcePersistentId<?> getSourcePid() {
+		return mySourcePid;
+	}
+
+	public void setSourcePid(IResourcePersistentId<?> theSourcePid) {
+		mySourcePid = theSourcePid;
+	}
+
 	@Override
 	public boolean equals(Object theO) {
 		if (this == theO) return true;
@@ -185,6 +226,8 @@ public class MdmLinkJson implements IModelJson {
 		final MdmLinkJson that = (MdmLinkJson) theO;
 		return Objects.equals(myGoldenResourceId, that.myGoldenResourceId)
 				&& Objects.equals(mySourceId, that.mySourceId)
+				&& mySourcePid.equals(that.mySourcePid)
+				&& myGoldenPid.equals(that.myGoldenPid)
 				&& myMatchResult == that.myMatchResult
 				&& myLinkSource == that.myLinkSource
 				&& Objects.equals(myCreated, that.myCreated)
@@ -202,6 +245,8 @@ public class MdmLinkJson implements IModelJson {
 		return Objects.hash(
 				myGoldenResourceId,
 				mySourceId,
+				mySourcePid,
+				myGoldenPid,
 				myMatchResult,
 				myLinkSource,
 				myCreated,
@@ -216,18 +261,21 @@ public class MdmLinkJson implements IModelJson {
 
 	@Override
 	public String toString() {
-		return "MdmLinkJson{" + "myGoldenResourceId='"
+		return "MdmLinkJson{"
-				+ myGoldenResourceId + '\'' + ", mySourceId='"
+				+ "myGoldenResourceId='" + myGoldenResourceId + '\''
-				+ mySourceId + '\'' + ", myMatchResult="
+				+ ", myGoldenPid='" + myGoldenPid + '\''
-				+ myMatchResult + ", myLinkSource="
+				+ ", mySourceId='" + mySourceId + '\''
-				+ myLinkSource + ", myCreated="
+				+ ", mySourcePid='" + mySourcePid + '\''
-				+ myCreated + ", myUpdated="
+				+ ", myMatchResult=" + myMatchResult
-				+ myUpdated + ", myVersion='"
+				+ ", myLinkSource=" + myLinkSource
-				+ myVersion + '\'' + ", myEidMatch="
+				+ ", myCreated=" + myCreated
-				+ myEidMatch + ", myLinkCreatedNewResource="
+				+ ", myUpdated=" + myUpdated
-				+ myLinkCreatedNewResource + ", myVector="
+				+ ", myVersion='" + myVersion + '\''
-				+ myVector + ", myScore="
+				+ ", myEidMatch=" + myEidMatch
-				+ myScore + ", myRuleCount="
+				+ ", myLinkCreatedNewResource=" + myLinkCreatedNewResource
-				+ myRuleCount + '}';
+				+ ", myVector=" + myVector
+				+ ", myScore=" + myScore
+				+ ", myRuleCount=" + myRuleCount
+				+ '}';
 	}
 }

hapi-fhir-server-mdm/src/main/java/ca/uhn/fhir/mdm/svc/MdmSurvivorshipSvcImpl.java

@@ -20,6 +20,7 @@
 package ca.uhn.fhir.mdm.svc;
 
 import ca.uhn.fhir.context.FhirContext;
+import ca.uhn.fhir.interceptor.model.RequestPartitionId;
 import ca.uhn.fhir.jpa.api.dao.DaoRegistry;
 import ca.uhn.fhir.jpa.api.dao.IFhirResourceDao;
 import ca.uhn.fhir.jpa.api.svc.IIdHelperService;
@@ -31,7 +32,7 @@ import ca.uhn.fhir.mdm.api.params.MdmQuerySearchParameters;
 import ca.uhn.fhir.mdm.model.MdmTransactionContext;
 import ca.uhn.fhir.mdm.model.mdmevents.MdmLinkJson;
 import ca.uhn.fhir.mdm.util.GoldenResourceHelper;
-import ca.uhn.fhir.model.primitive.IdDt;
+import ca.uhn.fhir.rest.api.Constants;
 import ca.uhn.fhir.rest.api.server.SystemRequestDetails;
 import ca.uhn.fhir.rest.api.server.storage.IResourcePersistentId;
 import ca.uhn.fhir.util.TerserUtil;
@@ -118,7 +119,15 @@ public class MdmSurvivorshipSvcImpl implements IMdmSurvivorshipService {
 
 		// save it
 		IFhirResourceDao dao = myDaoRegistry.getResourceDao(goldenResource.fhirType());
-		dao.update(toSave, new SystemRequestDetails());
+
+		SystemRequestDetails requestDetails = new SystemRequestDetails();
+		// if using partitions, we should save to the correct partition
+		Object resourcePartitionIdObj = toSave.getUserData(Constants.RESOURCE_PARTITION_ID);
+		if (resourcePartitionIdObj instanceof RequestPartitionId) {
+			RequestPartitionId partitionId = (RequestPartitionId) resourcePartitionIdObj;
+			requestDetails.setRequestPartitionId(partitionId);
+		}
+		dao.update(toSave, requestDetails);
 
 		return (T) toSave;
 	}
@@ -135,24 +144,8 @@ public class MdmSurvivorshipSvcImpl implements IMdmSurvivorshipService {
 		Page<MdmLinkJson> linksQuery = myMdmLinkQuerySvc.queryLinks(searchParameters, theMdmTransactionContext);
 
 		return linksQuery.get().map(link -> {
-			String sourceId = link.getSourceId();
+			IResourcePersistentId<?> pid = link.getSourcePid();
-
+			return dao.readByPid(pid);
-			// +1 because of "/" in id: "ResourceType/Id"
-			final String sourceIdUnqualified = sourceId.substring(resourceType.length() + 1);
-
-			// myMdmLinkQuerySvc.queryLinks populates sourceId with the FHIR_ID, not the RES_ID, so if we don't
-			// add this conditional logic, on JPA, myIIdHelperService.newPidFromStringIdAndResourceName will fail with
-			// NumberFormatException
-			if (isNumericOrUuid(sourceIdUnqualified)) {
-				IResourcePersistentId<?> pid = getResourcePID(sourceIdUnqualified, resourceType);
-
-				// this might be a bit unperformant
-				// but it depends how many links there are
-				// per golden resource (unlikely to be thousands)
-				return dao.readByPid(pid);
-			} else {
-				return dao.read(new IdDt(sourceId), new SystemRequestDetails());
-			}
 		});
 	}