database: add awaitWhileOnline when persistence disabled

Author: nbransby

firebase-database/src/androidMain/kotlin/dev/teamhub/firebase/database/database.kt

@@ -1,5 +1,6 @@
 package dev.teamhub.firebase.database
 
+import com.google.android.gms.tasks.Task
 import com.google.firebase.database.Logger
 import com.google.firebase.database.ServerValue
 import com.google.firebase.database.ValueEventListener
@@ -8,11 +9,31 @@ import dev.teamhub.firebase.FirebaseApp
 import dev.teamhub.firebase.decode
 import dev.teamhub.firebase.encode
 import kotlinx.coroutines.channels.awaitClose
+import kotlinx.coroutines.coroutineScope
 import kotlinx.coroutines.flow.callbackFlow
+import kotlinx.coroutines.flow.conflate
+import kotlinx.coroutines.flow.filter
+import kotlinx.coroutines.flow.produceIn
+import kotlinx.coroutines.selects.select
+import kotlinx.coroutines.tasks.asDeferred
 import kotlinx.coroutines.tasks.await
 import kotlinx.serialization.DeserializationStrategy
 import kotlinx.serialization.SerializationStrategy
 
+suspend fun <T> Task<T>.awaitWhileOnline(): T = coroutineScope {
+    val notConnected = Firebase.database
+        .reference(".info/connected")
+        .snapshots
+        .filter { !it.value<Boolean>() }
+        .conflate()
+        .produceIn(this)
+
+    select<T> {
+        asDeferred().onAwait { it }
+        notConnected.onReceive { throw DatabaseException("Database not connected") }
+    }
+}
+
 actual val Firebase.database
     get() = FirebaseDatabase(com.google.firebase.database.FirebaseDatabase.getInstance())
 
@@ -26,18 +47,26 @@ actual fun Firebase.database(app: FirebaseApp, url: String) =
     FirebaseDatabase(com.google.firebase.database.FirebaseDatabase.getInstance(app.android, url))
 
 actual class FirebaseDatabase internal constructor(val android: com.google.firebase.database.FirebaseDatabase) {
-    actual fun reference(path: String) = DatabaseReference(android.getReference(path))
 
-    actual fun setPersistenceEnabled(enabled: Boolean) = android.setPersistenceEnabled(enabled)
+    private var persistenceEnabled = true
+
+    actual fun reference(path: String) =
+        DatabaseReference(android.getReference(path), persistenceEnabled)
+
+    actual fun setPersistenceEnabled(enabled: Boolean) =
+        android.setPersistenceEnabled(enabled).also { persistenceEnabled = enabled }
 
     actual fun setLoggingEnabled(enabled: Boolean) =
         android.setLogLevel(Logger.Level.DEBUG.takeIf { enabled } ?: Logger.Level.NONE)
 }
 
-actual class DatabaseReference internal constructor(val android: com.google.firebase.database.DatabaseReference) {
+actual class DatabaseReference internal constructor(
+    val android: com.google.firebase.database.DatabaseReference,
+    val persistenceEnabled: Boolean
+) {
 
-    actual fun push() = DatabaseReference(android.push())
-    actual fun onDisconnect() = OnDisconnect(android.onDisconnect())
+    actual fun push() = DatabaseReference(android.push(), persistenceEnabled)
+    actual fun onDisconnect() = OnDisconnect(android.onDisconnect(), persistenceEnabled)
 
     actual val snapshots get() = callbackFlow {
         val listener = object : ValueEventListener {
@@ -53,16 +82,23 @@ actual class DatabaseReference internal constructor(val android: com.google.fire
         awaitClose { android.removeEventListener(listener) }
     }
 
-    actual suspend fun  setValue(value: Any?) =
-        android.setValue(encode(value)).await().run { Unit }
+    actual suspend fun setValue(value: Any?) = android.setValue(encode(value))
+        .run { if(persistenceEnabled) await() else awaitWhileOnline() }
+        .run { Unit }
 
     actual suspend inline fun <reified T> setValue(strategy: SerializationStrategy<T>, value: T) =
-        android.setValue(encode(strategy, value)).await().run { Unit }
+        android.setValue(encode(strategy, value))
+            .run { if(persistenceEnabled) await() else awaitWhileOnline() }
+            .run { Unit }
 
     actual suspend fun updateChildren(update: Map<String, Any?>) =
-        android.updateChildren(update.mapValues { (_, it) -> encode(value = it) }).await().run { Unit }
+        android.updateChildren(update.mapValues { (_, it) -> encode(value = it) })
+            .run { if(persistenceEnabled) await() else awaitWhileOnline() }
+            .run { Unit }
 
-    actual suspend fun removeValue() = android.removeValue().await().run { Unit }
+    actual suspend fun removeValue() = android.removeValue()
+        .run { if(persistenceEnabled) await() else awaitWhileOnline() }
+        .run { Unit }
 }
 
 @Suppress("UNCHECKED_CAST")
@@ -80,19 +116,33 @@ actual class DataSnapshot internal constructor(val android: com.google.firebase.
     actual val children: Iterable<DataSnapshot> get() = android.children.map { DataSnapshot(it) }
 }
 
-actual class OnDisconnect internal constructor(val android: com.google.firebase.database.OnDisconnect) {
+actual class OnDisconnect internal constructor(
+    val android: com.google.firebase.database.OnDisconnect,
+    val persistenceEnabled: Boolean
+) {
+
+    actual suspend fun removeValue() = android.removeValue()
+        .run { if(persistenceEnabled) await() else awaitWhileOnline() }
+        .run { Unit }
 
-    actual suspend fun removeValue() = android.removeValue().await().run { Unit }
-    actual suspend fun cancel() = android.cancel().await().run { Unit }
+    actual suspend fun cancel() = android.cancel()
+        .run { if(persistenceEnabled) await() else awaitWhileOnline() }
+        .run { Unit }
 
     actual suspend inline fun <reified T : Any> setValue(value: T) =
-        android.setValue(encode(value = value)).await().run { Unit }
+        android.setValue(encode(value = value))
+            .run { if(persistenceEnabled) await() else awaitWhileOnline() }
+            .run { Unit }
 
     actual suspend inline fun <reified T> setValue(strategy: SerializationStrategy<T>, value: T) =
-        android.setValue(encode(strategy, value)).await().run { Unit }
+        android.setValue(encode(strategy, value))
+            .run { if(persistenceEnabled) await() else awaitWhileOnline() }
+            .run { Unit}
 
     actual suspend fun updateChildren(update: Map<String, Any?>) =
-        android.updateChildren(update.mapValues { (_, it) -> encode(value = it) }).await().run { Unit }
+        android.updateChildren(update.mapValues { (_, it) -> encode(value = it) })
+            .run { if(persistenceEnabled) await() else awaitWhileOnline() }
+            .run { Unit }
 }
 
 actual typealias DatabaseException = com.google.firebase.database.DatabaseException