From 43885de0f6102daa971b907a6544ff1ab99f067d Mon Sep 17 00:00:00 2001
From: Uwe Pachler <uwe.pachler@gmail.com>
Date: Wed, 9 Oct 2024 18:50:45 +0200
Subject: [PATCH] add async post_hook function support

---
 progenitor-client/src/progenitor_client.rs |  8 ++++++++
 progenitor-impl/src/lib.rs                 | 11 +++++++++++
 progenitor-impl/src/method.rs              |  9 +++++++++
 3 files changed, 28 insertions(+)

diff --git a/progenitor-client/src/progenitor_client.rs b/progenitor-client/src/progenitor_client.rs
index 1ce617f6..3738fdd4 100644
--- a/progenitor-client/src/progenitor_client.rs
+++ b/progenitor-client/src/progenitor_client.rs
@@ -258,6 +258,9 @@ pub enum Error<E = ()> {
 
     /// An error occurred in the processing of a request pre-hook.
     PreHookError(String),
+
+    /// An error occurred in the processing of a request post-hook.
+    PostHookError(String),
 }
 
 impl<E> Error<E> {
@@ -266,6 +269,7 @@ impl<E> Error<E> {
         match self {
             Error::InvalidRequest(_) => None,
             Error::PreHookError(_) => None,
+            Error::PostHookError(_) => None,
             Error::CommunicationError(e) => e.status(),
             Error::ErrorResponse(rv) => Some(rv.status()),
             Error::InvalidUpgrade(e) => e.status(),
@@ -283,6 +287,7 @@ impl<E> Error<E> {
         match self {
             Error::InvalidRequest(s) => Error::InvalidRequest(s),
             Error::PreHookError(s) => Error::PreHookError(s),
+            Error::PostHookError(s) => Error::PostHookError(s),
             Error::CommunicationError(e) => Error::CommunicationError(e),
             Error::ErrorResponse(ResponseValue {
                 inner: _,
@@ -346,6 +351,9 @@ where
             Error::PreHookError(s) => {
                 write!(f, "Pre-hook Error: {}", s)
             }
+            Error::PostHookError(s) => {
+                write!(f, "Post-hook Error: {}", s)
+            }
         }
     }
 }
diff --git a/progenitor-impl/src/lib.rs b/progenitor-impl/src/lib.rs
index 564c74da..d81dee40 100644
--- a/progenitor-impl/src/lib.rs
+++ b/progenitor-impl/src/lib.rs
@@ -64,6 +64,7 @@ pub struct GenerationSettings {
     pre_hook: Option<TokenStream>,
     pre_hook_async: Option<TokenStream>,
     post_hook: Option<TokenStream>,
+    post_hook_async: Option<TokenStream>,
     extra_derives: Vec<String>,
 
     unknown_crates: UnknownPolicy,
@@ -152,6 +153,16 @@ impl GenerationSettings {
         self
     }
 
+    /// Hook invoked prior to receiving the HTTP response.
+    /// Any hook requires the inner type to be set via
+    /// [Self::with_inner_type], otherwise the code will not compile.
+    /// The signature for the post hook function should be
+    /// `post_hook_func(&InnerType, &Result<reqwest::Response,reqwest::Error>) -> ()`
+    pub fn with_post_hook_async(&mut self, post_hook_async: TokenStream) -> &mut Self {
+        self.post_hook_async = Some(post_hook_async);
+        self
+    }
+
     /// Additional derive macros applied to generated types.
     pub fn with_derive(&mut self, derive: impl ToString) -> &mut Self {
         self.extra_derives.push(derive.to_string());
diff --git a/progenitor-impl/src/method.rs b/progenitor-impl/src/method.rs
index 2c2a731b..740ab1ee 100644
--- a/progenitor-impl/src/method.rs
+++ b/progenitor-impl/src/method.rs
@@ -1147,6 +1147,14 @@ impl Generator {
                 (#hook)(&#client.inner, &#result_ident);
             }
         });
+        let post_hook_async = self.settings.post_hook_async.as_ref().map(|hook| {
+            quote! {
+                match (#hook)(&#client.inner, &#result_ident).await {
+                    Ok(_) => (),
+                    Err(e) => return Err(Error::PostHookError(e.to_string())),
+                }
+            }
+        });
 
         let method_func = format_ident!("{}", method.method.as_str());
 
@@ -1172,6 +1180,7 @@ impl Generator {
                 .execute(#request_ident)
                 .await;
             #post_hook
+            #post_hook_async
 
             let #response_ident = #result_ident?;