From 37dc85a29e18afdf991313fd86ece7f09ce72181 Mon Sep 17 00:00:00 2001
From: LFC <bayinamine@gmail.com>
Date: Wed, 9 Nov 2022 19:13:01 +0800
Subject: [PATCH] feat: add gRPC reflection service (#431)

Co-authored-by: luofucong <luofucong@greptime.com>
---
 Cargo.lock               | 16 ++++++++++++++++
 src/api/build.rs         |  4 ++++
 src/api/src/v1.rs        |  2 ++
 src/servers/Cargo.toml   |  1 +
 src/servers/src/error.rs |  9 ++++++++-
 src/servers/src/grpc.rs  |  9 ++++++++-
 6 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 5703faee5d..85f85fb66a 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -5069,6 +5069,7 @@ dependencies = [
  "tokio-stream",
  "tokio-test",
  "tonic",
+ "tonic-reflection",
  "tower",
  "tower-http",
 ]
@@ -5969,6 +5970,21 @@ dependencies = [
  "syn",
 ]
 
+[[package]]
+name = "tonic-reflection"
+version = "0.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0455f730d540a1484bffc3c55c94100b18a662597b982c2e9073f2c55c602616"
+dependencies = [
+ "bytes",
+ "prost 0.11.0",
+ "prost-types 0.11.1",
+ "tokio",
+ "tokio-stream",
+ "tonic",
+ "tonic-build",
+]
+
 [[package]]
 name = "tower"
 version = "0.4.13"
diff --git a/src/api/build.rs b/src/api/build.rs
index e15f2f80bc..4f398f2d71 100644
--- a/src/api/build.rs
+++ b/src/api/build.rs
@@ -1,5 +1,9 @@
+use std::path::PathBuf;
+
 fn main() {
+    let default_out_dir = PathBuf::from(std::env::var("OUT_DIR").unwrap());
     tonic_build::configure()
+        .file_descriptor_set_path(default_out_dir.join("greptime_fd.bin"))
         .compile(
             &[
                 "greptime/v1/insert.proto",
diff --git a/src/api/src/v1.rs b/src/api/src/v1.rs
index 1c5ceaade8..588d28140a 100644
--- a/src/api/src/v1.rs
+++ b/src/api/src/v1.rs
@@ -1,6 +1,8 @@
 #![allow(clippy::derive_partial_eq_without_eq)]
 tonic::include_proto!("greptime.v1");
 
+pub const GREPTIME_FD_SET: &[u8] = tonic::include_file_descriptor_set!("greptime_fd");
+
 pub mod codec {
     tonic::include_proto!("greptime.v1.codec");
 }
diff --git a/src/servers/Cargo.toml b/src/servers/Cargo.toml
index f17575a234..e62b3abb09 100644
--- a/src/servers/Cargo.toml
+++ b/src/servers/Cargo.toml
@@ -41,6 +41,7 @@ table = { path = "../table" }
 tokio = { version = "1.20", features = ["full"] }
 tokio-stream = { version = "0.1", features = ["net"] }
 tonic = "0.8"
+tonic-reflection = "0.5"
 tower = { version = "0.4", features = ["full"] }
 tower-http = { version = "0.3", features = ["full"] }
 
diff --git a/src/servers/src/error.rs b/src/servers/src/error.rs
index c2d0791657..39f46de3d1 100644
--- a/src/servers/src/error.rs
+++ b/src/servers/src/error.rs
@@ -155,6 +155,12 @@ pub enum Error {
 
     #[snafu(display("Failed to decode region id, source: {}", source))]
     DecodeRegionId { source: api::DecodeError },
+
+    #[snafu(display("Failed to build gRPC reflection service, source: {}", source))]
+    GrpcReflectionService {
+        source: tonic_reflection::server::Error,
+        backtrace: Backtrace,
+    },
 }
 
 pub type Result<T> = std::result::Result<T, Error>;
@@ -172,7 +178,8 @@ impl ErrorExt for Error {
             | StartGrpc { .. }
             | AlreadyStarted { .. }
             | InvalidPromRemoteReadQueryResult { .. }
-            | TcpBind { .. } => StatusCode::Internal,
+            | TcpBind { .. }
+            | GrpcReflectionService { .. } => StatusCode::Internal,
 
             InsertScript { source, .. }
             | ExecuteScript { source, .. }
diff --git a/src/servers/src/grpc.rs b/src/servers/src/grpc.rs
index ddddeab180..b7a4599610 100644
--- a/src/servers/src/grpc.rs
+++ b/src/servers/src/grpc.rs
@@ -16,7 +16,7 @@ use tokio::sync::Mutex;
 use tokio_stream::wrappers::TcpListenerStream;
 use tonic::{Request, Response, Status};
 
-use crate::error::{AlreadyStartedSnafu, Result, StartGrpcSnafu, TcpBindSnafu};
+use crate::error::{self, AlreadyStartedSnafu, Result, StartGrpcSnafu, TcpBindSnafu};
 use crate::grpc::handler::BatchHandler;
 use crate::query_handler::{GrpcAdminHandlerRef, GrpcQueryHandlerRef};
 use crate::server::Server;
@@ -104,9 +104,16 @@ impl Server for GrpcServer {
             (listener, addr)
         };
 
+        let reflection_service = tonic_reflection::server::Builder::configure()
+            .register_encoded_file_descriptor_set(api::v1::GREPTIME_FD_SET)
+            .with_service_name("greptime.v1.Greptime")
+            .build()
+            .context(error::GrpcReflectionServiceSnafu)?;
+
         // Would block to serve requests.
         tonic::transport::Server::builder()
             .add_service(self.create_service())
+            .add_service(reflection_service)
             .serve_with_incoming_shutdown(TcpListenerStream::new(listener), rx.map(drop))
             .await
             .context(StartGrpcSnafu)?;