From d59c7305b0387d3ebf5966326a31c98e1d0eb4e4 Mon Sep 17 00:00:00 2001
From: Alfredo Garcia <oxarbitrage@gmail.com>
Date: Fri, 9 Feb 2024 20:41:26 -0300
Subject: [PATCH] add(grpc): `get_results` (#8255)

* add `get_results` grpc call

* nitpick
---
 zebra-grpc/build.rs            |  4 ++-
 zebra-grpc/proto/scanner.proto | 29 +++++++++++++++++-
 zebra-grpc/src/server.rs       | 56 ++++++++++++++++++++++++++++++++--
 3 files changed, 85 insertions(+), 4 deletions(-)

diff --git a/zebra-grpc/build.rs b/zebra-grpc/build.rs
index f6c46fa24..0deb3a637 100644
--- a/zebra-grpc/build.rs
+++ b/zebra-grpc/build.rs
@@ -1,6 +1,8 @@
 //! Compile proto files
 
 fn main() -> Result<(), Box<dyn std::error::Error>> {
-    tonic_build::compile_protos("proto/scanner.proto")?;
+    tonic_build::configure()
+        .btree_map(["."])
+        .compile(&["proto/scanner.proto"], &[""])?;
     Ok(())
 }
diff --git a/zebra-grpc/proto/scanner.proto b/zebra-grpc/proto/scanner.proto
index 72ff40882..9eed480fd 100644
--- a/zebra-grpc/proto/scanner.proto
+++ b/zebra-grpc/proto/scanner.proto
@@ -16,6 +16,9 @@ service Scanner {
     // Deletes a set of keys and their results from the scanner.
     // This request stop the scanner from scanning blocks for the these keys.
     rpc DeleteKeys(DeleteKeysRequest) returns (Empty);
+
+    // Get all data we have stored for the given keys.
+    rpc GetResults(GetResultsRequest) returns (GetResultsResponse);
 }
 
 // A response to a GetInfo call.
@@ -34,4 +37,28 @@ message ClearResultsRequest {
 message DeleteKeysRequest {
     // Keys to delete from scanner.
     repeated string keys = 1;
-}
\ No newline at end of file
+}
+
+// A request for getting results for a set of keys.
+message GetResultsRequest {
+    // Keys for which to get results.
+    repeated string keys = 1;
+}
+
+// A set of responses for each provided key of a GetResults call.
+message GetResultsResponse {
+    // Results for each key.
+    map<string, Results> results = 1;
+}
+
+// A result for a single key.
+message Results {
+    // A height, transaction id map
+    map<uint32, TransactionHash> transactions = 1;
+}
+
+// A vector of transaction hashes
+message TransactionHash {
+    // A transaction id hash
+    repeated string hash = 1;
+}
diff --git a/zebra-grpc/src/server.rs b/zebra-grpc/src/server.rs
index c1d3e0bbd..65e2b8a98 100644
--- a/zebra-grpc/src/server.rs
+++ b/zebra-grpc/src/server.rs
@@ -1,6 +1,6 @@
 //! The gRPC server implementation
 
-use std::net::SocketAddr;
+use std::{collections::BTreeMap, net::SocketAddr};
 
 use futures_util::future::TryFutureExt;
 use tonic::{transport::Server, Response, Status};
@@ -12,7 +12,8 @@ use zebra_node_services::scan_service::{
 
 use crate::scanner::{
     scanner_server::{Scanner, ScannerServer},
-    ClearResultsRequest, DeleteKeysRequest, Empty, InfoReply,
+    ClearResultsRequest, DeleteKeysRequest, Empty, GetResultsRequest, GetResultsResponse,
+    InfoReply, Results, TransactionHash,
 };
 
 type BoxError = Box<dyn std::error::Error + Send + Sync + 'static>;
@@ -120,6 +121,57 @@ where
 
         Ok(Response::new(Empty {}))
     }
+
+    async fn get_results(
+        &self,
+        request: tonic::Request<GetResultsRequest>,
+    ) -> Result<Response<GetResultsResponse>, Status> {
+        let keys = request.into_inner().keys;
+
+        if keys.is_empty() {
+            let msg = "must provide at least 1 key to get results";
+            return Err(Status::invalid_argument(msg));
+        }
+
+        let ScanServiceResponse::Results(response) = self
+            .scan_service
+            .clone()
+            .ready()
+            .and_then(|service| service.call(ScanServiceRequest::Results(keys.clone())))
+            .await
+            .map_err(|err| Status::unknown(format!("scan service returned error: {err}")))?
+        else {
+            return Err(Status::unknown(
+                "scan service returned an unexpected response",
+            ));
+        };
+
+        // If there are no results for a key, we still want to return it with empty results.
+        let empty_map = BTreeMap::new();
+
+        let results = keys
+            .into_iter()
+            .map(|key| {
+                let values = response.get(&key).unwrap_or(&empty_map);
+
+                // Skip heights with no transactions, they are scanner markers and should not be returned.
+                let transactions = Results {
+                    transactions: values
+                        .iter()
+                        .filter(|(_, transactions)| !transactions.is_empty())
+                        .map(|(height, transactions)| {
+                            let txs = transactions.iter().map(ToString::to_string).collect();
+                            (height.0, TransactionHash { hash: txs })
+                        })
+                        .collect(),
+                };
+
+                (key, transactions)
+            })
+            .collect::<BTreeMap<_, _>>();
+
+        Ok(Response::new(GetResultsResponse { results }))
+    }
 }
 
 /// Initializes the zebra-scan gRPC server