Repository: juhaku/utoipa
Pull request #34 (2022-03-01 20:57:08)

Add Modify trait
* Add Modify trait to allow runtime OpenAPI modification
* Add unit tests
* Add documentation -> also improve documentation generally
* Rename Schema to Components. Name was previously bit misleading.


Patch:
diff --git a/src/lib.rs b/src/lib.rs
index feada9b..eff8d28 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -146,10 +146,15 @@
 //!
 //! println!("{}", ApiDoc::openapi().to_pretty_json().unwrap());
 //! ```
+//! # Go beyond the surface
 //!
-//! See how to serve OpenAPI doc via Swagger UI check [`swagger_ui`] module for more details.
+//! * See how to serve OpenAPI doc via Swagger UI check [`swagger_ui`] module for more details.
 //! You can also browse to [examples](https://github.com/juhaku/utoipa/tree/master/examples)
 //! for more comprehensinve examples.
+//! * Modify generated OpenAPI at runtime check [`Modify`] trait for more details.
+//! * More about OpenAPI security in [security documentation][security].
+//!
+//! [security]: openapi/security/index.html
 
 /// Rust implementation of Openapi Spec V3
 pub mod openapi;
@@ -158,7 +163,7 @@ pub mod swagger_ui;
 
 pub use utoipa_gen::*;
 
-/// OpenApi trait for implementing OpenAPI specification in Rust.
+/// Trait for implementing OpenAPI specification in Rust.
 ///
 /// This trait is derivable and can be used with `#[derive]` attribute. The derived implementation
 /// will use Cargo provided environment variables to implement the default information. For a details of
@@ -195,7 +200,7 @@ pub use utoipa_gen::*;
 ///                 ),
 ///             utoipa::openapi::path::Paths::new(),
 ///         )
-///         .with_components(utoipa::openapi::Schema::new())
+///         .with_components(utoipa::openapi::Components::new())
 ///     }
 /// }
 /// ```
@@ -204,7 +209,7 @@ pub trait OpenApi {
     fn openapi() -> openapi::OpenApi;
 }
 
-/// Component trait for implementing OpenAPI Schema object.
+/// Trait for implementing OpenAPI Schema object.
 ///
 /// This trait is deriveable and can be used with `[#derive]` attribute. For a details of
 /// `#[derive(Component)]` refer to [derive documentation][derive].
@@ -264,7 +269,7 @@ pub trait Component {
     fn component() -> openapi::schema::Component;
 }
 
-/// Path trait for implementing OpenAPI PathItem object with path.
+/// Trait for implementing OpenAPI PathItem object with path.
 ///
 /// This trait is implemented via [`#[utoipa::path(...)]`][derive] attribute macro and there
 /// is no need to implement this trait manually.
@@ -347,3 +352,45 @@ pub trait Path {
 
     fn path_item(defalt_tag: Option<&str>) -> openapi::path::PathItem;
 }
