Add FhirEngine interface method 'withTransaction'(google#2535)

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

@@ -58,55 +58,7 @@ import org.hl7.fhir.r4.model.ResourceType
  * val fhirEngine = FhirEngineProvider.getInstance(this)
  * ```
  */
-interface FhirEngine {
-  /**
-   * Creates one or more FHIR [Resource]s in the local storage. FHIR Engine requires all stored
-   * resources to have a logical [Resource.id]. If the `id` is specified in the resource passed to
-   * [create], the resource created in `FhirEngine` will have the same `id`. If no `id` is
-   * specified, `FhirEngine` will generate a UUID as that resource's `id` and include it in the
-   * returned list of IDs.
-   * @param resource The FHIR resources to create.
-   * @param isLocalOnly - Setting the value to [true] instructs engine that the resource and its
-   *   subsequent updates should never be synced to the server.
-   * @return A list of logical IDs of the newly created resources.
-   */
-  suspend fun create(vararg resource: Resource, isLocalOnly: Boolean = false): List<String>
-
-  /**
-   * Loads a FHIR resource given its [ResourceType] and logical ID.
-   *
-   * @param type The type of the resource to load.
-   * @param id The logical ID of the resource.
-   * @return The requested FHIR resource.
-   * @throws ResourceNotFoundException if the resource is not found.
-   */
-  @Throws(ResourceNotFoundException::class)
-  suspend fun get(type: ResourceType, id: String): Resource
-
-  /**
-   * Updates one or more FHIR [Resource]s in the local storage.
-   *
-   * @param resource The FHIR resources to update.
-   */
-  suspend fun update(vararg resource: Resource)
-
-  /**
-   * Removes a FHIR resource given its [ResourceType] and logical ID.
-   *
-   * @param type The type of the resource to delete.
-   * @param id The logical ID of the resource.
-   */
-  suspend fun delete(type: ResourceType, id: String)
-
-  /**
-   * Searches the database and returns a list of resources matching the [Search] specifications.
-   *
-   * @param search The search criteria to apply.
-   * @return A list of [SearchResult] objects containing the matching resources and any included
-   *   references.
-   */
-  suspend fun <R : Resource> search(search: Search): List<SearchResult<R>>
-
+interface FhirEngine : CrudFhirEngine {
   /**
    * Synchronizes upload results with the database.
    *
@@ -207,6 +159,71 @@ interface FhirEngine {
   suspend fun purge(type: ResourceType, ids: Set<String>, forcePurge: Boolean = false)
 }
 
+interface CrudFhirEngine {
+  /**
+   * Creates one or more FHIR [Resource]s in the local storage. FHIR Engine requires all stored
+   * resources to have a logical [Resource.id]. If the `id` is specified in the resource passed to
+   * [create], the resource created in `FhirEngine` will have the same `id`. If no `id` is
+   * specified, `FhirEngine` will generate a UUID as that resource's `id` and include it in the
+   * returned list of IDs.
+   *
+   * @param resource The FHIR resources to create.
+   * @return A list of logical IDs of the newly created resources.
+   */
+  suspend fun create(vararg resource: Resource): List<String>
+
+  /**
+   * Loads a FHIR resource given its [ResourceType] and logical ID.
+   *
+   * @param type The type of the resource to load.
+   * @param id The logical ID of the resource.
+   * @return The requested FHIR resource.
+   * @throws ResourceNotFoundException if the resource is not found.
+   */
+  @Throws(ResourceNotFoundException::class)
+  suspend fun get(type: ResourceType, id: String): Resource
+
+  /**
+   * Updates one or more FHIR [Resource]s in the local storage.
+   *
+   * @param resource The FHIR resources to update.
+   */
+  suspend fun update(vararg resource: Resource)
+
+  /**
+   * Removes a FHIR resource given its [ResourceType] and logical ID.
+   *
+   * @param type The type of the resource to delete.
+   * @param id The logical ID of the resource.
+   */
+  suspend fun delete(type: ResourceType, id: String)
+
+  /**
+   * Searches the database and returns a list of resources matching the [Search] specifications.
+   *
+   * Example:
+   * ```
+   * fhirEngine.search<Patient> {
+   *  filter(Patient.GIVEN, {
+   *    value = "Kiran"
+   *    modifier = StringFilterModifier.MATCHES_EXACTLY
+   *  })
+   * }
+   * ```
+   *
+   * @param search The search criteria to apply.
+   * @return A list of [SearchResult] objects containing the matching resources and any included
+   *   references.
+   */
+  suspend fun <R : Resource> search(search: Search): List<SearchResult<R>>
+
+  /**
+   * Adds support for performing actions on `FhirEngine` as a single atomic transaction where the
+   * entire set of changes succeed or fail as a single entity
+   */
+  suspend fun withTransaction(block: suspend CrudFhirEngine.() -> Unit)
+}
+
 /**
  * Retrieves a FHIR resource of type [R] with the given [id] from the local storage.
  *

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

@@ -39,9 +39,11 @@ import com.google.android.fhir.search.SearchQuery
 import com.google.android.fhir.toLocalChange
 import java.time.Instant
 import java.util.UUID
+import kotlinx.coroutines.async
+import kotlinx.coroutines.awaitAll
+import kotlinx.coroutines.coroutineScope
 import org.hl7.fhir.r4.model.Resource
 import org.hl7.fhir.r4.model.ResourceType
-import java.util.Collections
 
 /**
  * The implementation for the persistence layer using Room. See docs for
@@ -144,7 +146,9 @@ internal class DatabaseImpl(
   }
 
   override suspend fun <R : Resource> insertLocalOnly(vararg resource: R): List<String> {
-    return db.withTransaction { resourceDao.insertAllRemote(resource.toList()).map { it.toString() }.toList() }
+    return db.withTransaction {
+      resourceDao.insertAllRemote(resource.toList()).map { it.toString() }.toList()
+    }
   }
 
   override suspend fun <R : Resource> insertRemote(vararg resource: R) {
@@ -211,11 +215,15 @@ internal class DatabaseImpl(
     return db.withTransaction {
       resourceDao
         .getResources(SimpleSQLiteQuery(query.query, query.args.toTypedArray()))
-        .map { ResourceWithUUID(it.uuid, iParser.parseResource(it.serializedResource) as R) }
+        .pmap { ResourceWithUUID(it.uuid, iParser.parseResource(it.serializedResource) as R) }
         .distinctBy { it.uuid }
     }
   }
 
+  private suspend fun <A, B> Iterable<A>.pmap(f: suspend (A) -> B): List<B> = coroutineScope {
+    map { async { f(it) } }.awaitAll()
+  }
+
   override suspend fun searchForwardReferencedResources(
     query: SearchQuery,
   ): List<ForwardIncludeSearchResult> {

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

@@ -17,6 +17,7 @@
 package com.google.android.fhir.impl
 
 import android.content.Context
+import com.google.android.fhir.CrudFhirEngine
 import com.google.android.fhir.FhirEngine
 import com.google.android.fhir.FhirEngineProvider
 import com.google.android.fhir.LocalChange
@@ -109,6 +110,10 @@ internal class FhirEngineImpl(private val database: Database, private val contex
     }
   }
 
+  override suspend fun withTransaction(block: suspend CrudFhirEngine.() -> Unit) {
+    database.withTransaction { block.invoke(this@FhirEngineImpl) }
+  }
+
   private suspend fun saveResolvedResourcesToDatabase(resolved: List<Resource>?) {
     resolved?.let {
       database.deleteUpdates(it)

engine/src/main/java/com/google/android/fhir/sync/upload/request/UploadRequestGenerator.kt

@@ -62,7 +62,7 @@ internal object UploadRequestGeneratorFactory {
           mode.httpVerbToUseForCreate,
           mode.httpVerbToUseForUpdate,
           mode.bundleSize,
-          useETagForUpload = false
+          useETagForUpload = false,
         )
     }
 }

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

@@ -20,6 +20,7 @@ import androidx.work.Data
 import ca.uhn.fhir.context.FhirContext
 import ca.uhn.fhir.context.FhirVersionEnum
 import ca.uhn.fhir.parser.IParser
+import com.google.android.fhir.CrudFhirEngine
 import com.google.android.fhir.FhirEngine
 import com.google.android.fhir.LocalChange
 import com.google.android.fhir.LocalChangeToken
@@ -176,6 +177,8 @@ internal object TestFhirEngineImpl : FhirEngine {
     download().collect()
   }
 
+  override suspend fun withTransaction(block: suspend CrudFhirEngine.() -> Unit) {}
+
   override suspend fun count(search: Search): Long {
     return 0
   }

engine/src/test/java/com/google/android/fhir/impl/FhirEngineImplTest.kt

@@ -47,13 +47,18 @@ import kotlinx.coroutines.runBlocking
 import kotlinx.coroutines.test.runTest
 import org.hl7.fhir.exceptions.FHIRException
 import org.hl7.fhir.r4.model.Address
+import org.hl7.fhir.r4.model.Appointment
 import org.hl7.fhir.r4.model.CanonicalType
 import org.hl7.fhir.r4.model.Coding
 import org.hl7.fhir.r4.model.DateTimeType
+import org.hl7.fhir.r4.model.Encounter
 import org.hl7.fhir.r4.model.Enumerations
 import org.hl7.fhir.r4.model.HumanName
 import org.hl7.fhir.r4.model.Meta
+import org.hl7.fhir.r4.model.Observation
 import org.hl7.fhir.r4.model.Patient
+import org.hl7.fhir.r4.model.Practitioner
+import org.hl7.fhir.r4.model.Reference
 import org.hl7.fhir.r4.model.ResourceType
 import org.junit.Assert.assertThrows
 import org.junit.Before
@@ -808,6 +813,132 @@ class FhirEngineImplTest {
     assertThat(services.database.getLocalChangesCount()).isEqualTo(0)
   }
 
+  @Test
+  fun `withTransaction saves all changes successfully in order`() = runTest {
+    val patient01ID = "patient-01"
+    val patient01 =
+      Patient().apply {
+        id = patient01ID
+        gender = Enumerations.AdministrativeGender.FEMALE
+      }
+    val patient01AppointmentID = "appointment-01"
+    val patient01Appointment =
+      Appointment().apply {
+        id = patient01AppointmentID
+        status = Appointment.AppointmentStatus.BOOKED
+        addParticipant(
+          Appointment.AppointmentParticipantComponent().apply {
+            actor = Reference("${patient01.resourceType}/$patient01ID")
+          },
+        )
+      }
+    fhirEngine.create(patient01, patient01Appointment)
+
+    // Fulfill appointment with related Encounter/Observation
+    val patient01AppointmentEncounterID = "enc-01"
+    val patient01AppointmentEncounter =
+      Encounter().apply {
+        id = patient01AppointmentEncounterID
+        subject = Reference("${patient01.resourceType}/$patient01ID")
+        addAppointment(Reference("${patient01Appointment.resourceType}/$patient01AppointmentID"))
+      }
+    val patient01AppointmentEncounterObservationID = "obs-01"
+    val patient01AppointmentEncounterObservation =
+      Observation().apply {
+        id = patient01AppointmentEncounterObservationID
+        encounter =
+          Reference(
+            "${patient01AppointmentEncounter.resourceType}/$patient01AppointmentEncounterID",
+          )
+      }
+    val updatedAppointment =
+      patient01Appointment.copy().apply { status = Appointment.AppointmentStatus.FULFILLED }
+
+    fhirEngine.withTransaction {
+      this.create(patient01AppointmentEncounter, patient01AppointmentEncounterObservation)
+      this.update(updatedAppointment)
+    }
+
+    assertThat(
+        fhirEngine.get<Encounter>(patient01AppointmentEncounterID).appointmentFirstRep.reference,
+      )
+      .isEqualTo("Appointment/$patient01AppointmentID")
+    assertThat(
+        fhirEngine.get<Observation>(patient01AppointmentEncounterObservationID).encounter.reference,
+      )
+      .isEqualTo("Encounter/$patient01AppointmentEncounterID")
+    assertThat(fhirEngine.get<Appointment>(patient01AppointmentID).status)
+      .isEqualTo(Appointment.AppointmentStatus.FULFILLED)
+  }
+
+  @Test
+  fun `withTransaction reverts all changes when an error occurs`() = runTest {
+    val patient01ID = "patient-01"
+    val patient01 =
+      Patient().apply {
+        id = patient01ID
+        gender = Enumerations.AdministrativeGender.FEMALE
+      }
+    val patient01AppointmentID = "appointment-01"
+    val patient01Appointment =
+      Appointment().apply {
+        id = patient01AppointmentID
+        status = Appointment.AppointmentStatus.BOOKED
+        addParticipant(
+          Appointment.AppointmentParticipantComponent().apply {
+            actor = Reference("${patient01.resourceType}/$patient01ID")
+          },
+        )
+      }
+    fhirEngine.create(patient01, patient01Appointment)
+
+    // Fulfill appointment with related Encounter/Observation
+    val patient01AppointmentEncounterID = "enc-01"
+    val patient01AppointmentEncounter =
+      Encounter().apply {
+        id = patient01AppointmentEncounterID
+        subject = Reference("${patient01.resourceType}/$patient01ID")
+        addAppointment(Reference("${patient01Appointment.resourceType}/$patient01AppointmentID"))
+      }
+    val patient01AppointmentEncounterObservationID = "obs-01"
+    val patient01AppointmentEncounterObservation =
+      Observation().apply {
+        id = patient01AppointmentEncounterObservationID
+        encounter =
+          Reference(
+            "${patient01AppointmentEncounter.resourceType}/$patient01AppointmentEncounterID",
+          )
+      }
+
+    try {
+      fhirEngine.withTransaction {
+        this.create(patient01AppointmentEncounter, patient01AppointmentEncounterObservation)
+        // Get non-existent practitioner to force ResourceNotFoundException
+        val nonExistentPractitioner =
+          this.get(ResourceType.Practitioner, "non_existent_practitioner_id") as Practitioner
+        val updatedAppointment =
+          patient01Appointment.copy().apply {
+            status = Appointment.AppointmentStatus.FULFILLED
+            addParticipant(
+              Appointment.AppointmentParticipantComponent().apply {
+                actor = Reference("Practitioner/${nonExistentPractitioner.logicalId}")
+              },
+            )
+          }
+        this.update(updatedAppointment)
+      }
+    } catch (_: ResourceNotFoundException) {}
+
+    assertThrows(ResourceNotFoundException::class.java) {
+      runBlocking { fhirEngine.get<Encounter>(patient01AppointmentEncounterID) }
+    }
+    assertThrows(ResourceNotFoundException::class.java) {
+      runBlocking { fhirEngine.get<Observation>(patient01AppointmentEncounterObservationID) }
+    }
+    assertThat(fhirEngine.get<Appointment>(patient01AppointmentID).status)
+      .isEqualTo(Appointment.AppointmentStatus.BOOKED)
+  }
+
   companion object {
     private const val TEST_PATIENT_1_ID = "test_patient_1"
     private var TEST_PATIENT_1 =