WIP: Use a request builder pattern to allow clients to set headers and other per-request info

crates/twirp-build/Cargo.toml

@@ -18,4 +18,3 @@ prost-build = "0.13"
 prettyplease = { version = "0.2" }
 quote = "1.0"
 syn = "2.0"
-proc-macro2 = "1.0"

crates/twirp-build/src/lib.rs

@@ -166,19 +166,25 @@ impl prost_build::ServiceGenerator for ServiceGenerator {
         let mut client_methods = Vec::with_capacity(service.methods.len());
         for m in &service.methods {
             let name = &m.name;
+            let name_request = format_ident!("{}_request", name);
             let input_type = &m.input_type;
             let output_type = &m.output_type;
             let request_path = format!("{}/{}", service.fqn, m.proto_name);
 
             client_trait_methods.push(quote! {
-                async fn #name(&self, req: #input_type) -> Result<#output_type, twirp::ClientError>;
+                async fn #name(&self, req: #input_type) -> Result<#output_type, twirp::ClientError> {
+                    self.#name_request(req)?.send().await
+                }
+            });
+            client_trait_methods.push(quote! {
+                fn #name_request(&self, req: #input_type) -> Result<twirp::RequestBuilder<#input_type, #output_type>, twirp::ClientError>;
             });
 
             client_methods.push(quote! {
-                async fn #name(&self, req: #input_type) -> Result<#output_type, twirp::ClientError> {
-                    self.request(#request_path, req).await
+                fn #name_request(&self, req: #input_type) -> Result<twirp::RequestBuilder<#input_type, #output_type>, twirp::ClientError> {
+                    self.request(#request_path, req)
                 }
-            })
+            });
         }
         let client_trait = quote! {
             #[twirp::async_trait::async_trait]

crates/twirp/src/client.rs

@@ -2,6 +2,7 @@ use std::sync::Arc;
 use std::vec;
 
 use async_trait::async_trait;
+use http::{HeaderName, HeaderValue};
 use reqwest::header::{InvalidHeaderValue, CONTENT_TYPE};
 use reqwest::StatusCode;
 use thiserror::Error;
@@ -155,23 +156,12 @@ impl Client {
         }
     }
 
-    /// Make an HTTP twirp request.
-    pub async fn request<I, O>(&self, path: &str, body: I) -> Result<O>
+    /// Executes a `Request`.
+    pub(super) async fn execute<O>(&self, req: reqwest::Request) -> Result<O>
     where
-        I: prost::Message,
         O: prost::Message + Default,
     {
-        let mut url = self.inner.base_url.join(path)?;
-        if let Some(host) = &self.host {
-            url.set_host(Some(host))?
-        };
-        let path = url.path().to_string();
-        let req = self
-            .http_client
-            .post(url)
-            .header(CONTENT_TYPE, CONTENT_TYPE_PROTOBUF)
-            .body(serialize_proto_message(body))
-            .build()?;
+        let path = req.url().path().to_string();
 
         // Create and execute the middleware handlers
         let next = Next::new(&self.http_client, &self.inner.middlewares);
@@ -204,6 +194,68 @@ impl Client {
             }),
         }
     }
