use std::collections::{BTreeMap, BTreeSet};
use std::path::PathBuf;
use utoipa::openapi::path::{Operation, PathItem};
use utoipa::openapi::{Components, OpenApi, Ref, RefOr};
use super::{Fragment, OpenApiSplitter, SplitResult};
fn iter_operations(path_item: &PathItem) -> impl Iterator<Item = &Operation> {
[
path_item.get.as_ref(),
path_item.put.as_ref(),
path_item.post.as_ref(),
path_item.delete.as_ref(),
path_item.options.as_ref(),
path_item.head.as_ref(),
path_item.patch.as_ref(),
path_item.trace.as_ref(),
]
.into_iter()
.flatten()
}
#[derive(Debug, Clone)]
pub struct SplitSchemasByTag {
common_file: PathBuf,
schemas_dir: Option<PathBuf>,
}
impl SplitSchemasByTag {
pub fn new(common_file: impl Into<PathBuf>) -> Self {
Self {
common_file: common_file.into(),
schemas_dir: None,
}
}
pub fn with_schemas_dir(mut self, dir: impl Into<PathBuf>) -> Self {
self.schemas_dir = Some(dir.into());
self
}
fn analyze_schema_usage(&self, spec: &OpenApi) -> BTreeMap<String, BTreeSet<String>> {
let mut schema_to_tags: BTreeMap<String, BTreeSet<String>> = BTreeMap::new();
for path_item in spec.paths.paths.values() {
for operation in iter_operations(path_item) {
let tags = operation.tags.clone().unwrap_or_default();
if tags.is_empty() {
continue;
}
if let Some(ref request_body) = operation.request_body {
for content in request_body.content.values() {
if let Some(ref schema) = content.schema {
self.collect_schema_refs(schema, &tags, &mut schema_to_tags);
}
}
}
for response in operation.responses.responses.values() {
if let RefOr::T(resp) = response {
for content in resp.content.values() {
if let Some(ref schema) = content.schema {
self.collect_schema_refs(schema, &tags, &mut schema_to_tags);
}
}
}
}
if let Some(ref parameters) = operation.parameters {
for param in parameters {
if let Some(ref schema) = param.schema {
self.collect_schema_refs(schema, &tags, &mut schema_to_tags);
}
}
}
}
}
schema_to_tags
}
fn collect_schema_refs(
&self,
schema: &RefOr<utoipa::openapi::Schema>,
tags: &[String],
schema_to_tags: &mut BTreeMap<String, BTreeSet<String>>,
) {
match schema {
RefOr::Ref(r) => {
if let Some(name) = extract_schema_name(&r.ref_location) {
let entry = schema_to_tags.entry(name).or_default();
for tag in tags {
entry.insert(tag.clone());
}
}
}
RefOr::T(_) => {
}
}
}
fn target_file_for_schema(&self, _schema_name: &str, tags: &BTreeSet<String>) -> PathBuf {
let base_dir = self.schemas_dir.clone().unwrap_or_default();
if tags.len() == 1 {
let tag = tags.iter().next().expect("checked len == 1");
base_dir.join(format!("{tag}.yaml"))
} else {
if self.schemas_dir.is_some() {
base_dir.join(&self.common_file)
} else {
self.common_file.clone()
}
}
}
fn create_external_ref(file_path: &std::path::Path, schema_name: &str) -> String {
format!(
"{}#/components/schemas/{}",
file_path.display(),
schema_name
)
}
}
impl OpenApiSplitter for SplitSchemasByTag {
type Fragment = Components;
fn split(&self, mut spec: OpenApi) -> SplitResult<Self::Fragment> {
let schema_to_tags = self.analyze_schema_usage(&spec);
let mut file_to_schemas: BTreeMap<PathBuf, BTreeSet<String>> = BTreeMap::new();
for (schema_name, tags) in &schema_to_tags {
let target = self.target_file_for_schema(schema_name, tags);
file_to_schemas
.entry(target)
.or_default()
.insert(schema_name.clone());
}
if file_to_schemas.len() <= 1 {
return SplitResult::new(spec);
}
let mut result = SplitResult::new(spec.clone());
let original_components = spec.components.take().unwrap_or_default();
let mut remaining_schemas = original_components.schemas.clone();
for (file_path, schema_names) in &file_to_schemas {
let mut fragment_components = Components::new();
for schema_name in schema_names {
if let Some(schema) = remaining_schemas.remove(schema_name) {
fragment_components
.schemas
.insert(schema_name.clone(), schema);
}
}
if !fragment_components.schemas.is_empty() {
result.add_fragment(Fragment::new(file_path.clone(), fragment_components));
}
}
let mut new_components = Components::new();
for (file_path, schema_names) in &file_to_schemas {
for schema_name in schema_names {
let external_ref = Self::create_external_ref(file_path, schema_name);
new_components
.schemas
.insert(schema_name.clone(), RefOr::Ref(Ref::new(external_ref)));
}
}
for (name, schema) in remaining_schemas {
new_components.schemas.insert(name, schema);
}
new_components.security_schemes = original_components.security_schemes;
new_components.responses = original_components.responses;
result.main.components = Some(new_components);
result
}
}
#[derive(Clone)]
pub struct ExtractSchemasByPredicate<F>
where
F: Fn(&str) -> bool,
{
target_file: PathBuf,
predicate: F,
}
impl<F> ExtractSchemasByPredicate<F>
where
F: Fn(&str) -> bool,
{
pub fn new(target_file: impl Into<PathBuf>, predicate: F) -> Self {
Self {
target_file: target_file.into(),
predicate,
}
}
}
impl<F> OpenApiSplitter for ExtractSchemasByPredicate<F>
where
F: Fn(&str) -> bool,
{
type Fragment = Components;
fn split(&self, mut spec: OpenApi) -> SplitResult<Self::Fragment> {
let Some(mut components) = spec.components.take() else {
return SplitResult::new(spec);
};
let schemas_to_extract: Vec<String> = components
.schemas
.keys()
.filter(|name| (self.predicate)(name))
.cloned()
.collect();
if schemas_to_extract.is_empty() {
spec.components = Some(components);
return SplitResult::new(spec);
}
let mut extracted = Components::new();
for name in &schemas_to_extract {
if let Some(schema) = components.schemas.remove(name) {
extracted.schemas.insert(name.clone(), schema);
}
}
for name in &schemas_to_extract {
let external_ref = format!(
"{}#/components/schemas/{}",
self.target_file.display(),
name
);
components
.schemas
.insert(name.clone(), RefOr::Ref(Ref::new(external_ref)));
}
spec.components = Some(components);
let mut result = SplitResult::new(spec);
result.add_fragment(Fragment::new(self.target_file.clone(), extracted));
result
}
}
fn extract_schema_name(ref_location: &str) -> Option<String> {
const SCHEMA_PREFIX: &str = "#/components/schemas/";
ref_location
.strip_prefix(SCHEMA_PREFIX)
.map(|s| s.to_string())
}
#[cfg(test)]
mod tests {
use super::*;
use utoipa::openapi::path::OperationBuilder;
use utoipa::openapi::path::PathItemBuilder;
use utoipa::openapi::{ContentBuilder, ObjectBuilder, OpenApiBuilder, ResponseBuilder};
fn create_test_spec() -> OpenApi {
let user_schema = ObjectBuilder::new()
.property(
"id",
ObjectBuilder::new().schema_type(utoipa::openapi::Type::Integer),
)
.property(
"name",
ObjectBuilder::new().schema_type(utoipa::openapi::Type::String),
)
.build();
let error_schema = ObjectBuilder::new()
.property(
"code",
ObjectBuilder::new().schema_type(utoipa::openapi::Type::Integer),
)
.property(
"message",
ObjectBuilder::new().schema_type(utoipa::openapi::Type::String),
)
.build();
let order_schema = ObjectBuilder::new()
.property(
"id",
ObjectBuilder::new().schema_type(utoipa::openapi::Type::Integer),
)
.property(
"total",
ObjectBuilder::new().schema_type(utoipa::openapi::Type::Number),
)
.build();
let mut components = Components::new();
components
.schemas
.insert("User".to_string(), RefOr::T(user_schema.into()));
components
.schemas
.insert("Error".to_string(), RefOr::T(error_schema.into()));
components
.schemas
.insert("Order".to_string(), RefOr::T(order_schema.into()));
let get_users = OperationBuilder::new()
.tags(Some(vec!["users".to_string()]))
.response(
"200",
ResponseBuilder::new()
.content(
"application/json",
ContentBuilder::new()
.schema(Some(RefOr::Ref(Ref::new("#/components/schemas/User"))))
.build(),
)
.build(),
)
.build();
let get_orders = OperationBuilder::new()
.tags(Some(vec!["orders".to_string()]))
.response(
"200",
ResponseBuilder::new()
.content(
"application/json",
ContentBuilder::new()
.schema(Some(RefOr::Ref(Ref::new("#/components/schemas/Order"))))
.build(),
)
.build(),
)
.response(
"400",
ResponseBuilder::new()
.content(
"application/json",
ContentBuilder::new()
.schema(Some(RefOr::Ref(Ref::new("#/components/schemas/Error"))))
.build(),
)
.build(),
)
.build();
let get_user_orders = OperationBuilder::new()
.tags(Some(vec!["users".to_string(), "orders".to_string()]))
.response(
"400",
ResponseBuilder::new()
.content(
"application/json",
ContentBuilder::new()
.schema(Some(RefOr::Ref(Ref::new("#/components/schemas/Error"))))
.build(),
)
.build(),
)
.build();
let mut paths = utoipa::openapi::Paths::new();
paths.paths.insert(
"/users".to_string(),
PathItemBuilder::new()
.operation(utoipa::openapi::HttpMethod::Get, get_users)
.build(),
);
paths.paths.insert(
"/orders".to_string(),
PathItemBuilder::new()
.operation(utoipa::openapi::HttpMethod::Get, get_orders)
.build(),
);
paths.paths.insert(
"/users/{id}/orders".to_string(),
PathItemBuilder::new()
.operation(utoipa::openapi::HttpMethod::Get, get_user_orders)
.build(),
);
OpenApiBuilder::new()
.paths(paths)
.components(Some(components))
.build()
}
#[test]
fn should_extract_schema_name() {
assert_eq!(
extract_schema_name("#/components/schemas/User"),
Some("User".to_string())
);
assert_eq!(
extract_schema_name("#/components/schemas/MyError"),
Some("MyError".to_string())
);
assert_eq!(extract_schema_name("#/components/responses/Error"), None);
assert_eq!(extract_schema_name("User"), None);
}
#[test]
fn should_split_by_predicate() {
let spec = create_test_spec();
let splitter = ExtractSchemasByPredicate::new("errors.yaml", |name| name.contains("Error"));
let result = splitter.split(spec);
assert_eq!(result.fragment_count(), 1);
let fragment = &result.fragments[0];
assert_eq!(fragment.path, PathBuf::from("errors.yaml"));
assert!(fragment.content.schemas.contains_key("Error"));
assert!(!fragment.content.schemas.contains_key("User"));
assert!(!fragment.content.schemas.contains_key("Order"));
let main_components = result
.main
.components
.as_ref()
.expect("should have components");
match main_components.schemas.get("Error") {
Some(RefOr::Ref(r)) => {
assert!(r.ref_location.contains("errors.yaml"));
}
_ => panic!("Expected external reference for Error"),
}
}
#[test]
fn should_not_split_when_predicate_matches_nothing() {
let spec = create_test_spec();
let splitter =
ExtractSchemasByPredicate::new("nothing.yaml", |name| name.contains("NonExistent"));
let result = splitter.split(spec);
assert!(result.is_unsplit());
}
#[test]
fn should_analyze_schema_usage() {
let spec = create_test_spec();
let splitter = SplitSchemasByTag::new("common.yaml");
let usage = splitter.analyze_schema_usage(&spec);
assert!(
usage
.get("User")
.map(|t| t.contains("users"))
.unwrap_or(false)
);
assert!(
usage
.get("Order")
.map(|t| t.contains("orders"))
.unwrap_or(false)
);
let error_tags = usage.get("Error").expect("Error should be tracked");
assert!(error_tags.contains("orders"));
}
#[test]
fn should_not_split_when_all_schemas_map_to_one_file() {
let user_schema = ObjectBuilder::new()
.property(
"id",
ObjectBuilder::new().schema_type(utoipa::openapi::Type::Integer),
)
.build();
let profile_schema = ObjectBuilder::new()
.property(
"bio",
ObjectBuilder::new().schema_type(utoipa::openapi::Type::String),
)
.build();
let mut components = Components::new();
components
.schemas
.insert("User".to_string(), RefOr::T(user_schema.into()));
components
.schemas
.insert("Profile".to_string(), RefOr::T(profile_schema.into()));
let get_users = OperationBuilder::new()
.tags(Some(vec!["users".to_string()]))
.response(
"200",
ResponseBuilder::new()
.content(
"application/json",
ContentBuilder::new()
.schema(Some(RefOr::Ref(Ref::new("#/components/schemas/User"))))
.build(),
)
.build(),
)
.build();
let get_profile = OperationBuilder::new()
.tags(Some(vec!["users".to_string()]))
.response(
"200",
ResponseBuilder::new()
.content(
"application/json",
ContentBuilder::new()
.schema(Some(RefOr::Ref(Ref::new("#/components/schemas/Profile"))))
.build(),
)
.build(),
)
.build();
let mut paths = utoipa::openapi::Paths::new();
paths.paths.insert(
"/users".to_string(),
PathItemBuilder::new()
.operation(utoipa::openapi::HttpMethod::Get, get_users)
.build(),
);
paths.paths.insert(
"/profile".to_string(),
PathItemBuilder::new()
.operation(utoipa::openapi::HttpMethod::Get, get_profile)
.build(),
);
let spec = OpenApiBuilder::new()
.paths(paths)
.components(Some(components))
.build();
let splitter = SplitSchemasByTag::new("common.yaml");
let result = splitter.split(spec);
assert!(result.is_unsplit());
}
#[test]
fn should_collect_schemas_from_parameters() {
use utoipa::openapi::path::ParameterBuilder;
use utoipa::openapi::path::ParameterIn;
let id_schema = ObjectBuilder::new()
.schema_type(utoipa::openapi::Type::String)
.build();
let mut components = Components::new();
components
.schemas
.insert("UserId".to_string(), RefOr::T(id_schema.into()));
let get_user = OperationBuilder::new()
.tags(Some(vec!["users".to_string()]))
.parameter(
ParameterBuilder::new()
.name("id")
.parameter_in(ParameterIn::Path)
.schema(Some(RefOr::Ref(Ref::new("#/components/schemas/UserId"))))
.build(),
)
.response("200", ResponseBuilder::new().description("OK").build())
.build();
let mut paths = utoipa::openapi::Paths::new();
paths.paths.insert(
"/users/{id}".to_string(),
PathItemBuilder::new()
.operation(utoipa::openapi::HttpMethod::Get, get_user)
.build(),
);
let spec = OpenApiBuilder::new()
.paths(paths)
.components(Some(components))
.build();
let splitter = SplitSchemasByTag::new("common.yaml");
let usage = splitter.analyze_schema_usage(&spec);
assert!(
usage
.get("UserId")
.map(|t| t.contains("users"))
.unwrap_or(false)
);
}
#[test]
fn should_analyze_non_get_operations() {
let user_schema = ObjectBuilder::new()
.property(
"id",
ObjectBuilder::new().schema_type(utoipa::openapi::Type::Integer),
)
.build();
let order_schema = ObjectBuilder::new()
.property(
"id",
ObjectBuilder::new().schema_type(utoipa::openapi::Type::Integer),
)
.build();
let mut components = Components::new();
components
.schemas
.insert("User".to_string(), RefOr::T(user_schema.into()));
components
.schemas
.insert("Order".to_string(), RefOr::T(order_schema.into()));
let update_user = OperationBuilder::new()
.tags(Some(vec!["users".to_string()]))
.response(
"200",
ResponseBuilder::new()
.content(
"application/json",
ContentBuilder::new()
.schema(Some(RefOr::Ref(Ref::new("#/components/schemas/User"))))
.build(),
)
.build(),
)
.build();
let delete_order = OperationBuilder::new()
.tags(Some(vec!["orders".to_string()]))
.response(
"200",
ResponseBuilder::new()
.content(
"application/json",
ContentBuilder::new()
.schema(Some(RefOr::Ref(Ref::new("#/components/schemas/Order"))))
.build(),
)
.build(),
)
.build();
let mut paths = utoipa::openapi::Paths::new();
paths.paths.insert(
"/users/{id}".to_string(),
PathItemBuilder::new()
.operation(utoipa::openapi::HttpMethod::Put, update_user)
.build(),
);
paths.paths.insert(
"/orders/{id}".to_string(),
PathItemBuilder::new()
.operation(utoipa::openapi::HttpMethod::Delete, delete_order)
.build(),
);
let spec = OpenApiBuilder::new()
.paths(paths)
.components(Some(components))
.build();
let splitter = SplitSchemasByTag::new("common.yaml");
let usage = splitter.analyze_schema_usage(&spec);
assert!(
usage
.get("User")
.map(|t| t.contains("users"))
.unwrap_or(false)
);
assert!(
usage
.get("Order")
.map(|t| t.contains("orders"))
.unwrap_or(false)
);
}
#[test]
fn should_preserve_security_schemes_after_split() {
use utoipa::openapi::security::{HttpAuthScheme, HttpBuilder, SecurityScheme};
let spec = create_test_spec();
let mut spec_with_security = spec;
let mut security_schemes = BTreeMap::new();
security_schemes.insert(
"bearer_auth".to_string(),
SecurityScheme::Http(
HttpBuilder::new()
.scheme(HttpAuthScheme::Bearer)
.bearer_format("JWT")
.build(),
),
);
if let Some(ref mut components) = spec_with_security.components {
components.security_schemes = security_schemes;
}
let splitter = SplitSchemasByTag::new("common.yaml");
let result = splitter.split(spec_with_security);
let main_components = result
.main
.components
.as_ref()
.expect("should have components");
assert!(main_components.security_schemes.contains_key("bearer_auth"));
}
#[test]
fn should_skip_operations_without_tags() {
let user_schema = ObjectBuilder::new()
.property(
"id",
ObjectBuilder::new().schema_type(utoipa::openapi::Type::Integer),
)
.build();
let untagged_schema = ObjectBuilder::new()
.property(
"data",
ObjectBuilder::new().schema_type(utoipa::openapi::Type::String),
)
.build();
let mut components = Components::new();
components
.schemas
.insert("User".to_string(), RefOr::T(user_schema.into()));
components
.schemas
.insert("Untagged".to_string(), RefOr::T(untagged_schema.into()));
let get_user = OperationBuilder::new()
.tags(Some(vec!["users".to_string()]))
.response(
"200",
ResponseBuilder::new()
.content(
"application/json",
ContentBuilder::new()
.schema(Some(RefOr::Ref(Ref::new("#/components/schemas/User"))))
.build(),
)
.build(),
)
.build();
let get_health = OperationBuilder::new()
.response(
"200",
ResponseBuilder::new()
.content(
"application/json",
ContentBuilder::new()
.schema(Some(RefOr::Ref(Ref::new("#/components/schemas/Untagged"))))
.build(),
)
.build(),
)
.build();
let mut paths = utoipa::openapi::Paths::new();
paths.paths.insert(
"/users".to_string(),
PathItemBuilder::new()
.operation(utoipa::openapi::HttpMethod::Get, get_user)
.build(),
);
paths.paths.insert(
"/health".to_string(),
PathItemBuilder::new()
.operation(utoipa::openapi::HttpMethod::Get, get_health)
.build(),
);
let spec = OpenApiBuilder::new()
.paths(paths)
.components(Some(components))
.build();
let splitter = SplitSchemasByTag::new("common.yaml");
let usage = splitter.analyze_schema_usage(&spec);
assert!(usage.contains_key("User"));
assert!(!usage.contains_key("Untagged"));
}
#[test]
fn should_handle_spec_without_components() {
let mut paths = utoipa::openapi::Paths::new();
paths.paths.insert(
"/health".to_string(),
PathItemBuilder::new()
.operation(
utoipa::openapi::HttpMethod::Get,
OperationBuilder::new()
.tags(Some(vec!["health".to_string()]))
.response("200", ResponseBuilder::new().description("OK").build())
.build(),
)
.build(),
);
let spec = OpenApiBuilder::new().paths(paths).build();
let splitter = ExtractSchemasByPredicate::new("errors.yaml", |name| name.contains("Error"));
let result = splitter.split(spec);
assert!(result.is_unsplit());
}
#[test]
fn should_collect_schemas_from_request_bodies() {
use utoipa::openapi::request_body::RequestBodyBuilder;
let create_user_schema = ObjectBuilder::new()
.property(
"name",
ObjectBuilder::new().schema_type(utoipa::openapi::Type::String),
)
.build();
let user_schema = ObjectBuilder::new()
.property(
"id",
ObjectBuilder::new().schema_type(utoipa::openapi::Type::Integer),
)
.build();
let mut components = Components::new();
components.schemas.insert(
"CreateUser".to_string(),
RefOr::T(create_user_schema.into()),
);
components
.schemas
.insert("User".to_string(), RefOr::T(user_schema.into()));
let create_user = OperationBuilder::new()
.tags(Some(vec!["users".to_string()]))
.request_body(Some(
RequestBodyBuilder::new()
.content(
"application/json",
ContentBuilder::new()
.schema(Some(RefOr::Ref(Ref::new(
"#/components/schemas/CreateUser",
))))
.build(),
)
.build(),
))
.response(
"201",
ResponseBuilder::new()
.content(
"application/json",
ContentBuilder::new()
.schema(Some(RefOr::Ref(Ref::new("#/components/schemas/User"))))
.build(),
)
.build(),
)
.build();
let mut paths = utoipa::openapi::Paths::new();
paths.paths.insert(
"/users".to_string(),
PathItemBuilder::new()
.operation(utoipa::openapi::HttpMethod::Post, create_user)
.build(),
);
let spec = OpenApiBuilder::new()
.paths(paths)
.components(Some(components))
.build();
let splitter = SplitSchemasByTag::new("common.yaml");
let usage = splitter.analyze_schema_usage(&spec);
assert!(
usage
.get("CreateUser")
.map(|t| t.contains("users"))
.unwrap_or(false)
);
assert!(
usage
.get("User")
.map(|t| t.contains("users"))
.unwrap_or(false)
);
}
#[test]
fn should_place_files_in_schemas_dir() {
let spec = create_test_spec();
let splitter = SplitSchemasByTag::new("common.yaml").with_schemas_dir("schemas");
let result = splitter.split(spec);
for fragment in &result.fragments {
assert!(
fragment.path.starts_with("schemas"),
"Fragment path {:?} should start with 'schemas'",
fragment.path
);
}
let common_fragment = result.fragments.iter().find(|f| {
f.path
.file_name()
.map(|n| n.to_string_lossy().contains("common"))
.unwrap_or(false)
});
if let Some(fragment) = common_fragment {
assert_eq!(
fragment.path,
PathBuf::from("schemas/common.yaml"),
"Common file should be in schemas directory"
);
}
}
}