DGS-24052 Handle anyOf/allOf in JSON transforms (#2237)

rayokota · web-flow · commit a0496cf6bfdd · 2026-04-27T08:17:13.000-07:00
* DGS-24052 Handle anyOf/allOf in JSON transforms * Update changelog * Fix style * Fix mypy
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -5,6 +5,7 @@
 ### Fixes
 
 - Fix URL joining in Python client (#2228)
+- Handle anyOf/allOf in JSON transforms (#2237)
 
 
 ## v2.14.0 - 2026-04-01
diff --git a/src/confluent_kafka/schema_registry/_async/mock_schema_registry_client.py b/src/confluent_kafka/schema_registry/_async/mock_schema_registry_client.py
@@ -152,7 +152,7 @@ def clear(self):
 
 class _AssociationStore(object):
 
-    def __init__(self):
+    def __init__(self) -> None:
         self.lock = Lock()
         # Key: resource_id -> List[Association]
         self.associations_by_resource_id: Dict[str, List[Association]] = defaultdict(list)
diff --git a/src/confluent_kafka/schema_registry/_sync/mock_schema_registry_client.py b/src/confluent_kafka/schema_registry/_sync/mock_schema_registry_client.py
@@ -152,7 +152,7 @@ def clear(self):
 
 class _AssociationStore(object):
 
-    def __init__(self):
+    def __init__(self) -> None:
         self.lock = Lock()
         # Key: resource_id -> List[Association]
         self.associations_by_resource_id: Dict[str, List[Association]] = defaultdict(list)
diff --git a/src/confluent_kafka/schema_registry/common/json_schema.py b/src/confluent_kafka/schema_registry/common/json_schema.py
@@ -88,20 +88,38 @@ def transform(
         finally:
             schema["type"] = original_type  # restore original type
     all_of = schema.get("allOf")
-    if all_of is not None:
-        subschema = _validate_subschemas(all_of, message, ref_registry, ref_resolver)
-        if subschema is not None:
-            return transform(ctx, subschema, ref_registry, ref_resolver, path, message, field_transform)
     any_of = schema.get("anyOf")
-    if any_of is not None:
-        subschema = _validate_subschemas(any_of, message, ref_registry, ref_resolver)
-        if subschema is not None:
-            return transform(ctx, subschema, ref_registry, ref_resolver, path, message, field_transform)
     one_of = schema.get("oneOf")
-    if one_of is not None:
-        subschema = _validate_subschemas(one_of, message, ref_registry, ref_resolver)
-        if subschema is not None:
-            return transform(ctx, subschema, ref_registry, ref_resolver, path, message, field_transform)
+    if all_of is not None or any_of is not None or one_of is not None:
+        if all_of is not None:
+            for subschema in all_of:
+                message = transform(ctx, subschema, ref_registry, ref_resolver, path, message, field_transform)
+        elif one_of is not None:
+            for subschema in one_of:
+                resolved = _validate_subschema(subschema, message, ref_registry, ref_resolver)
+                if resolved is not None:
+                    message = transform(ctx, resolved, ref_registry, ref_resolver, path, message, field_transform)
+                    break
+        elif any_of is not None:
+            for subschema in any_of:
+                resolved = _validate_subschema(subschema, message, ref_registry, ref_resolver)
+                if resolved is not None:
+                    message = transform(ctx, resolved, ref_registry, ref_resolver, path, message, field_transform)
+        # Also visit sibling properties/items at this level
+        # (siblings to allOf/anyOf/oneOf).
+        props = schema.get("properties")
+        if props is not None and isinstance(message, dict):
+            for prop_name, prop_schema in props.items():
+                if isinstance(prop_schema, dict):
+                    _transform_field(
+                        ctx, path, prop_name, message, prop_schema, ref_registry, ref_resolver, field_transform
+                    )
+        items = schema.get("items")
+        if items is not None and isinstance(message, list):
+            message = [
+                transform(ctx, items, ref_registry, ref_resolver, path, item, field_transform) for item in message
+            ]
+        return message
     items = schema.get("items")
     if items is not None:
         if isinstance(message, list):
@@ -197,23 +215,39 @@ def _validate_subschemas(
         The validated schema if the message is valid against the subschemas, otherwise None.
     """
     for subschema in subschemas:
-        if isinstance(subschema, dict):
-            try:
-                ref = subschema.get("$ref")
-                if ref is not None:
-                    resolved = resolver.lookup(ref)
-                    subschema = resolved.contents
-                    # Pass _resolver (not resolver) to use the new referencing library's
-                    # Resolver with correct context for nested $ref resolution
-                    validate(instance=message, schema=subschema, registry=registry, _resolver=resolved.resolver)
-                else:
-                    validate(instance=message, schema=subschema, registry=registry)
-                return subschema
-            except ValidationError:
-                pass
+        resolved = _validate_subschema(subschema, message, registry, resolver)
+        if resolved is not None:
+            return resolved
     return None
 
 
+def _validate_subschema(
+    subschema: JsonSchema,
+    message: JsonMessage,
+    registry: Registry,
+    resolver: Resolver,
+) -> Optional[JsonSchema]:
+    """
+    Validate the message against a single subschema.
+    Returns the resolved subschema (with $ref followed) if valid, otherwise None.
+    """
+    if not isinstance(subschema, dict):
+        return None
+    try:
+        ref = subschema.get("$ref")
+        if ref is not None:
+            resolved = resolver.lookup(ref)
+            subschema = resolved.contents
+            # Pass _resolver (not resolver) to use the new referencing library's
+            # Resolver with correct context for nested $ref resolution
+            validate(instance=message, schema=subschema, registry=registry, _resolver=resolved.resolver)
+        else:
+            validate(instance=message, schema=subschema, registry=registry)
+        return subschema
+    except ValidationError:
+        return None
+
+
 def get_type(schema: JsonSchema) -> FieldType:
     if isinstance(schema, bool):
         return FieldType.COMBINED
diff --git a/tests/schema_registry/_async/test_json_serdes.py b/tests/schema_registry/_async/test_json_serdes.py
@@ -582,6 +582,139 @@ async def test_json_cel_field_transform_with_def():
     assert obj2 == newobj
 
 
+async def test_json_cel_field_transform_all_of():
+    conf = {'url': _BASE_URL}
+    client = AsyncSchemaRegistryClient.new_client(conf)
+    ser_conf = {'auto.register.schemas': False, 'use.latest.version': True}
+    schema = {
+        "type": "object",
+        "properties": {
+            "pins": {
+                "type": "object",
+                "allOf": [
+                    {"properties": {"pin": {"confluent:tags": ["PII"], "type": ["string", "null"]}}},
+                    {"properties": {"npin": {"confluent:tags": ["PII"], "type": ["string", "null"]}}},
+                ],
+            }
+        },
+    }
+
+    rule = Rule(
+        "test-cel",
+        "",
+        RuleKind.TRANSFORM,
+        RuleMode.WRITE,
+        "CEL_FIELD",
+        ["PII"],
+        None,
+        "value + '-suffix'",
+        None,
+        None,
+        False,
+    )
+    await client.register_schema(_SUBJECT, Schema(json.dumps(schema), "JSON", [], None, RuleSet(None, [rule])))
+
+    obj = {'pins': {'pin': 'P123456789', 'npin': 'NP00012345678'}}
+    ser = await AsyncJSONSerializer(json.dumps(schema), client, conf=ser_conf)
+    ser_ctx = SerializationContext(_TOPIC, MessageField.VALUE)
+    obj_bytes = await ser(obj, ser_ctx)
+
+    obj2 = {'pins': {'pin': 'P123456789-suffix', 'npin': 'NP00012345678-suffix'}}
+    deser = await AsyncJSONDeserializer(None, schema_registry_client=client)
+    newobj = await deser(obj_bytes, ser_ctx)
+    assert obj2 == newobj
+
+
+async def test_json_cel_field_transform_nested_any_of():
+    conf = {'url': _BASE_URL}
+    client = AsyncSchemaRegistryClient.new_client(conf)
+    ser_conf = {'auto.register.schemas': False, 'use.latest.version': True}
+    schema = {
+        "type": "object",
+        "properties": {
+            "pins": {
+                "type": "object",
+                "anyOf": [
+                    {"properties": {"pin": {"confluent:tags": ["PII"], "type": ["string", "null"]}}},
+                    {"properties": {"npin": {"confluent:tags": ["PII"], "type": ["string", "null"]}}},
+                ],
+            }
+        },
+    }
+
+    rule = Rule(
+        "test-cel",
+        "",
+        RuleKind.TRANSFORM,
+        RuleMode.WRITE,
+        "CEL_FIELD",
+        ["PII"],
+        None,
+        "value + '-suffix'",
+        None,
+        None,
+        False,
+    )
+    await client.register_schema(_SUBJECT, Schema(json.dumps(schema), "JSON", [], None, RuleSet(None, [rule])))
+
+    obj = {'pins': {'pin': 'P123456789', 'npin': 'NP00012345678'}}
+    ser = await AsyncJSONSerializer(json.dumps(schema), client, conf=ser_conf)
+    ser_ctx = SerializationContext(_TOPIC, MessageField.VALUE)
+    obj_bytes = await ser(obj, ser_ctx)
+
+    obj2 = {'pins': {'pin': 'P123456789-suffix', 'npin': 'NP00012345678-suffix'}}
+    deser = await AsyncJSONDeserializer(None, schema_registry_client=client)
+    newobj = await deser(obj_bytes, ser_ctx)
+    assert obj2 == newobj
+
+
+async def test_json_cel_field_transform_sibling_any_of():
+    conf = {'url': _BASE_URL}
+    client = AsyncSchemaRegistryClient.new_client(conf)
+    ser_conf = {'auto.register.schemas': False, 'use.latest.version': True}
+    schema = {
+        "type": "object",
+        "properties": {
+            "pins": {
+                "type": "object",
+                "anyOf": [
+                    {"required": ["pin"]},
+                    {"required": ["npin"]},
+                ],
+                "properties": {
+                    "pin": {"confluent:tags": ["PII"], "type": ["string", "null"]},
+                    "npin": {"confluent:tags": ["PII"], "type": ["string", "null"]},
+                },
+            }
+        },
+    }
+
+    rule = Rule(
+        "test-cel",
+        "",
+        RuleKind.TRANSFORM,
+        RuleMode.WRITE,
+        "CEL_FIELD",
+        ["PII"],
+        None,
+        "value + '-suffix'",
+        None,
+        None,
+        False,
+    )
+    await client.register_schema(_SUBJECT, Schema(json.dumps(schema), "JSON", [], None, RuleSet(None, [rule])))
+
+    obj = {'pins': {'pin': 'P123456789', 'npin': 'NP00012345678'}}
+    ser = await AsyncJSONSerializer(json.dumps(schema), client, conf=ser_conf)
+    ser_ctx = SerializationContext(_TOPIC, MessageField.VALUE)
+    obj_bytes = await ser(obj, ser_ctx)
+
+    obj2 = {'pins': {'pin': 'P123456789-suffix', 'npin': 'NP00012345678-suffix'}}
+    deser = await AsyncJSONDeserializer(None, schema_registry_client=client)
+    newobj = await deser(obj_bytes, ser_ctx)
+    assert obj2 == newobj
+
+
 async def test_json_cel_field_transform_complex():
     conf = {'url': _BASE_URL}
     client = AsyncSchemaRegistryClient.new_client(conf)
diff --git a/tests/schema_registry/_sync/test_json_serdes.py b/tests/schema_registry/_sync/test_json_serdes.py
