fix: Generate valid CRDs for fields with Optional enums and doc-comments

kube-core/src/schema.rs → kube-core/src/schema/mod.rs

@@ -11,7 +11,21 @@ use schemars::{
 };
 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},
+    sync::LazyLock,
+};
+
+mod transform_optional_enum;
+pub use transform_optional_enum::OptionalEnum;
+
+/// 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
+    })
+});
 
 /// schemars [`Visitor`] that rewrites a [`Schema`] to conform to Kubernetes' "structural schema" rules
 ///
@@ -30,42 +44,6 @@ use std::collections::{BTreeMap, BTreeSet, btree_map::Entry};
 #[derive(Debug, Clone)]
 pub struct StructuralSchemaRewriter;
 
-/// Recursively restructures JSON Schema objects so that the Option<Enum> object
-/// is returned per k8s CRD schema expectations.
-///
-/// In kube 2.x the schema output behavior for `Option<Enum>` types changed.
-///
-/// Previously given an enum like:
-///
-/// ```rust
-/// enum LogLevel {
-///     Debug,
-///     Info,
-///     Error,
-/// }
-/// ```
-///
-/// The following would be generated for Optional<LogLevel>:
-///
-/// ```json
-/// { "enum": ["Debug", "Info", "Error"], "type": "string", "nullable": true }
-/// ```
-///
-/// Now, schemars generates `anyOf` for `Option<LogLevel>` like:
-///
-/// ```json
-/// {
-///   "anyOf": [
-///     { "enum": ["Debug", "Info", "Error"], "type": "string" },
-///     { "enum": [null], "nullable": true }
-///   ]
-/// }
-/// ```
-///
-/// This transform implementation prevents this specific case from happening.
-#[derive(Debug, Clone, Default)]
-pub struct OptionalEnum;
-
 /// Recursively restructures JSON Schema objects so that the `Option<T>` object
 /// where `T` uses `x-kubernetes-int-or-string` is returned per k8s CRD schema expectations.
 ///
@@ -374,41 +352,6 @@ impl Transform for StructuralSchemaRewriter {
     }
 }
 
-impl Transform for OptionalEnum {
-    fn transform(&mut self, schema: &mut schemars::Schema) {
-        transform_subschemas(self, schema);
-
-        let Some(obj) = schema.as_object_mut().filter(|o| o.len() == 1) 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("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));
-        }
-    }
-}
-
 impl Transform for OptionalIntOrString {
     fn transform(&mut self, schema: &mut schemars::Schema) {
         transform_subschemas(self, schema);

kube-core/src/schema/transform_optional_enum.rs

@@ -0,0 +1,248 @@
+use std::ops::DerefMut as _;
+
+use schemars::transform::{Transform, transform_subschemas};
+
+use crate::schema::{NULL_SCHEMA, Schema, SchemaObject, SubschemaValidation};
+
+/// Recursively restructures JSON Schema objects so that the Option<Enum> object
+/// is returned per k8s CRD schema expectations.
+///
+/// In kube 2.x the schema output behavior for `Option<Enum>` types changed.
+///
+/// Previously given an enum like:
+///
+/// ```rust
+/// enum LogLevel {
+///     Debug,
+///     Info,
+///     Error,
+/// }
+/// ```
+///
+/// The following would be generated for Optional<LogLevel>:
+///
+/// ```json
+/// { "enum": ["Debug", "Info", "Error"], "type": "string", "nullable": true }
+/// ```
+///
+/// Now, schemars generates `anyOf` for `Option<LogLevel>` like:
+///
+/// ```json
+/// {
+///   "anyOf": [
+///     { "enum": ["Debug", "Info", "Error"], "type": "string" },
+///     { "enum": [null], "nullable": true }
+///   ]
+/// }
+/// ```
+///
+/// This transform implementation prevents this specific case from happening.
+#[derive(Debug, Clone, Default)]
+pub struct OptionalEnum;
+
+impl Transform for OptionalEnum {
+    fn transform(&mut self, transform_schema: &mut schemars::Schema) {
+        transform_subschemas(self, transform_schema);
+
+        let Some(mut schema) = serde_json::from_value(transform_schema.clone().to_value()).ok() else {
+            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_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);
+    }
+}

kube-derive/tests/crd_schema_test.rs

@@ -78,6 +78,7 @@ struct FooSpec {
     #[x_kube(merge_strategy = ListMerge::Set)]
     x_kubernetes_set: Vec<String>,
 
+    /// This comment will be present unless the Gender enum has doc-comments
     optional_enum: Option<Gender>,
 }
 
@@ -415,6 +416,7 @@ fn test_crd_schema_matches_expected() {
                                                 "x-kubernetes-list-type": "set",
                                             },
                                             "optionalEnum": {
+                                                "description": "This comment will be present unless the Gender enum has doc-comments",
                                                 "nullable": true,
                                                 "type": "string",
                                                 "enum": [