#[allow(unused_imports)] use schemars::generate::SchemaSettings;
use schemars::{
JsonSchema,
transform::{Transform, transform_subschemas},
};
use serde::{Deserialize, Serialize};
use serde_json::{Value, json};
use std::collections::{BTreeMap, BTreeSet, btree_map::Entry};
#[derive(Debug, Clone)]
pub struct StructuralSchemaRewriter;
#[derive(Debug, Clone, Default)]
pub struct OptionalEnum;
#[derive(Debug, Clone, Default)]
pub struct OptionalIntOrString;
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema)]
#[serde(untagged)]
enum Schema {
Bool(bool),
Object(SchemaObject),
}
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Default, JsonSchema)]
#[serde(rename_all = "camelCase", default)]
struct SchemaObject {
#[serde(flatten, deserialize_with = "skip_if_default")]
metadata: Option<Box<Metadata>>,
#[serde(rename = "type", skip_serializing_if = "Option::is_none")]
instance_type: Option<SingleOrVec<InstanceType>>,
#[serde(skip_serializing_if = "Option::is_none")]
format: Option<String>,
#[serde(rename = "enum", skip_serializing_if = "Option::is_none")]
enum_values: Option<Vec<Value>>,
#[serde(flatten, deserialize_with = "skip_if_default")]
subschemas: Option<Box<SubschemaValidation>>,
#[serde(flatten, deserialize_with = "skip_if_default")]
array: Option<Box<ArrayValidation>>,
#[serde(flatten, deserialize_with = "skip_if_default")]
object: Option<Box<ObjectValidation>>,
#[serde(flatten)]
extensions: BTreeMap<String, Value>,
#[serde(flatten)]
other: Value,
}
fn allow_null<'de, D>(de: D) -> Result<Option<Value>, D::Error>
where
D: serde::Deserializer<'de>,
{
Value::deserialize(de).map(Option::Some)
}
fn skip_if_default<'de, D, T>(deserializer: D) -> Result<Option<Box<T>>, D::Error>
where
D: serde::Deserializer<'de>,
T: Deserialize<'de> + Default + PartialEq,
{
let value = T::deserialize(deserializer)?;
if value == T::default() {
Ok(None)
} else {
Ok(Some(Box::new(value)))
}
}
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Default, JsonSchema)]
#[serde(rename_all = "camelCase", default)]
struct Metadata {
#[serde(skip_serializing_if = "Option::is_none")]
description: Option<String>,
#[serde(skip_serializing_if = "Option::is_none", deserialize_with = "allow_null")]
default: Option<Value>,
}
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Default, JsonSchema)]
#[serde(rename_all = "camelCase", default)]
struct SubschemaValidation {
#[serde(skip_serializing_if = "Option::is_none")]
any_of: Option<Vec<Schema>>,
#[serde(skip_serializing_if = "Option::is_none")]
one_of: Option<Vec<Schema>>,
}
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Default, JsonSchema)]
#[serde(rename_all = "camelCase", default)]
struct ArrayValidation {
#[serde(skip_serializing_if = "Option::is_none")]
items: Option<SingleOrVec<Schema>>,
#[serde(skip_serializing_if = "Option::is_none")]
additional_items: Option<Box<Schema>>,
#[serde(skip_serializing_if = "Option::is_none")]
max_items: Option<u32>,
#[serde(skip_serializing_if = "Option::is_none")]
min_items: Option<u32>,
#[serde(skip_serializing_if = "Option::is_none")]
unique_items: Option<bool>,
#[serde(skip_serializing_if = "Option::is_none")]
contains: Option<Box<Schema>>,
}
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Default, JsonSchema)]
#[serde(rename_all = "camelCase", default)]
struct ObjectValidation {
#[serde(skip_serializing_if = "Option::is_none")]
max_properties: Option<u32>,
#[serde(skip_serializing_if = "Option::is_none")]
min_properties: Option<u32>,
#[serde(skip_serializing_if = "BTreeSet::is_empty")]
required: BTreeSet<String>,
#[serde(skip_serializing_if = "BTreeMap::is_empty")]
properties: BTreeMap<String, Schema>,
#[serde(skip_serializing_if = "BTreeMap::is_empty")]
pattern_properties: BTreeMap<String, Schema>,
#[serde(skip_serializing_if = "Option::is_none")]
additional_properties: Option<Box<Schema>>,
#[serde(skip_serializing_if = "Option::is_none")]
property_names: Option<Box<Schema>>,
}
#[derive(Serialize, Deserialize, Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, JsonSchema)]
#[serde(rename_all = "camelCase")]
enum InstanceType {
Null,
Boolean,
Object,
Array,
Number,
String,
Integer,
}
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, PartialOrd, Ord, JsonSchema)]
#[serde(untagged)]
enum SingleOrVec<T> {
Single(Box<T>),
Vec(Vec<T>),
}
impl Transform for StructuralSchemaRewriter {
fn transform(&mut self, transform_schema: &mut schemars::Schema) {
schemars::transform::transform_subschemas(self, transform_schema);
let mut schema: SchemaObject = match serde_json::from_value(transform_schema.clone().to_value()).ok()
{
Some(schema) => schema,
None => return,
};
if let Some(subschemas) = &mut schema.subschemas {
if let Some(one_of) = subschemas.one_of.as_mut() {
hoist_subschema_properties(one_of, &mut schema.object, &mut schema.instance_type, true);
hoist_subschema_enum_values(one_of, &mut schema.enum_values, &mut schema.instance_type);
if one_of.is_empty() {
subschemas.one_of = None;
}
}
if let Some(any_of) = &mut subschemas.any_of {
hoist_subschema_properties(any_of, &mut schema.object, &mut schema.instance_type, false);
}
}
if let Some(object) = &mut schema.object
&& !object.properties.is_empty()
{
match object.additional_properties.as_deref() {
Some(&Schema::Bool(true)) => {
object.additional_properties = None;
schema
.extensions
.insert("x-kubernetes-preserve-unknown-fields".into(), true.into());
}
Some(&Schema::Bool(false)) => {
object.additional_properties = None;
}
_ => {}
}
}
if let Some(array) = &mut schema.array {
array.unique_items = None;
}
if let Ok(schema) = serde_json::to_value(schema)
&& let Ok(transformed) = serde_json::from_value(schema)
{
*transform_schema = transformed;
}
}
}
impl Transform for OptionalEnum {
fn transform(&mut self, schema: &mut schemars::Schema) {
transform_subschemas(self, schema);
let Some(obj) = schema.as_object_mut() else {
return;
};
let arr = obj
.get("anyOf")
.iter()
.flat_map(|any_of| any_of.as_array())
.last()
.cloned()
.unwrap_or_default();
let [first, second] = arr.as_slice() else {
return;
};
let (Some(first), Some(second)) = (first.as_object(), second.as_object()) else {
return;
};
if !first.contains_key("nullable")
&& second.get("enum") == Some(&json!([null]))
&& second.get("nullable") == Some(&json!(true))
{
obj.remove("anyOf");
obj.append(&mut first.clone());
obj.insert("nullable".to_string(), Value::Bool(true));
}
}
}
impl Transform for OptionalIntOrString {
fn transform(&mut self, schema: &mut schemars::Schema) {
transform_subschemas(self, schema);
let Some(obj) = schema.as_object_mut() else {
return;
};
let required: BTreeSet<String> = obj
.get("required")
.and_then(|v| v.as_array())
.map(|arr| arr.iter().filter_map(|v| v.as_str().map(String::from)).collect())
.unwrap_or_default();
let Some(properties) = obj.get_mut("properties").and_then(|v| v.as_object_mut()) else {
return;
};
for (name, prop_schema) in properties.iter_mut() {
if required.contains(name) {
continue;
}
let Some(prop_obj) = prop_schema.as_object_mut() else {
continue;
};
if prop_obj.get("x-kubernetes-int-or-string") == Some(&json!(true))
&& !prop_obj.contains_key("nullable")
{
prop_obj.insert("nullable".to_string(), Value::Bool(true));
}
}
}
}
fn hoist_subschema_enum_values(
subschemas: &mut Vec<Schema>,
common_enum_values: &mut Option<Vec<serde_json::Value>>,
instance_type: &mut Option<SingleOrVec<InstanceType>>,
) {
subschemas.retain(|variant| {
if let Schema::Object(SchemaObject {
instance_type: variant_type,
enum_values: Some(variant_enum_values),
..
}) = variant
{
if let Some(variant_type) = variant_type {
match instance_type {
None => *instance_type = Some(variant_type.clone()),
Some(tpe) => {
if tpe != variant_type {
panic!("Enum variant set {variant_enum_values:?} has type {variant_type:?} but was already defined as {instance_type:?}. The instance type must be equal for all subschema variants.")
}
}
}
}
common_enum_values
.get_or_insert_with(Vec::new)
.extend(variant_enum_values.iter().cloned());
false
} else {
true
}
})
}
fn hoist_subschema_properties(
subschemas: &mut Vec<Schema>,
common_obj: &mut Option<Box<ObjectValidation>>,
instance_type: &mut Option<SingleOrVec<InstanceType>>,
push_description_to_property: bool,
) {
for variant in subschemas {
if let Schema::Object(SchemaObject {
instance_type: variant_type,
object: Some(variant_obj),
metadata: variant_metadata,
..
}) = variant
{
let common_obj = common_obj.get_or_insert_with(Box::<ObjectValidation>::default);
if let Some(variant_metadata) = variant_metadata
&& let Some(description) = std::mem::take(&mut variant_metadata.description)
&& let Some(Schema::Object(variant_object)) = only_item(variant_obj.properties.values_mut())
{
let metadata = variant_object
.metadata
.get_or_insert_with(Box::<Metadata>::default);
if push_description_to_property {
metadata.description = Some(description);
}
}
let variant_properties = std::mem::take(&mut variant_obj.properties);
for (property_name, property) in variant_properties {
match common_obj.properties.entry(property_name) {
Entry::Vacant(entry) => {
entry.insert(property);
}
Entry::Occupied(entry) => {
if &property != entry.get() {
panic!(
"Property {:?} has the schema {:?} but was already defined as {:?} in another subschema. The schemas for a property used in multiple subschemas must be identical",
entry.key(),
&property,
entry.get()
);
}
}
}
}
variant_obj.additional_properties = None;
merge_metadata(instance_type, variant_type.take());
}
else if let Schema::Object(SchemaObject {
metadata: variant_metadata,
instance_type: variant_type,
enum_values: None,
subschemas: None,
array: None,
object: None,
..
}) = variant
{
std::mem::take(&mut *variant_type);
std::mem::take(&mut *variant_metadata);
}
}
}
fn only_item<I: Iterator>(mut i: I) -> Option<I::Item> {
let item = i.next()?;
if i.next().is_some() {
return None;
}
Some(item)
}
fn merge_metadata(
instance_type: &mut Option<SingleOrVec<InstanceType>>,
variant_type: Option<SingleOrVec<InstanceType>>,
) {
match (instance_type, variant_type) {
(_, None) => {}
(common_type @ None, variant_type) => {
*common_type = variant_type;
}
(Some(common_type), Some(variant_type)) => {
if *common_type != variant_type {
panic!(
"variant defined type {variant_type:?}, conflicting with existing type {common_type:?}"
);
}
}
}
}