Add utoipa modifier to update all transient Response_<> to `<>Respo…

…nse`

crates/server/src/lib.rs

@@ -13,7 +13,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#![feature(never_type, decl_macro)]
+#![feature(never_type, decl_macro, entry_insert)]
 
 mod state;
 pub use state::*;

crates/server/src/openapi.rs

@@ -31,7 +31,8 @@ use utoipa::{
     modifiers(
         &UpdatePathsToIncludeDefaultVersion,
         &IncludeErrorProneDatatypes,
-        &SecuritySchemes
+        &SecuritySchemes,
+        &ResponseModifiers,
     ),
     info(
         title = "charted-server",
@@ -94,7 +95,9 @@ use utoipa::{
             charted_types::payloads::user::PatchUserPayload,
 
             // ==== Response Datatypes ====
-            crate::routing::v1::EntrypointResponse,
+            crate::routing::v1::info::Info,
+            crate::routing::v1::main::Main,
+            crate::routing::v1::Entrypoint,
 
             // ==== Helm ====
             charted_types::helm::StringOrImportValue,

crates/server/src/openapi/modifiers.rs

@@ -14,11 +14,11 @@
 // limitations under the License.
 
 use charted_core::api::Version;
-use std::collections::BTreeMap;
+use std::collections::{BTreeMap, HashMap, HashSet};
 use utoipa::{
     openapi::{
         security::{ApiKey, ApiKeyValue, HttpAuthScheme, HttpBuilder, SecurityScheme},
-        ComponentsBuilder, OpenApi,
+        ComponentsBuilder, OpenApi, Ref, RefOr, Schema,
     },
     Modify,
 };
@@ -139,22 +139,131 @@ impl Modify for SecuritySchemes {
     }
 }
 
+/// [Modifier][Modify] that replaces `Response_<type>` as `<type>Response`.
+pub struct ResponseModifiers;
+
+// While the implementation can be cleaned and optimized, for now it works.
+impl Modify for ResponseModifiers {
+    fn modify(&self, openapi: &mut OpenApi) {
+        let mut components = openapi.components.take().unwrap();
+
+        // 1. First, we need to update all `Response_<>` with `<>Response` and update the
+        //    `data` field to be the type that we want
+        let mut scheduled_to_be_removed = HashSet::new();
+        let mut scheduled_to_be_created = HashMap::new();
+        let schemas = components.schemas.clone();
+
+        for (key, schema) in schemas
+            .iter()
+            .filter_map(|(key, schema)| key.starts_with("Response_").then_some((key, schema)))
+        {
+            // First, we need to schedule the deletion of `key` since we will
+            // no longer be using it
+            scheduled_to_be_removed.insert(key.as_str());
+
+            // Update the schema's description
+            //
+            // We only expect `"type": "object"` as response types are only objects
+            // and cannot be anything else
+            let RefOr::T(Schema::Object(mut object)) = schema.clone() else {
+                unreachable!();
+            };
+
+            let (_, mut ty) = key.split_once('_').unwrap();
+            if ty.ends_with("Response") {
+                ty = ty.trim_end_matches("Response");
+            }
+
+            object.description = Some(format!("Response datatype for the `{ty}` type"));
+            assert!(object.properties.remove("data").is_some());
+
+            object
+                .properties
+                .insert("data".into(), RefOr::Ref(Ref::from_schema_name(ty)));
+
+            scheduled_to_be_created.insert(format!("{ty}Response"), RefOr::T(Schema::Object(object)));
+        }
+
+        scheduled_to_be_removed
+            .iter()
+            .map(|x| components.schemas.remove(*x))
+            .for_each(drop);
+
+        scheduled_to_be_created
+            .iter()
+            .map(|(key, value)| components.schemas.insert(key.to_owned(), value.clone()))
+            .for_each(drop);
+
+        // 2. Now, we need to go to every path and check if there is a ref
+        //    with the `Response_<>` suffix
+        {
+            for item in openapi.paths.paths.values_mut() {
+                macro_rules! do_update {
+                    ($kind:ident as $op:expr) => {
+                        if let Some(ref mut op) = $op {
+                            for resp in op
+                                .responses
+                                .responses
+                                .values_mut()
+                                .filter_map(|resp| match resp {
+                                    RefOr::T(resp) => Some(resp),
+                                    _ => None,
+                                })
+                            {
+                                for content in resp.content.values_mut() {
+                                    if let Some(RefOr::Ref(ref_)) = content.schema.as_ref() {
+                                        if ref_.ref_location.contains("Response_") {
+                                            let reference = ref_.ref_location.split("/").last().unwrap();
+
+                                            let (_, ty) = reference.split_once('_').unwrap();
+                                            assert!(!ty.contains("_"));
+
+                                            content.schema =
+                                                Some(RefOr::Ref(Ref::from_schema_name(format!("{ty}Response"))));
+                                        }
+                                    }
+                                }
+                            }
+
+                            item.$kind = ($op).clone();
+                        }
+                    };
+                }
+
+                do_update!(get as item.get);
+                do_update!(put as item.put);
+                do_update!(head as item.head);
+                do_update!(post as item.post);
+                do_update!(patch as item.patch);
+                do_update!(trace as item.trace);
+                do_update!(delete as item.delete);
+                do_update!(delete as item.delete);
+                do_update!(options as item.options);
+            }
+        }
+
+        openapi.components = Some(components);
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
-    use utoipa::OpenApi;
-
-    /// a dummy route that exists
-    #[utoipa::path(get, path = "/v1/weow")]
-    #[allow(dead_code)]
-    fn dummy_route() {}
-
-    #[utoipa::path(get, path = "/v1")]
-    #[allow(dead_code)]
-    fn other_dummy_route() {}
+    use charted_core::api;
+    use charted_types::User;
+    use utoipa::{openapi::HttpMethod, OpenApi};
 
     #[test]
     fn update_paths_to_include_default_api_version() {
+        // a dummy route that exists
+        #[utoipa::path(get, path = "/v1/weow")]
+        #[allow(dead_code)]
+        fn dummy_route() {}
+
+        #[utoipa::path(get, path = "/v1")]
+        #[allow(dead_code)]
+        fn other_dummy_route() {}
+
         #[derive(OpenApi)]
         #[openapi(paths(dummy_route, other_dummy_route), modifiers(&UpdatePathsToIncludeDefaultVersion))]
         struct Document;
@@ -164,4 +273,60 @@ mod tests {
 
         assert_eq!(&paths, &["/", "/v1", "/v1/weow", "/weow"]);
     }
+
+    // This test combats using `Response_<>` as `<>Response`, where `<>` is
+    // the type that is registered as a schema.
+    #[test]
+    fn update_response_types() {
+        #[utoipa::path(get, path = "/", responses((status = 200, body = api::Response<User>)))]
+        #[allow(unused)]
+        fn test_path() {}
+
+        #[derive(OpenApi)]
+        #[openapi(paths(test_path), modifiers(&ResponseModifiers))]
+        struct Document;
+
+        let openapi = Document::openapi();
+
+        // Check that `UserResponse` has the modified description and reference
+        {
+            let components = openapi.components.unwrap();
+            let RefOr::T(Schema::Object(object)) = components.schemas.get("UserResponse").unwrap() else {
+                unreachable!();
+            };
+
+            assert_eq!(
+                object.description,
+                Some("Response datatype for the `User` type".to_string())
+            );
+
+            let Some(RefOr::Ref(ref_)) = object.properties.get("data") else {
+                unreachable!();
+            };
+
+            assert_eq!(ref_.ref_location, "#/components/schemas/User");
+        }
+
+        // Check if `GET /` has the updated response type
+        {
+            let paths = openapi.paths.clone();
+            let Some(path) = paths.get_path_operation("/", HttpMethod::Get) else {
+                unreachable!();
+            };
+
+            let Some(RefOr::T(resp)) = path.responses.responses.get("200") else {
+                unreachable!();
+            };
+
+            let Some(content) = resp.content.get("application/json") else {
+                unreachable!();
+            };
+
+            let Some(RefOr::Ref(ref ref_)) = content.schema else {
+                unreachable!()
+            };
+
+            assert_eq!(ref_.ref_location, "#/components/responses/UserResponse");
+        }
+    }
 }

crates/server/src/routing/v1/info.rs

@@ -20,7 +20,7 @@ use utoipa::ToSchema;
 
 /// Represents the response for the `GET /info` REST handler.
 #[derive(Serialize, ToSchema)]
-pub struct InfoResponse {
+pub struct Info {
     /// The distribution the server is running off from
     pub distribution: Distribution,
 
@@ -40,9 +40,9 @@ pub struct InfoResponse {
     pub vendor: &'static str,
 }
 
-impl Default for InfoResponse {
-    fn default() -> InfoResponse {
-        InfoResponse {
+impl Default for Info {
+    fn default() -> Self {
+        Self {
             distribution: Distribution::detect(),
             commit_sha: COMMIT_HASH,
             build_date: BUILD_DATE,
@@ -63,12 +63,12 @@ impl Default for InfoResponse {
         (
             status = 200,
             description = "Successful response",
-            body = inline(api::Response<InfoResponse>),
+            body = api::Response<Info>,
             content_type = "application/json"
         )
     )
 )]
 #[cfg_attr(debug_assertions, axum::debug_handler)]
-pub async fn info() -> api::Response<InfoResponse> {
+pub async fn info() -> api::Response<Info> {
     api::from_default(StatusCode::OK)
 }

crates/server/src/routing/v1/main.rs

@@ -23,7 +23,7 @@ use utoipa::{
 
 /// Response object for the `GET /` REST controller.
 #[derive(Serialize, ToSchema)]
-pub struct MainResponse {
+pub struct Main {
     /// The message, which will always be "Hello, world!"
     pub message: &'static str,
 
@@ -34,20 +34,20 @@ pub struct MainResponse {
     pub docs: String,
 }
 
-impl Default for MainResponse {
+impl Default for Main {
     fn default() -> Self {
-        MainResponse {
+        Self {
             message: "Hello, world! 👋",
             tagline: "You know, for Helm charts?",
             docs: format!("https://charts.noelware.org/docs/server/{VERSION}"),
         }
     }
 }
 
-impl<'r> ToResponse<'r> for MainResponse {
+impl<'r> ToResponse<'r> for Main {
     fn response() -> (&'r str, RefOr<Response>) {
         (
-            "MainResponse",
+            "Main",
             RefOr::T(
                 ResponseBuilder::new()
                     .description("Response for the `/` REST handler")
@@ -73,12 +73,12 @@ impl<'r> ToResponse<'r> for MainResponse {
         (
             status = 200,
             description = "Successful response",
-            body = inline(api::Response<MainResponse>),
+            body = api::Response<Main>,
             content_type = "application/json"
         )
     )
 )]
 #[cfg_attr(debug_assertions, axum::debug_handler)]
-pub async fn main() -> api::Response<MainResponse> {
+pub async fn main() -> api::Response<Main> {
     api::from_default(StatusCode::OK)
 }

crates/server/src/routing/v1/mod.rs

@@ -30,15 +30,15 @@ use utoipa::ToSchema;
 
 /// Generic entrypoint message for any API route like `/users`.
 #[derive(Serialize, ToSchema)]
-pub struct EntrypointResponse {
+pub struct Entrypoint {
     /// Humane message to greet you.
     pub message: Cow<'static, str>,
 
     /// URI to the documentation for this entrypoint.
     pub docs: Cow<'static, str>,
 }
 
-impl EntrypointResponse {
+impl Entrypoint {
     pub fn new(entity: impl AsRef<str>) -> Self {
         let entity = entity.as_ref();
         Self {