use std::cell::RefCell;
use cfg_if::cfg_if;
use schemars::{generate::SchemaSettings, Schema, SchemaGenerator};
use serde_json::Value;
use crate::error::Error;
thread_local! {
static GEN_CTX: RefCell<GenContext> = RefCell::new(GenContext::new());
}
pub fn in_context<R, F>(cb: F) -> R
where
F: FnOnce(&mut GenContext) -> R,
{
GEN_CTX.with(|ctx| cb(&mut ctx.borrow_mut()))
}
pub fn on_error(handler: impl Fn(Error) + 'static) {
in_context(|ctx| ctx.error_handler = Some(Box::new(handler)));
}
pub fn extract_schemas(extract: bool) {
in_context(|ctx| {
ctx.set_extract_schemas(extract);
});
}
pub fn inferred_empty_response_status(status: u16) {
in_context(|ctx| {
ctx.no_content_status = status;
});
}
pub fn infer_responses(infer: bool) {
in_context(|ctx| {
ctx.infer_responses = infer;
});
}
pub fn all_error_responses(infer: bool) {
in_context(|ctx| {
ctx.all_error_responses = infer;
});
}
pub fn reset_context() {
in_context(|ctx| *ctx = GenContext::new());
}
pub struct GenContext {
pub schema: SchemaGenerator,
pub(crate) infer_responses: bool,
pub(crate) all_error_responses: bool,
pub(crate) extract_schemas: bool,
pub(crate) no_content_status: u16,
pub(crate) show_error: fn(&Error) -> bool,
error_handler: Option<Box<dyn Fn(Error)>>,
}
impl GenContext {
fn new() -> Self {
cfg_if! {
if #[cfg(feature = "axum")] {
let no_content_status = 200;
} else {
let no_content_status = 204;
}
}
let mut this = Self {
schema: SchemaGenerator::new(SchemaSettings::draft07()),
infer_responses: true,
all_error_responses: false,
extract_schemas: true,
show_error: default_error_filter,
error_handler: None,
no_content_status,
};
this.set_extract_schemas(true);
this
}
pub(crate) fn reset_error_filter(&mut self) {
self.show_error = default_error_filter;
}
fn set_extract_schemas(&mut self, extract: bool) {
if extract {
self.schema = SchemaGenerator::new(SchemaSettings::draft07().with(|s| {
s.inline_subschemas = false;
s.definitions_path = "#/components/schemas/".into();
}));
self.extract_schemas = true;
} else {
self.schema = SchemaGenerator::new(SchemaSettings::draft07().with(|s| {
s.inline_subschemas = true;
}));
self.extract_schemas = false;
}
}
#[tracing::instrument(skip_all)]
pub fn error(&mut self, error: Error) {
if let Some(handler) = &self.error_handler {
if !(self.show_error)(&error) {
return;
}
handler(error);
}
}
#[must_use]
pub fn resolve_schema<'s>(&'s self, schema_or_ref: &'s Schema) -> &'s Schema {
match &schema_or_ref.as_object().and_then(|o| o.get("$ref")) {
Some(Value::String(r)) => self
.schema
.definitions()
.get(r.strip_prefix("#/components/schemas/").unwrap_or(r))
.and_then(|s| Into::<serde_json::Result<&Schema>>::into(s.try_into()).ok())
.unwrap_or(schema_or_ref),
_ => schema_or_ref,
}
}
}
fn default_error_filter(_: &Error) -> bool {
true
}