Review Comments: Refactored code to find references using LocalChange…

…ResourceReferenceEntity

engine/build.gradle.kts

@@ -140,6 +140,7 @@ dependencies {
  testImplementation(Dependencies.AndroidxTest.core)
  testImplementation(Dependencies.AndroidxTest.workTestingRuntimeKtx)
  testImplementation(Dependencies.Kotlin.kotlinCoroutinesTest)
+ testImplementation(Dependencies.Kotlin.kotlinTestJunit)
  testImplementation(Dependencies.junit)
  testImplementation(Dependencies.jsonAssert)
  testImplementation(Dependencies.mockitoInline)

engine/src/main/java/com/google/android/fhir/FhirEngineProvider.kt

@@ -18,6 +18,7 @@ package com.google.android.fhir
 
 import android.content.Context
 import com.google.android.fhir.DatabaseErrorStrategy.UNSPECIFIED
+import com.google.android.fhir.db.Database
 import com.google.android.fhir.sync.DataSource
 import com.google.android.fhir.sync.FhirDataStore
 import com.google.android.fhir.sync.HttpAuthenticator
@@ -67,6 +68,11 @@ object FhirEngineProvider {
  return getOrCreateFhirService(context).fhirDataStore
  }
 
+ @Synchronized
+ internal fun getFhirDatabase(context: Context): Database {
+ return getOrCreateFhirService(context).database
+ }
+
  @Synchronized
  private fun getOrCreateFhirService(context: Context): FhirServices {
  if (fhirServices == null) {

engine/src/main/java/com/google/android/fhir/db/Database.kt

@@ -185,9 +185,24 @@ internal interface Database {
  * will delete resource entry from LocalChangeEntity table.
  */
  suspend fun purge(type: ResourceType, id: String, forcePurge: Boolean = false)
+
+ /**
+ * @return List of [LocalChangeResourceReference] associated with the [LocalChangeEntity.id]s. A
+ * single [LocalChangeEntity] may have one or more [LocalChangeResourceReference] associated
+ * with it.
+ */
+ suspend fun getLocalChangeResourceReferences(
+ localChangeIds: List<Long>,
+ ): List<LocalChangeResourceReference>
 }
 
 data class ResourceWithUUID<R>(
  val uuid: UUID,
  val resource: R,
 )
+
+data class LocalChangeResourceReference(
+ val localChangeId: Long,
+ val resourceReferenceValue: String,
+ val resourceReferencePath: String?,
+)

engine/src/main/java/com/google/android/fhir/db/impl/DatabaseImpl.kt

@@ -26,6 +26,7 @@ import ca.uhn.fhir.util.FhirTerser
 import com.google.android.fhir.DatabaseErrorStrategy
 import com.google.android.fhir.LocalChange
 import com.google.android.fhir.LocalChangeToken
+import com.google.android.fhir.db.LocalChangeResourceReference
 import com.google.android.fhir.db.ResourceNotFoundException
 import com.google.android.fhir.db.ResourceWithUUID
 import com.google.android.fhir.db.impl.DatabaseImpl.Companion.UNENCRYPTED_DATABASE_NAME
@@ -408,6 +409,18 @@ internal class DatabaseImpl(
  }
  }
 
+ override suspend fun getLocalChangeResourceReferences(
+ localChangeIds: List<Long>,
+ ): List<LocalChangeResourceReference> {
+ return localChangeDao.getReferencesForLocalChanges(localChangeIds).map {
+ LocalChangeResourceReference(
+ it.localChangeId,
+ it.resourceReferenceValue,
+ it.resourceReferencePath,
+ )
+ }
+ }
+
  companion object {
  /**
  * The name for unencrypted database.

engine/src/main/java/com/google/android/fhir/db/impl/dao/LocalChangeDao.kt

@@ -1,5 +1,5 @@
 /*
- * Copyright 2023 Google LLC
+ * Copyright 2023-2024 Google LLC
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -348,6 +348,17 @@ internal abstract class LocalChangeDao {
  localChangeId: Long,
  ): List<LocalChangeResourceReferenceEntity>
 
+ @Query(
+ """
+ SELECT *
+ FROM LocalChangeResourceReferenceEntity
+ WHERE localChangeId = (:localChangeId)
+ """,
+ )
+ abstract suspend fun getReferencesForLocalChanges(
+ localChangeId: List<Long>,
+ ): List<LocalChangeResourceReferenceEntity>
+
  @Insert(onConflict = OnConflictStrategy.REPLACE)
  abstract suspend fun insertLocalChangeResourceReferences(
  resourceReferences: List<LocalChangeResourceReferenceEntity>,

engine/src/main/java/com/google/android/fhir/sync/FhirSyncWorker.kt

@@ -76,7 +76,11 @@ abstract class FhirSyncWorker(appContext: Context, workerParams: WorkerParameter
  UploadConfiguration(
  Uploader(
  dataSource = dataSource,
- patchGenerator = PatchGeneratorFactory.byMode(getUploadStrategy().patchGeneratorMode),
+ patchGenerator =
+ PatchGeneratorFactory.byMode(
+ getUploadStrategy().patchGeneratorMode,
+ FhirEngineProvider.getFhirDatabase(applicationContext),
+ ),
  requestGenerator =
  UploadRequestGeneratorFactory.byMode(getUploadStrategy().requestGeneratorMode),
  ),

engine/src/main/java/com/google/android/fhir/sync/upload/patch/PatchGenerator.kt

@@ -1,5 +1,5 @@
 /*
- * Copyright 2023 Google LLC
+ * Copyright 2023-2024 Google LLC
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -17,6 +17,7 @@
 package com.google.android.fhir.sync.upload.patch
 
 import com.google.android.fhir.LocalChange
+import com.google.android.fhir.db.Database
 
 /**
  * Generates [Patch]es from [LocalChange]s and output [List<[PatchMapping]>] to keep a mapping of
@@ -34,16 +35,17 @@ internal interface PatchGenerator {
  * NOTE: different implementations may have requirements on the size of [localChanges] and output
  * certain numbers of [Patch]es.
  */
- fun generate(localChanges: List<LocalChange>): List<PatchMapping>
+ suspend fun generate(localChanges: List<LocalChange>): List<PatchMapping>
 }
 
 internal object PatchGeneratorFactory {
  fun byMode(
  mode: PatchGeneratorMode,
+ database: Database,
  ): PatchGenerator =
  when (mode) {
  is PatchGeneratorMode.PerChange -> PerChangePatchGenerator
- is PatchGeneratorMode.PerResource -> PerResourcePatchGenerator
+ is PatchGeneratorMode.PerResource -> PerResourcePatchGenerator(database)
  }
 }
 

engine/src/main/java/com/google/android/fhir/sync/upload/patch/PerChangePatchGenerator.kt

@@ -1,5 +1,5 @@
 /*
- * Copyright 2023 Google LLC
+ * Copyright 2023-2024 Google LLC
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -25,7 +25,7 @@ import com.google.android.fhir.LocalChange
  * maintain an audit trail.
  */
 internal object PerChangePatchGenerator : PatchGenerator {
- override fun generate(localChanges: List<LocalChange>): List<PatchMapping> =
+ override suspend fun generate(localChanges: List<LocalChange>): List<PatchMapping> =
  localChanges.map {
  PatchMapping(
  localChanges = listOf(it),

engine/src/main/java/com/google/android/fhir/sync/upload/patch/PerResourcePatchGenerator.kt

@@ -22,6 +22,8 @@ import com.github.fge.jackson.JsonLoader
 import com.github.fge.jsonpatch.JsonPatch
 import com.google.android.fhir.LocalChange
 import com.google.android.fhir.LocalChange.Type
+import com.google.android.fhir.db.Database
+import com.google.android.fhir.db.LocalChangeResourceReference
 
 /**
  * Generates a [Patch] for all [LocalChange]es made to a single FHIR resource.
@@ -30,10 +32,10 @@ import com.google.android.fhir.LocalChange.Type
  * maintain an audit trail, but instead, multiple changes made to the same FHIR resource on the
  * client can be recorded as a single change on the server.
  */
-internal object PerResourcePatchGenerator : PatchGenerator {
+internal class PerResourcePatchGenerator(val database: Database) : PatchGenerator {
 
- override fun generate(localChanges: List<LocalChange>): List<PatchMapping> {
- return generateSquashedChangesMapping(localChanges).orderByReferences()
+ override suspend fun generate(localChanges: List<LocalChange>): List<PatchMapping> {
+ return generateSquashedChangesMapping(localChanges).orderByReferences(database)
  }
 
  @androidx.annotation.VisibleForTesting
@@ -152,7 +154,7 @@ private typealias Node = String
  * - throws [IllegalStateException] otherwise
  */
 @androidx.annotation.VisibleForTesting
-internal fun List<PatchMapping>.orderByReferences(): List<PatchMapping> {
+internal suspend fun List<PatchMapping>.orderByReferences(database: Database): List<PatchMapping> {
  // if the resource A has outgoing references (B,C) and these referenced resources are getting
  // created on device,
  // then these referenced resources (B,C) needs to go before the resource A so that referential
@@ -161,11 +163,18 @@ internal fun List<PatchMapping>.orderByReferences(): List<PatchMapping> {
  "${patchMapping.generatedPatch.resourceType}/${patchMapping.generatedPatch.resourceId}"
  }
 
+ // get references for all the local changes
+ val localChangeIdToReferenceMap: Map<Long, List<LocalChangeResourceReference>> =
+ database
+ .getLocalChangeResourceReferences(flatMap { it.localChanges.flatMap { it.token.ids } })
+ .groupBy { it.localChangeId }
+
  val adjacencyList =
- createReferenceAdjacencyList(
+ createAdjacencyListForCreateReferences(
  resourceIdToPatchMapping.keys
  .filter { resourceIdToPatchMapping[it]?.generatedPatch?.type == Patch.Type.INSERT }
  .toSet(),
+ localChangeIdToReferenceMap,
  )
  return createTopologicalOrderedList(adjacencyList).mapNotNull { resourceIdToPatchMapping[it] }
 }
@@ -175,33 +184,36 @@ internal fun List<PatchMapping>.orderByReferences(): List<PatchMapping> {
  * type [Patch.Type.INSERT] .
  */
 @androidx.annotation.VisibleForTesting
-internal fun List<PatchMapping>.createReferenceAdjacencyList(
+internal fun List<PatchMapping>.createAdjacencyListForCreateReferences(
  insertResourceIds: Set<String>,
+ localChangeIdToReferenceMap: Map<Long, List<LocalChangeResourceReference>>,
 ): Map<Node, List<Node>> {
  val adjacencyList = mutableMapOf<Node, List<Node>>()
  forEach { patchMapping ->
  adjacencyList[
  "${patchMapping.generatedPatch.resourceType}/${patchMapping.generatedPatch.resourceId}",
- ] = patchMapping.findOutgoingReferences().filter { insertResourceIds.contains(it) }
+ ] =
+ patchMapping.findOutgoingReferences(localChangeIdToReferenceMap).filter {
+ insertResourceIds.contains(it)
+ }
  }
  return adjacencyList
 }
 
 // if the outgoing reference is to a resource that's just an update and not create, then don't link
 // to it. This may make the sub graphs smaller and also help avoid cyclic dependencies.
-@androidx.annotation.VisibleForTesting
-internal fun PatchMapping.findOutgoingReferences(): List<Node> {
- val references = mutableListOf<Node>()
+private fun PatchMapping.findOutgoingReferences(
+ localChangeIdToReferenceMap: Map<Long, List<LocalChangeResourceReference>>,
+): Set<Node> {
+ val references = mutableSetOf<Node>()
  when (generatedPatch.type) {
- Patch.Type.INSERT -> {
- JsonLoader.fromString(generatedPatch.payload).search(references)
- }
- Patch.Type.UPDATE -> {
- JsonLoader.fromString(generatedPatch.payload).forEach {
- if (it.get("path").asText().endsWith("/reference")) {
- references.add(it.get("value").asText())
- } else {
- it.get("value").search(references)
+ Patch.Type.INSERT,
+ Patch.Type.UPDATE, -> {
+ localChanges.forEach { localChange ->
+ localChange.token.ids.forEach { id ->
+ localChangeIdToReferenceMap[id]?.let {
+ references.addAll(it.map { it.resourceReferenceValue })
+ }
  }
  }
  }
@@ -212,13 +224,6 @@ internal fun PatchMapping.findOutgoingReferences(): List<Node> {
  return references
 }
 
-private fun JsonNode.search(reference: MutableList<String>) {
- if (has("reference") && get("reference").isTextual) {
- reference.add(get("reference").asText())
- }
- elements().forEach { it.search(reference) }
-}
-
 @androidx.annotation.VisibleForTesting
 internal fun createTopologicalOrderedList(adjacencyList: Map<Node, List<Node>>): List<Node> {
  val stack = ArrayDeque<String>()

engine/src/main/java/com/google/android/fhir/testing/Utilities.kt

@@ -192,7 +192,7 @@ object TestFhirEngineImpl : FhirEngine {
  resourceType = type.name,
  resourceId = id,
  payload = """{ "resourceType" : "$type", "id" : "$id" }""",
- token = LocalChangeToken(listOf()),
+ token = LocalChangeToken(listOf(1)),
  type = LocalChange.Type.INSERT,
  timestamp = Instant.now(),
  ),

engine/src/test/java/com/google/android/fhir/sync/upload/UploaderTest.kt

@@ -20,8 +20,10 @@ import ca.uhn.fhir.context.FhirContext
 import ca.uhn.fhir.context.FhirVersionEnum
 import com.google.android.fhir.LocalChange
 import com.google.android.fhir.LocalChangeToken
+import com.google.android.fhir.db.Database
 import com.google.android.fhir.db.impl.entities.LocalChangeEntity
 import com.google.android.fhir.logicalId
+import com.google.android.fhir.sync.upload.patch.PatchGenerator
 import com.google.android.fhir.sync.upload.patch.PatchGeneratorFactory
 import com.google.android.fhir.sync.upload.patch.PatchGeneratorMode
 import com.google.android.fhir.sync.upload.request.UploadRequestGeneratorFactory
@@ -42,13 +44,32 @@ import org.hl7.fhir.r4.model.OperationOutcome
 import org.hl7.fhir.r4.model.Patient
 import org.hl7.fhir.r4.model.ResourceType
 import org.hl7.fhir.r4.model.codesystems.HttpVerb
+import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
+import org.mockito.Mock
+import org.mockito.MockitoAnnotations
+import org.mockito.kotlin.any
+import org.mockito.kotlin.whenever
 import org.robolectric.RobolectricTestRunner
 
 @RunWith(RobolectricTestRunner::class)
 class UploaderTest {
 
+ @Mock private lateinit var database: Database
+ private lateinit var perResourcePatchGenerator: PatchGenerator
+ private lateinit var perChangePatchGenerator: PatchGenerator
+
+ @Before
+ fun setUp() {
+ MockitoAnnotations.openMocks(this)
+ runTest { whenever(database.getLocalChangeResourceReferences(any())).thenReturn(emptyList()) }
+
+ perResourcePatchGenerator =
+ PatchGeneratorFactory.byMode(PatchGeneratorMode.PerResource, database)
+ perChangePatchGenerator = PatchGeneratorFactory.byMode(PatchGeneratorMode.PerChange, database)
+ }
+
  @Test
  fun `bundle upload for per resource patch should output responses mapped correctly to the local changes`() =
  runTest {
@@ -513,9 +534,7 @@ class UploaderTest {
  }
 
  companion object {
- private val perResourcePatchGenerator =
- PatchGeneratorFactory.byMode(PatchGeneratorMode.PerResource)
- private val perChangePatchGenerator = PatchGeneratorFactory.byMode(PatchGeneratorMode.PerChange)
+
  private val urlUploadRequestGenerator =
  UploadRequestGeneratorFactory.byMode(
  UploadRequestGeneratorMode.UrlRequest(HttpVerb.PUT, HttpVerb.PATCH),