Generate service trait from proto service definition (#2812)

Author: guilload

quickwit/Cargo.lock

@@ -5,6 +5,7 @@ members = [
   "quickwit-aws",
   "quickwit-cli",
   "quickwit-cluster",
+  "quickwit-codegen",
   "quickwit-common",
   "quickwit-config",
   "quickwit-control-plane",
@@ -99,6 +100,8 @@ opentelemetry-otlp = "0.11.0"
 pin-project-lite = "0.2.9"
 pnet = { version = "0.31.0", features = ["std"] }
 predicates = "2"
+prettyplease = "0.1.23"
+proc-macro2 = "1.0.50"
 prometheus = { version = "0.13", features = ["process"] }
 proptest = "1"
 prost = { version = "0.11.6", default-features = false, features = [
@@ -107,7 +110,7 @@ prost = { version = "0.11.6", default-features = false, features = [
 prost-build = "0.11.6"
 prost-types = "0.11.6"
 pulsar = { git = "https://github.com/quickwit-oss/pulsar-rs.git", rev = "f9eff04", default-features = false, features = ["compression", "tokio-runtime", "auth-oauth2"] }
-quickwit_elastic_api_generation = { git = "https://github.com/quickwit-oss/elasticsearch-rs", rev = "c157b19" }
+quote = "1.0.23"
 rand = "0.8"
 rand_distr = "0.4"
 rayon = "1"
@@ -152,6 +155,7 @@ sqlx = { version = "0.6", features = [
   "migrate",
   "time",
 ] }
+syn = "1.0.107"
 tabled = { version = "0.8", features = ["color"] }
 tempfile = "3"
 termcolor = "1"
@@ -195,6 +199,7 @@ quickwit-control-plane = { version = "0.4.0", path = "./quickwit-control-plane"
 quickwit-core = { version = "0.4.0", path = "./quickwit-core" }
 quickwit-directories = { version = "0.4.0", path = "./quickwit-directories" }
 quickwit-doc-mapper = { version = "0.4.0", path = "./quickwit-doc-mapper" }
+quickwit_elastic_api_generation = { git = "https://github.com/quickwit-oss/elasticsearch-rs", rev = "c157b19" }
 quickwit-grpc-clients = { version = "0.4.0", path = "./quickwit-grpc-clients" }
 quickwit-indexing = { version = "0.4.0", path = "./quickwit-indexing" }
 quickwit-ingest-api = { version = "0.4.0", path = "./quickwit-ingest-api" }

quickwit/Cargo.toml

@@ -0,0 +1,29 @@
+[package]
+name = "quickwit-codegen"
+version = "0.4.0"
+authors = ["Quickwit, Inc. <[email protected]>"]
+edition = "2021"
+license = "AGPL-3.0-or-later"                           # For a commercial, license, contact [email protected]
+description = "Generate traits, adapters, and gRPC clients/servers from proto files"
+repository = "https://github.com/quickwit-oss/quickwit"
+homepage = "https://quickwit.io/"
+documentation = "https://quickwit.io/docs/"
+
+[dependencies]
+anyhow = { workspace = true }
+prettyplease = { workspace = true }
+proc-macro2 = { workspace = true }
+prost = { workspace = true }
+prost-build = { workspace = true }
+quote = { workspace = true }
+syn = { workspace = true }
+tonic = { workspace = true }
+tonic-build = { workspace = true }
+
+[dev-dependencies]
+async-trait = { workspace = true }
+dyn-clone = { workspace = true }
+serde = { workspace = true }
+thiserror = { workspace = true }
+tokio = { workspace = true }
+utoipa = { workspace = true }

quickwit/quickwit-codegen/Cargo.toml

@@ -0,0 +1,122 @@
+// Copyright (C) 2023 Quickwit, Inc.
+//
+// Quickwit is offered under the AGPL v3.0 and as commercial software.
+// For commercial licensing, contact us at [email protected].
+//
+// AGPL:
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as
+// published by the Free Software Foundation, either version 3 of the
+// License, or (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program. If not, see <http://www.gnu.org/licenses/>.
+
+use std::path::Path;
+
+use proc_macro2::TokenStream;
+use prost_build::{Method, Service, ServiceGenerator};
+use quote::{quote, ToTokens};
+use syn::Ident;
+
+pub struct Codegen;
+
+impl Codegen {
+    pub fn run(proto: &Path, out_dir: &Path, result_path: &str) -> anyhow::Result<()> {
+        let mut prost_config = prost_build::Config::default();
+        prost_config
+            .protoc_arg("--experimental_allow_proto3_optional")
+            .type_attribute(".", "#[derive(Serialize, Deserialize, utoipa::ToSchema)]")
+            .out_dir(out_dir);
+
+        let service_generator = Box::new(QuickwitServiceGenerator::new(result_path));
+
+        prost_config
+            .service_generator(service_generator)
+            .compile_protos(&[proto], &["protos"])?;
+        Ok(())
+    }
+}
+
+struct QuickwitServiceGenerator {
+    result_path: String,
+    tonic_svc_generator: Box<dyn ServiceGenerator>,
+}
+
+impl QuickwitServiceGenerator {
+    fn new(result_path: &str) -> Self {
+        Self {
+            result_path: result_path.to_string(),
+            tonic_svc_generator: tonic_build::configure().service_generator(),
+        }
+    }
+}
+
+impl ServiceGenerator for QuickwitServiceGenerator {
+    fn generate(&mut self, service: Service, buf: &mut String) {
+        let tokens = generate_all(&service, &self.result_path);
+        let ast: syn::File = syn::parse2(tokens).expect("Tokenstream should be a valid Syn AST.");
+        let pretty_code = prettyplease::unparse(&ast);
+        buf.push_str(&pretty_code);
+
+        self.tonic_svc_generator.generate(service, buf)
+    }
+
+    fn finalize(&mut self, buf: &mut String) {
+        self.tonic_svc_generator.finalize(buf);
+    }
+}
+
+fn generate_all(service: &Service, result_path: &str) -> TokenStream {
+    let service_trait_name = quote::format_ident!("{}", service.name);
+    let result_path = syn::parse_str::<syn::Path>(result_path)
+        .expect("Result path should be a valid result path such as `crate::Result`.");
+
+    let service_trait = generate_service_trait(&service_trait_name, &service.methods, &result_path);
+
+    quote! {
+        #service_trait
+    }
+}
+
+fn generate_service_trait(
+    trait_name: &Ident,
+    methods: &[Method],
+    result_path: &syn::Path,
+) -> TokenStream {
+    let trait_methods = generate_service_trait_methods(methods, result_path);
+
+    quote! {
+        #[async_trait]
+        pub trait #trait_name: std::fmt::Debug + dyn_clone::DynClone + Send + Sync + 'static {
+
+            #trait_methods
+        }
+
+        dyn_clone::clone_trait_object!(#trait_name);
+    }
+}
+
+fn generate_service_trait_methods(methods: &[Method], result_path: &syn::Path) -> TokenStream {
+    let mut stream = TokenStream::new();
+    for method in methods {
+        let method_name = quote::format_ident!("{}", method.name);
+        let request_type = syn::parse_str::<syn::Path>(&method.input_type)
+            .unwrap()
+            .to_token_stream();
+        let response_type = syn::parse_str::<syn::Path>(&method.output_type)
+            .unwrap()
+            .to_token_stream();
+
+        let method = quote! {
+            async fn #method_name(&mut self, request: #request_type) -> #result_path<#response_type>;
+        };
+        stream.extend(method);
+    }
+    stream
+}

quickwit/quickwit-codegen/src/codegen.rs

@@ -0,0 +1,22 @@
+// Copyright (C) 2023 Quickwit, Inc.
+//
+// Quickwit is offered under the AGPL v3.0 and as commercial software.
+// For commercial licensing, contact us at [email protected].
+//
+// AGPL:
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as
+// published by the Free Software Foundation, either version 3 of the
+// License, or (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program. If not, see <http://www.gnu.org/licenses/>.
+
+mod codegen;
+
+pub use codegen::Codegen;

quickwit/quickwit-codegen/src/lib.rs

@@ -0,0 +1,62 @@
+// Copyright (C) 2023 Quickwit, Inc.
+//
+// Quickwit is offered under the AGPL v3.0 and as commercial software.
+// For commercial licensing, contact us at [email protected].
+//
+// AGPL:
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as
+// published by the Free Software Foundation, either version 3 of the
+// License, or (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program. If not, see <http://www.gnu.org/licenses/>.
+
+use std::path::Path;
+
+use async_trait::async_trait;
+use quickwit_codegen::Codegen;
+
+mod hello;
+
+#[tokio::test]
+async fn test_hello_codegen() {
+    let proto = Path::new("tests/hello/hello.proto");
+    let out_dir = Path::new("tests/hello/");
+
+    Codegen::run(proto, out_dir, "crate::hello::HelloResult").unwrap();
+
+    use crate::hello::{Hello, HelloRequest, HelloResponse};
+
+    #[derive(Debug, Clone)]
+    struct HelloImpl;
+
+    #[async_trait]
+    impl Hello for HelloImpl {
+        async fn hello(
+            &mut self,
+            request: HelloRequest,
+        ) -> crate::hello::HelloResult<HelloResponse> {
+            Ok(HelloResponse {
+                message: format!("Hello, {}!", request.name),
+            })
+        }
+    }
+    let mut hello = HelloImpl;
+    assert_eq!(
+        hello
+            .hello(HelloRequest {
+                name: "World".to_string()
+            })
+            .await
+            .unwrap(),
+        HelloResponse {
+            message: "Hello, World!".to_string()
+        }
+    );
+}

quickwit/quickwit-codegen/tests/codegen.rs

@@ -0,0 +1,37 @@
+// Copyright (C) 2023 Quickwit, Inc.
+//
+// Quickwit is offered under the AGPL v3.0 and as commercial software.
+// For commercial licensing, contact us at [email protected].
+//
+// AGPL:
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as
+// published by the Free Software Foundation, either version 3 of the
+// License, or (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program. If not, see <http://www.gnu.org/licenses/>.
+
+// Service errors have to be handwritten before codegen.
+#[derive(Debug, thiserror::Error)]
+pub enum HelloError {
+    #[error("Internal error: {0}")]
+    InternalError(String),
+    #[error("Transport error: {0}")]
+    TransportError(#[from] tonic::Status),
+}
+
+// Service errors must implement `From<tonic::Status>` and `Into<tonic::Status>`.
+impl From<HelloError> for tonic::Status {
+    fn from(val: HelloError) -> Self {
+        match val {
+            HelloError::InternalError(message) => tonic::Status::internal(message),
+            HelloError::TransportError(status) => status,
+        }
+    }
+}

quickwit/quickwit-codegen/tests/hello/error.rs

@@ -0,0 +1,34 @@
+// Copyright (C) 2023 Quickwit, Inc.
+//
+// Quickwit is offered under the AGPL v3.0 and as commercial software.
+// For commercial licensing, contact us at [email protected].
+//
+// AGPL:
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as
+// published by the Free Software Foundation, either version 3 of the
+// License, or (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program. If not, see <http://www.gnu.org/licenses/>.
+
+syntax = "proto3";
+
+package hello;
+
+message HelloRequest {
+    string name = 1;
+}
+
+message HelloResponse {
+    string message = 1;
+}
+
+service Hello {
+    rpc Hello(HelloRequest) returns (HelloResponse);
+}