chore(schema): Split StructuralSchemaRewriter into multiple transformers

Signed-off-by: Nick Larsen <[email protected]>

Author: NickLarsenNZ

kube-core/src/schema/mod.rs

@@ -18,11 +18,10 @@ use std::{
     sync::LazyLock,
 };
 
-use crate::schema::{
-    transform_any_of::hoist_any_of_subschema_with_a_nullable_variant,
-    transform_one_of::hoist_one_of_enum_with_unit_variants,
-    transform_optional_enum_with_null::remove_optional_enum_null_variant,
-    transform_properties::hoist_properties_for_any_of_subschemas,
+pub use crate::schema::{
+    transform_any_of::AnyOfSchemaRewriter, transform_one_of::OneOfSchemaRewriter,
+    transform_optional_enum_with_null::OptionalEnumSchemaRewriter,
+    transform_properties::PropertiesSchemaRewriter,
 };
 
 /// This is the signature for the `null` variant produced by schemars.
@@ -47,6 +46,7 @@ static NULL_SCHEMA: LazyLock<Value> = LazyLock::new(|| {
 ///
 /// The [`Visitor`] functions may panic if the transform could not be applied. For example,
 /// there must not be any overlapping properties between `oneOf` branches.
+// TODO (@NickLarsenNZ): Rename this transformer and update the docs above
 #[derive(Debug, Clone)]
 pub struct StructuralSchemaRewriter;
 
@@ -282,12 +282,6 @@ impl Transform for StructuralSchemaRewriter {
             None => return,
         };
 
-        // Hoist parts of the schema to make it valid for k8s
-        hoist_one_of_enum_with_unit_variants(&mut schema);
-        hoist_any_of_subschema_with_a_nullable_variant(&mut schema);
-        hoist_properties_for_any_of_subschemas(&mut schema);
-        remove_optional_enum_null_variant(&mut schema);
-
         // check for maps without with properties (i.e. flattened maps)
         // and allow these to persist dynamically
         if let Some(object) = &mut schema.object {

kube-core/src/schema/transform_any_of.rs

@@ -1,5 +1,7 @@
 use std::ops::DerefMut;
 
+use schemars::transform::Transform;
+
 use crate::schema::{Schema, SchemaObject, SubschemaValidation, NULL_SCHEMA};
 
 /// Replace the schema with the anyOf subschema and set to nullable when the
@@ -14,7 +16,34 @@ use crate::schema::{Schema, SchemaObject, SubschemaValidation, NULL_SCHEMA};
 ///   `null` entry, and `nullable` set to true).
 ///
 /// NOTE: This should work regardless of whether other hoisting has been performed or not.
-pub(crate) fn hoist_any_of_subschema_with_a_nullable_variant(kube_schema: &mut SchemaObject) {
+#[derive(Debug, Clone)]
+pub struct AnyOfSchemaRewriter;
+
+impl Transform for AnyOfSchemaRewriter {
+    fn transform(&mut self, transform_schema: &mut schemars::Schema) {
+        // Apply this (self) transform to all subschemas
+        schemars::transform::transform_subschemas(self, transform_schema);
+
+        let mut schema: SchemaObject = match serde_json::from_value(transform_schema.clone().to_value()).ok()
+        {
+            // TODO (@NickLarsenNZ): At this point, we are seeing duplicate `title` when printing schema as json.
+            // This is because `title` is specified under both `extensions` and `other`.
+            Some(schema) => schema,
+            None => return,
+        };
+
+        hoist_any_of_subschema_with_a_nullable_variant(&mut schema);
+
+        // Convert back to schemars::Schema
+        if let Ok(schema) = serde_json::to_value(schema) {
+            if let Ok(transformed) = serde_json::from_value(schema) {
+                *transform_schema = transformed;
+            }
+        }
+    }
+}
+
+fn hoist_any_of_subschema_with_a_nullable_variant(kube_schema: &mut SchemaObject) {
     // Run some initial checks in case there is nothing to do
     let SchemaObject {
         subschemas: Some(subschemas),

kube-core/src/schema/transform_one_of.rs

@@ -1,5 +1,7 @@
 use std::ops::DerefMut;
 
+use schemars::transform::Transform;
+
 use crate::schema::{Schema, SchemaObject, SubschemaValidation};
 
 /// Merge oneOf subschema enums and consts into a schema level enum.
@@ -17,7 +19,34 @@ use crate::schema::{Schema, SchemaObject, SubschemaValidation};
 ///
 /// NOTE: This should work regardless of whether other hoisting has been
 /// performed or not.
-pub(crate) fn hoist_one_of_enum_with_unit_variants(kube_schema: &mut SchemaObject) {
+#[derive(Debug, Clone)]
+pub struct OneOfSchemaRewriter;
+
+impl Transform for OneOfSchemaRewriter {
+    fn transform(&mut self, transform_schema: &mut schemars::Schema) {
+        // Apply this (self) transform to all subschemas
+        schemars::transform::transform_subschemas(self, transform_schema);
+
+        let mut schema: SchemaObject = match serde_json::from_value(transform_schema.clone().to_value()).ok()
+        {
+            // TODO (@NickLarsenNZ): At this point, we are seeing duplicate `title` when printing schema as json.
+            // This is because `title` is specified under both `extensions` and `other`.
+            Some(schema) => schema,
+            None => return,
+        };
+
+        hoist_one_of_enum_with_unit_variants(&mut schema);
+
+        // Convert back to schemars::Schema
+        if let Ok(schema) = serde_json::to_value(schema) {
+            if let Ok(transformed) = serde_json::from_value(schema) {
+                *transform_schema = transformed;
+            }
+        }
+    }
+}
+
+fn hoist_one_of_enum_with_unit_variants(kube_schema: &mut SchemaObject) {
     // Run some initial checks in case there is nothing to do
     let SchemaObject {
         subschemas: Some(subschemas),

kube-core/src/schema/transform_optional_enum_with_null.rs

@@ -1,3 +1,4 @@
+use schemars::transform::Transform;
 use serde_json::Value;
 
 use super::SchemaObject;
@@ -8,7 +9,34 @@ use super::SchemaObject;
 ///
 /// NOTE: The trailing null is removed because it's not needed by Kubernetes
 /// and makes the CRD more compact by removing redundant information.
-pub(crate) fn remove_optional_enum_null_variant(kube_schema: &mut SchemaObject) {
+#[derive(Debug, Clone)]
+pub struct OptionalEnumSchemaRewriter;
+
+impl Transform for OptionalEnumSchemaRewriter {
+    fn transform(&mut self, transform_schema: &mut schemars::Schema) {
+        // Apply this (self) transform to all subschemas
+        schemars::transform::transform_subschemas(self, transform_schema);
+
+        let mut schema: SchemaObject = match serde_json::from_value(transform_schema.clone().to_value()).ok()
+        {
+            // TODO (@NickLarsenNZ): At this point, we are seeing duplicate `title` when printing schema as json.
+            // This is because `title` is specified under both `extensions` and `other`.
+            Some(schema) => schema,
+            None => return,
+        };
+
+        remove_optional_enum_null_variant(&mut schema);
+
+        // Convert back to schemars::Schema
+        if let Ok(schema) = serde_json::to_value(schema) {
+            if let Ok(transformed) = serde_json::from_value(schema) {
+                *transform_schema = transformed;
+            }
+        }
+    }
+}
+
+fn remove_optional_enum_null_variant(kube_schema: &mut SchemaObject) {
     let SchemaObject {
         enum_values: Some(enum_values),
         extensions,

kube-core/src/schema/transform_properties.rs

@@ -1,3 +1,5 @@
+use schemars::transform::Transform;
+
 use crate::schema::{InstanceType, Metadata, Schema, SchemaObject, SingleOrVec, NULL_SCHEMA};
 
 /// Take oneOf or anyOf subschema properties and move them them into the schema
@@ -11,7 +13,34 @@ use crate::schema::{InstanceType, Metadata, Schema, SchemaObject, SingleOrVec, N
 /// - Each subschema has the type "object"
 ///
 /// NOTE: This should work regardless of whether other hoisting has been performed or not.
-pub(crate) fn hoist_properties_for_any_of_subschemas(kube_schema: &mut SchemaObject) {
+#[derive(Debug, Clone)]
+pub struct PropertiesSchemaRewriter;
+
+impl Transform for PropertiesSchemaRewriter {
+    fn transform(&mut self, transform_schema: &mut schemars::Schema) {
+        // Apply this (self) transform to all subschemas
+        schemars::transform::transform_subschemas(self, transform_schema);
+
+        let mut schema: SchemaObject = match serde_json::from_value(transform_schema.clone().to_value()).ok()
+        {
+            // TODO (@NickLarsenNZ): At this point, we are seeing duplicate `title` when printing schema as json.
+            // This is because `title` is specified under both `extensions` and `other`.
+            Some(schema) => schema,
+            None => return,
+        };
+
+        hoist_properties_for_any_of_subschemas(&mut schema);
+
+        // Convert back to schemars::Schema
+        if let Ok(schema) = serde_json::to_value(schema) {
+            if let Ok(transformed) = serde_json::from_value(schema) {
+                *transform_schema = transformed;
+            }
+        }
+    }
+}
+
+fn hoist_properties_for_any_of_subschemas(kube_schema: &mut SchemaObject) {
     // Run some initial checks in case there is nothing to do
     let SchemaObject {
         subschemas: Some(subschemas),

kube-derive/src/custom_resource.rs

@@ -634,6 +634,10 @@ pub(crate) fn derive(input: proc_macro2::TokenStream) -> proc_macro2::TokenStrea
                     s.meta_schema = None;
                 })
                 .with_transform(#schemars::transform::AddNullable::default())
+                .with_transform(#kube_core::schema::OneOfSchemaRewriter)
+                .with_transform(#kube_core::schema::AnyOfSchemaRewriter)
+                .with_transform(#kube_core::schema::PropertiesSchemaRewriter)
+                .with_transform(#kube_core::schema::OptionalEnumSchemaRewriter)
                 .with_transform(#kube_core::schema::StructuralSchemaRewriter)
                 .into_generator();
             let schema = generate.into_root_schema_for::<Self>();