DGS-22899 Fix support for wrapped Avro unions (#402)

Author: rayokota

schemaregistry/serde/avro.ts

@@ -296,11 +296,15 @@ async function transform(ctx: RuleContext, schema: Type, msg: any, fieldTransfor
   switch (schema.typeName) {
     case 'union:unwrapped':
     case 'union:wrapped':
-      const subschema = resolveUnion(schema, msg)
+      let [subschema, submsg] = resolveUnion(schema, msg)
       if (subschema == null) {
-        return null
+        return msg
       }
-      return await transform(ctx, subschema, msg, fieldTransform)
+      submsg = await transform(ctx, subschema, submsg, fieldTransform)
+      if (schema.typeName === 'union:wrapped') {
+        return {[subschema.branchName!]: submsg}
+      }
+      return submsg
     case 'array':
       const arraySchema = schema as ArrayType
       const array = msg as any[]
@@ -408,15 +412,15 @@ function disjoint(slice1: Set<string>, map1: Set<string>): boolean {
   return true
 }
 
-function resolveUnion(schema: Type, msg: any): Type | null {
+function resolveUnion(schema: Type, msg: any): [Type | null, any] {
   let unionTypes = null
   if (schema.typeName === 'union:unwrapped') {
     const union = schema as UnwrappedUnionType
     unionTypes = union.types.slice()
     if (unionTypes != null) {
       for (let i = 0; i < unionTypes.length; i++) {
         if (unionTypes[i].isValid(msg)) {
-          return unionTypes[i]
+          return [unionTypes[i], msg]
         }
       }
     }
@@ -429,13 +433,15 @@ function resolveUnion(schema: Type, msg: any): Type | null {
         let name = keys[0]
         for (let i = 0; i < unionTypes.length; i++) {
           if (unionTypes[i].branchName === name) {
-            return unionTypes[i]
+            return [unionTypes[i], msg[name]]
           }
         }
+      } else {
+        throw new Error('wrapped unions require a name/value pair with the name as the type name')
       }
     }
   }
-  return null
+  return [null, msg]
 }
 
 function getInlineTags(info: SchemaInfo, deps: Map<string, string>): Map<string, Set<string>> {

schemaregistry/test/serde/avro.spec.ts

@@ -357,6 +357,61 @@ const complexNestedSchema = `
   ]
 }`;
 
+const wrappedUnionSchema = `{
+  "fields": [
+    {
+      "name": "id",
+      "type": "int"
+    },
+    {
+      "name": "result",
+      "type": [
+        "null",
+        {
+          "fields": [
+            {
+              "name": "code",
+              "type": "int"
+            },
+            {
+              "confluent:tags": [
+                "PII"
+              ],
+              "name": "secret",
+              "type": [
+                "null",
+                "string"
+              ]
+            }
+          ],
+          "name": "Data",
+          "type": "record"
+        },
+        {
+          "fields": [
+            {
+              "name": "code",
+              "type": "int"
+            },
+            {
+              "name": "reason",
+              "type": [
+                "null",
+                "string"
+              ]
+            }
+          ],
+          "name": "Error",
+          "type": "record"
+        }
+      ]
+    }
+  ],
+  "name": "Result",
+  "namespace": "com.acme",
+  "type": "record"
+}`;
+
 class FakeClock extends Clock {
   fixedNow: number = 0
 
@@ -1473,6 +1528,71 @@ describe('AvroSerializer', () => {
     expect(obj2.stringField).not.toEqual("hi");
     expect(obj2.bytesField).not.toEqual(Buffer.from([1, 2]));
   })
+  it('basic encryption with wrapped union', async () => {
+    let conf: ClientConfig = {
+      baseURLs: [baseURL],
+      cacheCapacity: 1000
+    }
+    let client = SchemaRegistryClient.newClient(conf)
+    let serConfig: AvroSerializerConfig = {
+      useLatestVersion: true,
+      ruleConfig: {
+        secret: 'mysecret'
+      }
+    }
+    let ser = new AvroSerializer(client, SerdeType.VALUE, serConfig)
+    let dekClient = fieldEncryptionExecutor.executor.client!
+
+    let encRule: Rule = {
+      name: 'test-encrypt',
+      kind: 'TRANSFORM',
+      mode: RuleMode.WRITEREAD,
+      type: 'ENCRYPT',
+      tags: ['PII'],
+      params: {
+        'encrypt.kek.name': 'kek1',
+        'encrypt.kms.type': 'local-kms',
+        'encrypt.kms.key.id': 'mykey',
+      },
+      onFailure: 'ERROR,NONE'
+    }
+    let ruleSet: RuleSet = {
+      domainRules: [encRule]
+    }
+
+    let info: SchemaInfo = {
+      schemaType: 'AVRO',
+      schema: wrappedUnionSchema,
+      ruleSet
+    }
+
+    await client.register(subject, info, false)
+
+    let obj = {
+      id: 123,
+      result: {
+        'com.acme.Data': {
+          code: 456,
+          secret: 'mypii'
+        }
+      }
+    }
+    let bytes = await ser.serialize(topic, obj)
+
+    // reset encrypted field
+    obj.result["com.acme.Data"].secret = 'mypii'
+
+    let deserConfig: AvroDeserializerConfig = {
+      ruleConfig: {
+        secret: 'mysecret'
+      }
+    }
+    let deser = new AvroDeserializer(client, SerdeType.VALUE, deserConfig)
+    fieldEncryptionExecutor.executor.client = dekClient
+    let obj2 = await deser.deserialize(topic, bytes)
+    expect(obj2.result["com.acme.Data"].code).toEqual(obj.result["com.acme.Data"].code);
+    expect(obj2.result["com.acme.Data"].secret).toEqual(obj.result["com.acme.Data"].secret);
+  })
   it('basic encryption with logical type', async () => {
     let conf: ClientConfig = {
       baseURLs: [baseURL],