fix(kube-core): Use the transformer originally written in kube-rs#1839

Schema tests included

Author: NickLarsenNZ

kube-core/src/schema.rs

@@ -3,15 +3,20 @@
 //! [`CustomResourceDefinition`]: `k8s_openapi::apiextensions_apiserver::pkg::apis::apiextensions::v1::CustomResourceDefinition`
 
 // Used in docs
-#[allow(unused_imports)] use schemars::generate::SchemaSettings;
+#[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};
+use std::{
+    collections::{BTreeMap, BTreeSet, btree_map::Entry},
+    ops::DerefMut as _,
+    sync::LazyLock,
+};
 
 /// schemars [`Visitor`] that rewrites a [`Schema`] to conform to Kubernetes' "structural schema" rules
 ///
@@ -375,40 +380,220 @@ impl Transform for StructuralSchemaRewriter {
 }
 
 impl Transform for OptionalEnum {
-    fn transform(&mut self, schema: &mut schemars::Schema) {
-        transform_subschemas(self, schema);
+    fn transform(&mut self, transform_schema: &mut schemars::Schema) {
+        transform_subschemas(self, transform_schema);
 
-        let Some(obj) = schema.as_object_mut().filter(|o| o.len() == 1) else {
+        let Some(mut schema) = serde_json::from_value(transform_schema.clone().to_value()).ok() else {
             return;
         };
 
-        let arr = obj
-            .get("anyOf")
-            .iter()
-            .flat_map(|any_of| any_of.as_array())
-            .last()
-            .cloned()
-            .unwrap_or_default();
+        hoist_any_of_subschema_with_a_nullable_variant(&mut schema);
 
-        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("enum")
-            && !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));
+        // 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;
+            }
         }
     }
 }
 
+/// This is the signature for the `null` variant produced by schemars.
+static NULL_SCHEMA: LazyLock<Value> = LazyLock::new(|| {
+    serde_json::json!({
+        "enum": [null],
+        "nullable": true
+    })
+});
+
+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_schema
+    else {
+        return;
+    };
+
+    let SubschemaValidation {
+        any_of: Some(any_of),
+        one_of: None,
+    } = subschemas.deref_mut()
+    else {
+        return;
+    };
+
+    if any_of.len() != 2 {
+        return;
+    }
+
+    let entry_is_null: [bool; 2] = any_of
+        .iter()
+        .map(|schema| serde_json::to_value(schema).expect("schema should be able to convert to JSON"))
+        .map(|value| value == *NULL_SCHEMA)
+        .collect::<Vec<_>>()
+        .try_into()
+        .expect("there should be exactly 2 elements. We checked earlier");
+
+    // Get the `any_of` subschema that isn't the null entry
+    let subschema_to_hoist = match entry_is_null {
+        [true, false] => &any_of[1],
+        [false, true] => &any_of[0],
+        _ => return,
+    };
+
+    // At this point, we can be reasonably sure we need to hoist the non-null
+    // anyOf subschema up to the schema level, and unset the anyOf field.
+    // From here, anything that looks wrong will panic instead of return.
+    // TODO (@NickLarsenNZ): Return errors instead of panicking, leave panicking up to the
+    // infallible schemars::Transform
+    let Schema::Object(to_hoist) = subschema_to_hoist else {
+        panic!("the non-null anyOf subschema is a bool. That is not expected here");
+    };
+
+    let mut to_hoist = to_hoist.clone();
+
+    // Move the metadata into the subschema before hoisting (unless it already has metadata set)
+    to_hoist.metadata = to_hoist.metadata.or_else(|| kube_schema.metadata.take());
+
+    // Replace the schema with the non-null subschema
+    *kube_schema = to_hoist;
+
+    // Set the schema to nullable (as we know we matched the null variant earlier)
+    kube_schema.extensions.insert("nullable".to_owned(), true.into());
+}
+
+#[cfg(test)]
+mod tests {
+    use assert_json_diff::assert_json_eq;
+
+    use super::*;
+
+    #[test]
+    fn optional_tagged_enum_with_unit_variants() {
+        let original_value = serde_json::json!({
+            "anyOf": [
+                {
+                    "description": "A very simple enum with empty variants",
+                    "oneOf": [
+                        {
+                            "type": "string",
+                            "enum": [
+                                "C",
+                                "D"
+                            ]
+                        },
+                        {
+                            "description": "First variant doc-comment",
+                            "type": "string",
+                            "enum": [
+                                "A"
+                            ]
+                        },
+                        {
+                            "description": "Second variant doc-comment",
+                            "type": "string",
+                            "enum": [
+                                "B"
+                            ]
+                        }
+                    ]
+                },
+                {
+                    "enum": [
+                        null
+                    ],
+                    "nullable": true
+                }
+            ]
+        });
+
+        let expected_converted_value = serde_json::json!({
+            "description": "A very simple enum with empty variants",
+            "nullable": true,
+            "oneOf": [
+                {
+                    "type": "string",
+                    "enum": [
+                        "C",
+                        "D"
+                    ]
+                },
+                {
+                    "description": "First variant doc-comment",
+                    "type": "string",
+                    "enum": [
+                    "A"
+                ]
+                },
+                {
+                    "description": "Second variant doc-comment",
+                    "type": "string",
+                    "enum": [
+                        "B"
+                    ]
+                }
+            ]
+        });
+
+        let original: SchemaObject = serde_json::from_value(original_value).expect("valid JSON");
+        let expected_converted: SchemaObject =
+            serde_json::from_value(expected_converted_value).expect("valid JSON");
+
+        let mut actual_converted = original.clone();
+        hoist_any_of_subschema_with_a_nullable_variant(&mut actual_converted);
+
+        assert_json_eq!(actual_converted, expected_converted);
+    }
+
+    #[test]
+    fn optional_tagged_enum_with_unit_variants_but_also_an_existing_description() {
+        let original_value = serde_json::json!({
+            "description": "This comment will be lost",
+            "anyOf": [
+                {
+                    "description": "A very simple enum with empty variants",
+                    "type": "string",
+                    "enum": [
+                        "C",
+                        "D",
+                        "A",
+                        "B"
+                    ]
+                },
+                {
+                    "enum": [
+                        null
+                    ],
+                    "nullable": true
+                }
+            ]
+        });
+
+        let expected_converted_value = serde_json::json!({
+            "description": "A very simple enum with empty variants",
+            "nullable": true,
+            "type": "string",
+            "enum": [
+                "C",
+                "D",
+                "A",
+                "B"
+            ]
+        });
+
+        let original: SchemaObject = serde_json::from_value(original_value).expect("valid JSON");
+        let expected_converted: SchemaObject =
+            serde_json::from_value(expected_converted_value).expect("valid JSON");
+
+        let mut actual_converted = original.clone();
+        hoist_any_of_subschema_with_a_nullable_variant(&mut actual_converted);
+
+        assert_json_eq!(actual_converted, expected_converted);
+    }
+}
+
 impl Transform for OptionalIntOrString {
     fn transform(&mut self, schema: &mut schemars::Schema) {
         transform_subschemas(self, schema);