+
+    /// Start building a `Request` with a path and a request body.
+    ///
+    /// Returns a `RequestBuilder`, which will allow setting headers before sending.
+    pub fn request<I, O>(&self, path: &str, body: I) -> Result<RequestBuilder<I, O>>
+    where
+        I: prost::Message,
+        O: prost::Message + Default,
+    {
+        let mut url = self.inner.base_url.join(path)?;
+        if let Some(host) = &self.host {
+            url.set_host(Some(host))?
+        };
+
+        let req = self
+            .http_client
+            .post(url)
+            .header(CONTENT_TYPE, CONTENT_TYPE_PROTOBUF)
+            .body(serialize_proto_message(body));
+        Ok(RequestBuilder::new(self.clone(), req))
+    }
+}
+
+pub struct RequestBuilder<I, O>
+where
+    O: prost::Message + Default,
+{
+    client: Client,
+    inner: reqwest::RequestBuilder,
+    _input: std::marker::PhantomData<I>,
+    _output: std::marker::PhantomData<O>,
+}
+
+impl<I, O> RequestBuilder<I, O>
+where
+    O: prost::Message + Default,
+{
+    pub fn new(client: Client, inner: reqwest::RequestBuilder) -> Self {
+        Self {
+            client,
+            inner,
+            _input: std::marker::PhantomData,
+            _output: std::marker::PhantomData,
+        }
+    }
+
+    /// Add a `Header` to this Request.
+    pub fn header<K, V>(mut self, key: K, value: V) -> RequestBuilder<I, O>
+    where
+        HeaderName: TryFrom<K>,
+        <HeaderName as TryFrom<K>>::Error: Into<http::Error>,
+        HeaderValue: TryFrom<V>,
+        <HeaderValue as TryFrom<V>>::Error: Into<http::Error>,
+    {
+        self.inner = self.inner.header(key, value);
+        self
+    }
+
+    pub async fn send(self) -> Result<O> {
+        let req = self.inner.build()?;
+        self.client.execute(req).await
+    }
 }
 
 // This concept of reqwest middleware is taken pretty much directly from:

crates/twirp/src/lib.rs

@@ -10,7 +10,7 @@ pub mod test;
 #[doc(hidden)]
 pub mod details;
 
-pub use client::{Client, ClientBuilder, ClientError, Middleware, Next, Result};
+pub use client::{Client, ClientBuilder, ClientError, Middleware, Next, RequestBuilder, Result};
 pub use context::Context;
 pub use error::*; // many constructors like `invalid_argument()`
 pub use http::Extensions;

crates/twirp/src/test.rs

@@ -121,7 +121,7 @@ pub trait TestApiClient {
 #[async_trait]
 impl TestApiClient for Client {
     async fn ping(&self, req: PingRequest) -> Result<PingResponse> {
-        self.request("test.TestAPI/Ping", req).await
+        self.request("test.TestAPI/Ping", req)?.send().await
     }
 
     async fn boom(&self, _req: PingRequest) -> Result<PingResponse> {

example/src/bin/client.rs

@@ -38,6 +38,14 @@ pub async fn main() -> Result<(), GenericError> {
         .await;
     eprintln!("{:?}", resp);
 
+    let resp = client
+        .with_host("localhost")
+        .make_hat_request(MakeHatRequest { inches: 1 })?
+        .header("x-custom-header", "a")
+        .send()
+        .await?;
+    eprintln!("{:?}", resp);
+
     Ok(())
 }
 
@@ -69,23 +77,39 @@ impl Middleware for PrintResponseHeaders {
     }
 }
 
+// NOTE: This is just to demonstrate manually implementing the client trait. You don't need to do this as this code will
+// be generated for you by twirp-build.
+//
+// This is here so that we can visualize changes to the generated client code
 #[allow(dead_code)]
 #[derive(Debug)]
 struct MockHaberdasherApiClient;
 
 #[async_trait]
 impl HaberdasherApiClient for MockHaberdasherApiClient {
-    async fn make_hat(
+    fn make_hat_request(
         &self,
         _req: MakeHatRequest,
-    ) -> Result<MakeHatResponse, twirp::client::ClientError> {
+    ) -> Result<twirp::RequestBuilder<MakeHatRequest, MakeHatResponse>, twirp::ClientError> {
+        todo!()
+    }
+    // implementing this one is optional
+    async fn make_hat(&self, _req: MakeHatRequest) -> Result<MakeHatResponse, twirp::ClientError> {
         todo!()
     }
 
+    fn get_status_request(
+        &self,
+        _req: GetStatusRequest,
+    ) -> Result<twirp::RequestBuilder<GetStatusRequest, GetStatusResponse>, twirp::ClientError>
+    {
+        todo!()
+    }
+    // implementing this one is optional
     async fn get_status(
         &self,
         _req: GetStatusRequest,
-    ) -> Result<GetStatusResponse, twirp::client::ClientError> {
+    ) -> Result<GetStatusResponse, twirp::ClientError> {
         todo!()
     }
 }