From e9729d379e3420f8197d14795c0fed936f852929 Mon Sep 17 00:00:00 2001
From: Mikhail Pankratov <mikhail.pankratov@yamaneco.io>
Date: Wed, 5 Feb 2025 22:30:59 +0900
Subject: [PATCH] trying to implement count function

---
 .../firestore/internal/NativeQueryWrapper.kt      |  8 ++++++++
 .../firestore/internal/NativeQueryWrapper.kt      |  1 +
 .../firestore/internal/NativeQueryWrapper.kt      | 15 +++++++++++++++
 .../firestore/internal/NativeQueryWrapper.kt      |  5 +++++
 4 files changed, 29 insertions(+)

diff --git a/firebase-firestore/src/androidMain/kotlin/dev/gitlive/firebase/firestore/internal/NativeQueryWrapper.kt b/firebase-firestore/src/androidMain/kotlin/dev/gitlive/firebase/firestore/internal/NativeQueryWrapper.kt
index 5dc17f72b..390cea3cf 100644
--- a/firebase-firestore/src/androidMain/kotlin/dev/gitlive/firebase/firestore/internal/NativeQueryWrapper.kt
+++ b/firebase-firestore/src/androidMain/kotlin/dev/gitlive/firebase/firestore/internal/NativeQueryWrapper.kt
@@ -1,6 +1,8 @@
 package dev.gitlive.firebase.firestore.internal
 
 import com.google.android.gms.tasks.TaskExecutors
+import com.google.android.gms.tasks.Tasks
+import com.google.firebase.firestore.AggregateSource
 import com.google.firebase.firestore.FieldPath
 import com.google.firebase.firestore.MetadataChanges
 import com.google.firebase.firestore.Query
@@ -41,6 +43,12 @@ internal actual open class NativeQueryWrapper internal actual constructor(actual
     actual suspend fun get(source: Source): QuerySnapshot =
         QuerySnapshot(native.get(source.toAndroidSource()).await())
 
+    actual suspend fun count(): Long {
+        val aggregateQuery = native.count()
+        val task = aggregateQuery.get(AggregateSource.SERVER)
+        return Tasks.await(task).count
+    }
+
     actual fun where(filter: Filter) = native.where(filter.toAndroidFilter())
 
     private fun Filter.toAndroidFilter(): com.google.firebase.firestore.Filter = when (this) {
diff --git a/firebase-firestore/src/commonMain/kotlin/dev/gitlive/firebase/firestore/internal/NativeQueryWrapper.kt b/firebase-firestore/src/commonMain/kotlin/dev/gitlive/firebase/firestore/internal/NativeQueryWrapper.kt
index 5d1d1e67c..c212c086d 100644
--- a/firebase-firestore/src/commonMain/kotlin/dev/gitlive/firebase/firestore/internal/NativeQueryWrapper.kt
+++ b/firebase-firestore/src/commonMain/kotlin/dev/gitlive/firebase/firestore/internal/NativeQueryWrapper.kt
@@ -17,6 +17,7 @@ internal expect open class NativeQueryWrapper internal constructor(native: Nativ
     val snapshots: Flow<QuerySnapshot>
     fun snapshots(includeMetadataChanges: Boolean = false): Flow<QuerySnapshot>
     suspend fun get(source: Source = Source.DEFAULT): QuerySnapshot
+    suspend fun count(): Long
 
     fun where(filter: Filter): NativeQuery
 
diff --git a/firebase-firestore/src/iosMain/kotlin/dev/gitlive/firebase/firestore/internal/NativeQueryWrapper.kt b/firebase-firestore/src/iosMain/kotlin/dev/gitlive/firebase/firestore/internal/NativeQueryWrapper.kt
index 30f01b7f1..465c87555 100644
--- a/firebase-firestore/src/iosMain/kotlin/dev/gitlive/firebase/firestore/internal/NativeQueryWrapper.kt
+++ b/firebase-firestore/src/iosMain/kotlin/dev/gitlive/firebase/firestore/internal/NativeQueryWrapper.kt
@@ -1,5 +1,6 @@
 package dev.gitlive.firebase.firestore.internal
 
+import cocoapods.FirebaseFirestoreInternal.FIRAggregateSource
 import cocoapods.FirebaseFirestoreInternal.FIRFilter
 import dev.gitlive.firebase.firestore.Direction
 import dev.gitlive.firebase.firestore.EncodedFieldPath
@@ -13,7 +14,9 @@ import dev.gitlive.firebase.firestore.awaitResult
 import dev.gitlive.firebase.firestore.toException
 import kotlinx.coroutines.channels.awaitClose
 import kotlinx.coroutines.flow.callbackFlow
+import kotlinx.coroutines.suspendCancellableCoroutine
 import platform.Foundation.NSNull
+import kotlin.coroutines.resume
 
 internal actual open class NativeQueryWrapper internal actual constructor(actual open val native: NativeQuery) {
 
@@ -22,6 +25,18 @@ internal actual open class NativeQueryWrapper internal actual constructor(actual
     actual suspend fun get(source: Source) =
         QuerySnapshot(awaitResult { native.getDocumentsWithSource(source.toIosSource(), it) })
 
+    actual suspend fun count(): Long = suspendCancellableCoroutine { continuation ->
+        val aggregateQuery = native.count()
+
+        aggregateQuery.aggregationWithSource(FIRAggregateSource.FIRAggregateSourceServer) { snapshot, error ->
+            if (error != null) {
+                continuation.resume(0L)
+            } else {
+                continuation.resume(snapshot?.count?.longValue ?: 0L)
+            }
+        }
+    }
+
     actual val snapshots get() = callbackFlow {
         val listener = native.addSnapshotListener { snapshot, error ->
             snapshot?.let { trySend(QuerySnapshot(snapshot)) }
diff --git a/firebase-firestore/src/jsMain/kotlin/dev/gitlive/firebase/firestore/internal/NativeQueryWrapper.kt b/firebase-firestore/src/jsMain/kotlin/dev/gitlive/firebase/firestore/internal/NativeQueryWrapper.kt
index bc49b3437..d96df064c 100644
--- a/firebase-firestore/src/jsMain/kotlin/dev/gitlive/firebase/firestore/internal/NativeQueryWrapper.kt
+++ b/firebase-firestore/src/jsMain/kotlin/dev/gitlive/firebase/firestore/internal/NativeQueryWrapper.kt
@@ -33,6 +33,11 @@ internal actual open class NativeQueryWrapper internal actual constructor(actual
 
     actual suspend fun get(source: Source) = rethrow { QuerySnapshot(js.get(source).await()) }
 
+    actual suspend fun count(): Long = rethrow {
+        val snapshot = js.get(Source.DEFAULT).await()
+        snapshot.docs.size.toLong()
+    }
+
     actual fun limit(limit: Number) = query(
         js,
         dev.gitlive.firebase.firestore.externals.limit(limit),