batch count

komp-app/src/grpc.rs

@@ -1244,6 +1244,34 @@ impl GrpcClient {
         self.search_client.get_filtered_row(request).await
     }
 
+    pub async fn count_table_rows_batch(
+        &mut self,
+        profile_name: String,
+        rows: Vec<(String, String, String)>,
+    ) -> Vec<std::result::Result<u64, String>> {
+        let mut counts = Vec::with_capacity(rows.len());
+        for chunk in rows.chunks(64) {
+            let result = async {
+                let request = self.authenticated_request(common::proto::komp_ac::search::BatchCountRequest {
+                    requests: chunk.iter().map(|(table, column, value)| SearchRequest {
+                        profile_name: profile_name.clone(),
+                        table_name: Some(table.clone()),
+                        must: vec![ColumnConstraint {
+                            column: column.clone(), query: value.clone(), mode: MatchMode::Exact as i32,
+                        }],
+                        ..Default::default()
+                    }).collect(),
+                })?;
+                self.search_client.batch_count(request).await
+            }.await;
+            match result {
+                Ok(results) => counts.extend(results),
+                Err(error) => counts.extend((0..chunk.len()).map(|_| Err(error.to_string()))),
+            }
+        }
+        counts
+    }
+
     pub async fn count_table_rows(
         &mut self,
         profile_name: String,

komp-app/src/search.rs

@@ -28,7 +28,54 @@ impl SearchGrpc {
         Ok(self.client.get_filtered_row(request).await?.into_inner())
     }
 
+    pub async fn batch_count(
+        &mut self,
+        request: Request<common::proto::komp_ac::search::BatchCountRequest>,
+    ) -> Result<Vec<std::result::Result<u64, String>>> {
+        let expected = request.get_ref().requests.len();
+        let response = self.client.batch_count(request).await?.into_inner();
+        decode_batch_counts(response, expected)
+    }
+
     pub async fn count(&mut self, request: Request<SearchRequest>) -> Result<SearchCountResponse> {
         Ok(self.client.count(request).await?.into_inner())
     }
 }
+
+fn decode_batch_counts(
+    response: common::proto::komp_ac::search::BatchCountResponse,
+    expected: usize,
+) -> Result<Vec<std::result::Result<u64, String>>> {
+    use common::proto::komp_ac::search::batch_count_result::Outcome;
+    anyhow::ensure!(response.results.len() == expected, "Server returned an incomplete count batch");
+    Ok(response.results.into_iter().map(|result| match result.outcome {
+        Some(Outcome::Count(count)) => Ok(count),
+        Some(Outcome::Error(error)) => Err(error),
+        None => Err("Server omitted a count result".to_string()),
+    }).collect())
+}
+
+#[cfg(test)]
+mod batch_count_tests {
+    use super::*;
+    use common::proto::komp_ac::search::{BatchCountResponse, BatchCountResult, batch_count_result::Outcome};
+
+    #[test]
+    fn batch_counts_preserve_zero_errors_and_missing_results() {
+        let response = BatchCountResponse { results: vec![
+            BatchCountResult { outcome: Some(Outcome::Count(0)) },
+            BatchCountResult { outcome: Some(Outcome::Error("denied".into())) },
+            BatchCountResult { outcome: None },
+            BatchCountResult { outcome: Some(Outcome::Count(7)) },
+        ] };
+        let counts = decode_batch_counts(response, 4).unwrap();
+        assert_eq!(counts[0], Ok(0));
+        assert_eq!(counts[1], Err("denied".into()));
+        assert!(counts[2].is_err());
+        assert_eq!(counts[3], Ok(7));
+        assert!(decode_batch_counts(BatchCountResponse { results: Vec::new() }, 1).is_err());
+        assert!(decode_batch_counts(BatchCountResponse { results: vec![
+            BatchCountResult { outcome: Some(Outcome::Count(0)) },
+        ] }, 0).is_err());
+    }
+}