+
+/// Trait that allows OpenApi modification at runtime.
+///
+/// Implement this trait if you wish to modify the OpenApi at runtime before it is being consumed
+/// *(Before `utoipa::OpenApi::openapi()` function returns)*.
+/// This is trait can be used to add or change already generated OpenApi spec to alter the generated
+/// specification by user defined condition. For example you can add definitions that should be loaded
+/// from some configuration at runtime what may not be available during compile time.
+///
+/// See more about [`OpenApi`][derive] derive at [derive documentation][derive].
+///
+/// [derive]: derive.OpenApi.html
+/// [security_schema]: openapi/security/enum.SecuritySchema.html
+///
+/// # Examples
+///
+/// Add custom JWT [`SecuritySchema`][security_schema] to [`OpenApi`][`openapi::OpenApi`].
+/// ```rust
+/// # use utoipa::{OpenApi, Modify};
+/// # use utoipa::openapi::security::{SecuritySchema, Http, HttpAuthenticationType};
+/// #[derive(OpenApi)]
+/// #[openapi(modifiers = [&SecurityAddon])]
+/// struct ApiDoc;
+///
+/// struct SecurityAddon;
+///
+/// impl Modify for SecurityAddon {
+///    fn modify(&self, openapi: &mut utoipa::openapi::OpenApi) {
+///        if let Some(components) = openapi.components.as_mut() {
+///            components.add_security_schema(
+///                "api_jwt_token",
+///                SecuritySchema::Http(
+///                    Http::new(HttpAuthenticationType::Bearer).with_bearer_format("JWT"),
+///                ),
+///            )
+///        }
+///    }
+///}
+/// ```
+pub trait Modify {
+    fn modify(&self, openapi: &mut openapi::OpenApi);
+}
diff --git a/src/openapi/mod.rs b/src/openapi/mod.rs
index 5855b5b..7923d38 100644
--- a/src/openapi/mod.rs
+++ b/src/openapi/mod.rs
@@ -12,7 +12,7 @@ pub use self::{
     path::{PathItem, PathItemType, Paths},
     response::{Response, Responses},
     schema::{
-        Array, Component, ComponentFormat, ComponentType, Object, OneOf, Property, Ref, Schema,
+        Array, Component, ComponentFormat, ComponentType, Components, Object, OneOf, Property, Ref,
         ToArray,
     },
     security::SecurityRequirement,
@@ -50,7 +50,7 @@ pub struct OpenApi {
     pub paths: BTreeMap<String, PathItem>,
 
     #[serde(skip_serializing_if = "Option::is_none")]
-    pub components: Option<Schema>,
+    pub components: Option<Components>,
 
     #[serde(skip_serializing_if = "Option::is_none")]
     pub security: Option<Vec<SecurityRequirement>>,
@@ -77,8 +77,8 @@ impl OpenApi {
         self
     }
 
-    pub fn with_components(mut self, schema: Schema) -> Self {
-        self.components = Some(schema);
+    pub fn with_components(mut self, components: Components) -> Self {
+        self.components = Some(components);
 
         self
     }
diff --git a/src/openapi/schema.rs b/src/openapi/schema.rs
index 4232d58..649fedb 100644
--- a/src/openapi/schema.rs
+++ b/src/openapi/schema.rs
@@ -9,13 +9,13 @@ use super::{security::SecuritySchema, Deprecated};
 #[non_exhaustive]
 #[derive(Serialize, Deserialize, Default, Clone)]
 #[serde(rename_all = "camelCase")]
-pub struct Schema {
+pub struct Components {
     schemas: HashMap<String, Component>,
 
     security_schemas: HashMap<String, SecuritySchema>,
 }
 
-impl Schema {
+impl Components {
     pub fn new() -> Self {
         Self {
             ..Default::default()
@@ -43,6 +43,30 @@ impl Schema {
 
         self
     }
+
+    pub fn add_security_schema<N: Into<String>, S: Into<SecuritySchema>>(
+        &mut self,
+        name: N,
+        security_schema: S,
+    ) {
+        self.security_schemas
+            .insert(name.into(), security_schema.into());
+    }
+
+    pub fn add_security_schemas_from_iter<
+        I: IntoIterator<Item = (N, S)>,
+        N: Into<String>,
+        S: Into<SecuritySchema>,
+    >(
+        &mut self,
+        schemas: I,
+    ) {
+        self.security_schemas.extend(
+            schemas
+                .into_iter()
+                .map(|(name, item)| (name.into(), item.into())),
+        );
+    }
 }
 
 #[non_exhaustive]
@@ -492,7 +516,7 @@ mod tests {
     #[cfg(feature = "serde_json")]
     fn create_schema_serializes_json() -> Result<(), serde_json::Error> {
         let openapi = OpenApi::new(Info::new("My api", "1.0.0"), Paths::new()).with_components(
-            Schema::new()
+            Components::new()
                 .with_component("Person", Ref::new("#/components/PersonModel"))
                 .with_component(
                     "Credential",
diff --git a/utoipa-gen/src/lib.rs b/utoipa-gen/src/lib.rs
index 95b8e47..b30e083 100644
--- a/utoipa-gen/src/lib.rs
+++ b/utoipa-gen/src/lib.rs
@@ -357,14 +357,16 @@ pub fn path(attr: TokenStream, item: TokenStream) -> TokenStream {
 #[proc_macro_derive(OpenApi, attributes(openapi))]
 /// OpenApi derive macro
 ///
-/// This is `#[derive]` implementation for [OpenApi][openapi] trait. The macro accepts one `openapi` argument.
+/// This is `#[derive]` implementation for [`OpenApi`][openapi] trait. The macro accepts one `openapi` argument.
 ///
-/// **Accepted `openapi` argument attributes**
+/// **Accepted argument attributes:**
 ///
 /// * **handlers**  List of method references having attribute [`#[utoipa::path]`][path] macro.
-/// * **components**  List of [Component][component]s in OpenAPI schema.
+/// * **components**  List of [`Component`][component]s in OpenAPI schema.
+/// * **modifiers** List of items implemeting [`Modify`][modify] trait for runtime OpenApi modification.
+///   See the [trait documentation][modify] for more details.
 ///
-/// OpenApi derive macro will also derive [Info][info] for OpenApi specification using Cargo
+/// OpenApi derive macro will also derive [`Info`][info] for OpenApi specification using Cargo
 /// environment variables.
 ///
 /// * env `CARGO_PKG_NAME` map to info `title`
@@ -411,6 +413,7 @@ pub fn path(attr: TokenStream, item: TokenStream) -> TokenStream {
 /// [openapi]: trait.OpenApi.html
 /// [component]: derive.Component.html
 /// [path]: attr.path.html
+/// [modify]: trait.Modify.html
 /// [info]: openapi/info/struct.Info.html
 pub fn openapi(input: TokenStream) -> TokenStream {
     let DeriveInput { attrs, ident, .. } = syn::parse_macro_input!(input);
diff --git a/utoipa-gen/src/openapi.rs b/utoipa-gen/src/openapi.rs
index 6f221a2..67fdf5c 100644
--- a/utoipa-gen/src/openapi.rs
+++ b/utoipa-gen/src/openapi.rs
@@ -4,6 +4,7 @@ use syn::{
     bracketed,
     parse::{Parse, ParseStream},
     punctuated::Punctuated,
+    token::{And, Comma},
     Attribute, Error, ExprPath, GenericParam, Generics, Token,
 };
 
@@ -21,6 +22,7 @@ const PATH_STRUCT_PREFIX: &str = "__path_";
 pub struct OpenApiAttr {
     handlers: Vec<ExprPath>,
     components: Vec<Component>,
+    modifiers: Punctuated<Modifier, Comma>,
 }
 
 #[cfg_attr(feature = "debug", derive(Debug))]
@@ -29,6 +31,31 @@ struct Component {
     generics: Generics,
 }
 
+#[cfg_attr(feature = "debug", derive(Debug))]
+struct Modifier {
+    and: And,
+    ident: Ident,
+}
+
+impl ToTokens for Modifier {
+    fn to_tokens(&self, tokens: &mut TokenStream) {
+        let and = &self.and;
+        let ident = &self.ident;
+        tokens.extend(quote! {
+            #and #ident
+        })
+    }
+}
+
+impl Parse for Modifier {
+    fn parse(input: ParseStream) -> syn::Result<Self> {
+        Ok(Self {
+            and: input.parse()?,
+            ident: input.parse()?,
+        })
+    }
+}
+
 impl Component {
     fn has_lifetime_generics(&self) -> bool {
         self.generics
@@ -57,16 +84,19 @@ impl Parse for OpenApiAttr {
 
             match attribute {
                 "handlers" => {
-                    openapi.handlers = parse_handlers(input).unwrap_or_abort();
+                    openapi.handlers = parse_handlers(input)?;
                 }
                 "components" => {
-                    openapi.components = parse_components(input).unwrap_or_abort();
+                    openapi.components = parse_components(input)?;
+                }
+                "modifiers" => {
+                    openapi.modifiers = parse_modifiers(input)?;
                 }
                 _ => {
                     return Err(Error::new(
                         ident.span(),
                         format!(
-                            "unexpected attribute: {}, expected: handlers, components",
+                            "unexpected attribute: {}, expected: handlers, components, modifiers",
                             ident
                         ),
                     ));
@@ -133,6 +163,15 @@ fn parse_components(input: ParseStream) -> syn::Result<Vec<Component>> {
     })
 }
 
+fn parse_modifiers(input: ParseStream) -> syn::Result<Punctuated<Modifier, Comma>> {
+    parse_utils::parse_next(input, || {
+        let content;
+        bracketed!(content in input);
+
+        Punctuated::<Modifier, Comma>::parse_terminated(&content)
+    })
+}
+
 pub(crate) struct OpenApi(pub OpenApiAttr, pub Ident);
 
 impl ToTokens for OpenApi {
@@ -141,8 +180,8 @@ impl ToTokens for OpenApi {
 
         let info = info::impl_info();
 
-        let schema = attributes.components.iter().fold(
-            quote! { utoipa::openapi::Schema::new() },
+        let components = attributes.components.iter().fold(
+            quote! { utoipa::openapi::Components::new() },
             |mut schema, component| {
                 let ident = &component.ty;
                 let span = ident.span();
@@ -172,14 +211,30 @@ impl ToTokens for OpenApi {
             },
         );
 
+        let modifiers = &self.0.modifiers;
+        let modifiers_len = modifiers.len();
+
+        modifiers.iter().for_each(|modifier| {
+            let assert_modifier = format_ident!("_Assert{}", modifier.ident);
+            let ident = &modifier.ident;
+            quote_spanned! {modifier.ident.span()=>
+                struct #assert_modifier where #ident : utoipa::Modify;
+            };
+        });
+
         let path_items = impl_paths(&attributes.handlers);
 
         tokens.extend(quote! {
             impl utoipa::OpenApi for #ident {
                 fn openapi() -> utoipa::openapi::OpenApi {
                     use utoipa::{Component, Path};
-                    utoipa::openapi::OpenApi::new(#info, #path_items)
-                        .with_components(#schema)
+                    let mut openapi = utoipa::openapi::OpenApi::new(#info, #path_items)
+                        .with_components(#components);
+
+                    let _mods: [&dyn utoipa::Modify; #modifiers_len] = [#modifiers];
+                    _mods.iter().for_each(|modifier| modifier.modify(&mut openapi));
+
+                    openapi
                 }
             }
         });


Test patch:
diff --git a/tests/modify_test.rs b/tests/modify_test.rs
new file mode 100644
index 0000000..78825bf
--- /dev/null
+++ b/tests/modify_test.rs
@@ -0,0 +1,41 @@
+#![cfg(feature = "json")]
+
+use utoipa::{
+    openapi::{
+        self,
+        security::{Http, HttpAuthenticationType, SecuritySchema},
+    },
+    Modify, OpenApi,
+};
+
+mod common;
+
+#[test]
+fn modify_openapi_add_security_schema() {
+    #[derive(Default, OpenApi)]
+    #[openapi(modifiers = [&SecurityAddon])]
+    struct ApiDoc;
+
+    struct SecurityAddon;
+
+    impl Modify for SecurityAddon {
+        fn modify(&self, openapi: &mut openapi::OpenApi) {
+            if let Some(schema) = openapi.components.as_mut() {
+                schema.add_security_schema(
+                    "api_jwt_token",
+                    SecuritySchema::Http(
+                        Http::new(HttpAuthenticationType::Bearer).with_bearer_format("JWT"),
+                    ),
+                )
+            }
+        }
+    }
+
+    let doc = serde_json::to_value(ApiDoc::openapi()).unwrap();
+
+    assert_value! {doc=>
+        "components.securitySchemas.api_jwt_token.scheme" = r###""bearer""###, "api_jwt_token scheme"
+        "components.securitySchemas.api_jwt_token.type" = r###""http""###, "api_jwt_token type"
+        "components.securitySchemas.api_jwt_token.bearerFormat" = r###""JWT""###, "api_jwt_token bearerFormat"
+    }
+}
diff --git a/tests/utoipa_gen_test.rs b/tests/utoipa_gen_test.rs
index d9b3420..36dbdec 100644
--- a/tests/utoipa_gen_test.rs
+++ b/tests/utoipa_gen_test.rs
@@ -2,7 +2,13 @@
 #![allow(dead_code)]
 
 use serde::{Deserialize, Serialize};
-use utoipa::{Component, OpenApi};
+use utoipa::{
+    openapi::{
+        self,
+        security::{Http, HttpAuthenticationType, SecuritySchema},
+    },
+    Component, Modify, OpenApi,
+};
 
 #[derive(Deserialize, Serialize, Component)]
 #[component(example = json!({"name": "bob the cat", "id": 1}))]
@@ -64,8 +70,8 @@ mod pet_api {
     }
 }
 
-#[derive(OpenApi, Default)]
-#[openapi(handlers = [pet_api::get_pet_by_id], components = [Pet])]
+#[derive(Default, OpenApi)]
+#[openapi(handlers = [pet_api::get_pet_by_id], components = [Pet], modifiers = [&Foo])]
 struct ApiDoc;
 
 #[test]
@@ -77,3 +83,18 @@ fn derive_openapi() {
     );
     println!("{}", ApiDoc::openapi().to_pretty_json().unwrap());
 }
+
+impl Modify for Foo {
+    fn modify(&self, openapi: &mut openapi::OpenApi) {
+        if let Some(schema) = openapi.components.as_mut() {
+            schema.add_security_schema(
+                "token_jwt",
+                SecuritySchema::Http(
+                    Http::new(HttpAuthenticationType::Bearer).with_bearer_format("JWT"),
+                ),
+            )
+        }
+    }
+}
+
+struct Foo;