Batch query (#77)

* add api client with grpc kotlin * get api client setup with batch queries * just take a list of requests * write a test for it
xmtp · May 31, 2023 · 71c8035 · 71c8035
1 parent effc4f5
commit 71c8035
Show file tree

Hide file tree

Showing 6 changed files with 68 additions and 1 deletion.
diff --git a/library/build.gradle b/library/build.gradle
@@ -80,7 +80,7 @@ dependencies {
     implementation 'org.web3j:crypto:5.0.0'
     implementation "net.java.dev.jna:jna:5.13.0@aar"
     api 'com.google.protobuf:protobuf-kotlin-lite:3.22.3'
-    api 'org.xmtp:proto-kotlin:3.23.2'
+    api 'org.xmtp:proto-kotlin:3.24.1'
 
     testImplementation 'junit:junit:4.13.2'
     androidTestImplementation 'app.cash.turbine:turbine:0.12.1'

diff --git a/library/src/androidTest/java/org/xmtp/android/library/LocalInstrumentedTest.kt b/library/src/androidTest/java/org/xmtp/android/library/LocalInstrumentedTest.kt
@@ -20,6 +20,7 @@ import org.xmtp.android.library.messages.generate
 import org.xmtp.android.library.messages.secp256K1Uncompressed
 import org.xmtp.android.library.messages.toPublicKeyBundle
 import org.xmtp.android.library.messages.walletAddress
+import org.xmtp.proto.message.api.v1.MessageApiOuterClass.QueryRequest
 import org.xmtp.proto.message.contents.Contact
 import org.xmtp.proto.message.contents.PrivateKeyOuterClass
 import java.util.Date
@@ -221,4 +222,29 @@ class LocalInstrumentedTest {
         val keys = PrivateKeyBundleV1Builder.buildFromBundle(bytes)
         Client().buildFrom(bundle = keys, options = options)
     }
+
+    @Test
+    fun testBatchQuery() {
+        val alice = PrivateKeyBuilder()
+        val identity = PrivateKey.newBuilder().build().generate()
+        val authorized = alice.createIdentity(identity)
+        val authToken = authorized.createAuthToken()
+        val api = GRPCApiClient(environment = XMTPEnvironment.LOCAL, secure = false)
+        api.setAuthToken(authToken)
+        val encryptedBundle = authorized.toBundle.encrypted(alice)
+        val envelope = Envelope.newBuilder().also {
+            it.contentTopic = Topic.userPrivateStoreKeyBundle(authorized.address).description
+            it.timestampNs = Date().time * 1_000_000
+            it.message = encryptedBundle.toByteString()
+        }.build()
+        runBlocking {
+            api.publish(envelopes = listOf(envelope))
+        }
+        Thread.sleep(2_000)
+        val request = QueryRequest.newBuilder().addContentTopics(Topic.userPrivateStoreKeyBundle(authorized.address).description).build()
+        val result =
+            runBlocking { api.batchQuery(requests = listOf(request)) }
+
+        assertEquals(result.responsesOrBuilderList.size, 1)
+    }
 }
diff --git a/library/src/androidTest/java/org/xmtp/android/library/TestHelpers.kt b/library/src/androidTest/java/org/xmtp/android/library/TestHelpers.kt
@@ -116,6 +116,14 @@ class FakeApiClient : ApiClient {
         return query(topic = topic, pagination = pagination).envelopesList
     }
 
+    override suspend fun batchQuery(requests: List<MessageApiOuterClass.QueryRequest>): MessageApiOuterClass.BatchQueryResponse {
+        val response = query(requests.first().getContentTopics(0))
+
+        return MessageApiOuterClass.BatchQueryResponse.newBuilder().also {
+            it.addResponses(response)
+        }.build()
+    }
+
     override suspend fun query(
         topic: String,
         pagination: Pagination?,

diff --git a/library/src/main/java/org/xmtp/android/library/ApiClient.kt b/library/src/main/java/org/xmtp/android/library/ApiClient.kt
@@ -10,6 +10,8 @@ import org.xmtp.android.library.messages.Pagination
 import org.xmtp.android.library.messages.Topic
 import org.xmtp.proto.message.api.v1.MessageApiGrpcKt
 import org.xmtp.proto.message.api.v1.MessageApiOuterClass
+import org.xmtp.proto.message.api.v1.MessageApiOuterClass.BatchQueryRequest
+import org.xmtp.proto.message.api.v1.MessageApiOuterClass.BatchQueryResponse
 import org.xmtp.proto.message.api.v1.MessageApiOuterClass.Cursor
 import org.xmtp.proto.message.api.v1.MessageApiOuterClass.Envelope
 import org.xmtp.proto.message.api.v1.MessageApiOuterClass.PublishRequest
@@ -29,6 +31,7 @@ interface ApiClient {
     ): QueryResponse
 
     suspend fun queryTopic(topic: Topic, pagination: Pagination? = null): QueryResponse
+    suspend fun batchQuery(requests: List<QueryRequest>): BatchQueryResponse
     suspend fun envelopes(topic: String, pagination: Pagination? = null): List<Envelope>
     suspend fun publish(envelopes: List<Envelope>): PublishResponse
     suspend fun subscribe(topics: List<String>): Flow<Envelope>
@@ -126,6 +129,22 @@ data class GRPCApiClient(override val environment: XMTPEnvironment, val secure:
         return query(topic.description, pagination)
     }
 
+    override suspend fun batchQuery(
+        requests: List<QueryRequest>,
+    ): BatchQueryResponse {
+        val batchRequest = BatchQueryRequest.newBuilder().addAllRequests(requests).build()
+        val headers = Metadata()
+
+        authToken?.let { token ->
+            headers.put(AUTHORIZATION_HEADER_KEY, "Bearer $token")
+        }
+        headers.put(CLIENT_VERSION_HEADER_KEY, Constants.VERSION)
+        if (appVersion != null) {
+            headers.put(APP_VERSION_HEADER_KEY, appVersion)
+        }
+        return client.batchQuery(batchRequest, headers = headers)
+    }
+
     override suspend fun publish(envelopes: List<Envelope>): PublishResponse {
         val request = PublishRequest.newBuilder().addAllEnvelopes(envelopes).build()
         val headers = Metadata()

diff --git a/library/src/main/java/org/xmtp/android/library/Client.kt b/library/src/main/java/org/xmtp/android/library/Client.kt
@@ -32,6 +32,8 @@ import org.xmtp.android.library.messages.toPublicKeyBundle
 import org.xmtp.android.library.messages.toV2
 import org.xmtp.android.library.messages.walletAddress
 import org.xmtp.proto.message.api.v1.MessageApiOuterClass
+import org.xmtp.proto.message.api.v1.MessageApiOuterClass.BatchQueryResponse
+import org.xmtp.proto.message.api.v1.MessageApiOuterClass.QueryRequest
 import java.nio.charset.StandardCharsets
 import java.text.SimpleDateFormat
 import java.time.Instant
@@ -240,6 +242,10 @@ class Client() {
         return apiClient.queryTopic(topic = topic, pagination = pagination)
     }
 
+    suspend fun batchQuery(requests: List<QueryRequest>): BatchQueryResponse {
+        return apiClient.batchQuery(requests)
+    }
+
     suspend fun subscribe(topics: List<String>): Flow<Envelope> {
         return apiClient.subscribe(topics = topics)
     }

diff --git a/library/src/test/java/org/xmtp/android/library/TestHelpers.kt b/library/src/test/java/org/xmtp/android/library/TestHelpers.kt
@@ -105,6 +105,14 @@ class FakeApiClient : ApiClient {
         return query(topic = topic.description, pagination)
     }
 
+    override suspend fun batchQuery(requests: List<MessageApiOuterClass.QueryRequest>): MessageApiOuterClass.BatchQueryResponse {
+        val response = query(requests.first().getContentTopics(0))
+
+        return MessageApiOuterClass.BatchQueryResponse.newBuilder().also {
+            it.addResponses(response)
+        }.build()
+    }
+
     suspend fun send(envelope: Envelope) {
         stream.emit(envelope)
     }