Ensure multiple rules are supported (#1870)

confluentinc · Dec 17, 2024 · 70872b9 · 70872b9
1 parent 057640c
commit 70872b9
Show file tree

Hide file tree

Showing 2 changed files with 81 additions and 3 deletions.
diff --git a/src/confluent_kafka/schema_registry/serde.py b/src/confluent_kafka/schema_registry/serde.py
@@ -335,10 +335,8 @@ def _execute_rules(
                 if rule.kind == RuleKind.CONDITION:
                     if not result:
                         raise RuleConditionError(rule)
-                    break
                 elif rule.kind == RuleKind.TRANSFORM:
                     message = result
-                    break
                 self._run_action(
                     ctx, rule_mode, rule,
                     self._get_on_failure(rule) if message is None else self._get_on_success(rule),
@@ -377,7 +375,7 @@ def _run_action(
         action_name = self._get_rule_action_name(rule, rule_mode, action)
         if action_name is None:
             action_name = default_action
-        rule_action = self._get_rule_action(self._rule_registry, action_name)
+        rule_action = self._get_rule_action(ctx, action_name)
         if rule_action is None:
             log.error("Could not find rule action of type %s", action_name)
             raise RuleError(f"Could not find rule action of type {action_name}")

diff --git a/tests/schema_registry/test_avro_serdes.py b/tests/schema_registry/test_avro_serdes.py
@@ -929,6 +929,86 @@ def test_avro_encryption():
     assert obj == obj2
 
 
+def test_avro_encryption_cel():
+    executor = FieldEncryptionExecutor.register_with_clock(FakeClock())
+
+    conf = {'url': _BASE_URL}
+    client = SchemaRegistryClient.new_client(conf)
+    ser_conf = {'auto.register.schemas': False, 'use.latest.version': True}
+    rule_conf = {'secret': 'mysecret'}
+    schema = {
+        'type': 'record',
+        'name': 'test',
+        'fields': [
+            {'name': 'intField', 'type': 'int'},
+            {'name': 'doubleField', 'type': 'double'},
+            {'name': 'stringField', 'type': 'string', 'confluent:tags': ['PII']},
+            {'name': 'booleanField', 'type': 'boolean'},
+            {'name': 'bytesField', 'type': 'bytes', 'confluent:tags': ['PII']},
+        ]
+    }
+
+    rule1 = Rule(
+        "test-cel",
+        "",
+        RuleKind.TRANSFORM,
+        RuleMode.WRITE,
+        "CEL_FIELD",
+        None,
+        None,
+        "name == 'stringField' ; value + '-suffix'",
+        None,
+        None,
+        False
+    )
+    rule2 = Rule(
+        "test-encrypt",
+        "",
+        RuleKind.TRANSFORM,
+        RuleMode.WRITEREAD,
+        "ENCRYPT",
+        ["PII"],
+        RuleParams({
+            "encrypt.kek.name": "kek1",
+            "encrypt.kms.type": "local-kms",
+            "encrypt.kms.key.id": "mykey"
+        }),
+        None,
+        None,
+        "ERROR,NONE",
+        False
+    )
+    client.register_schema(_SUBJECT, Schema(
+        json.dumps(schema),
+        "AVRO",
+        [],
+        None,
+        RuleSet(None, [rule1, rule2])
+    ))
+
+    obj = {
+        'intField': 123,
+        'doubleField': 45.67,
+        'stringField': 'hi',
+        'booleanField': True,
+        'bytesField': b'foobar',
+    }
+    ser = AvroSerializer(client, schema_str=None, conf=ser_conf, rule_conf=rule_conf)
+    dek_client = executor.client
+    ser_ctx = SerializationContext(_TOPIC, MessageField.VALUE)
+    obj_bytes = ser(obj, ser_ctx)
+
+    # reset encrypted fields
+    assert obj['stringField'] != 'hi-suffix'
+    obj['stringField'] = 'hi-suffix'
+    obj['bytesField'] = b'foobar'
+
+    deser = AvroDeserializer(client, rule_conf=rule_conf)
+    executor.client = dek_client
+    obj2 = deser(obj_bytes, ser_ctx)
+    assert obj == obj2
+
+
 def test_avro_encryption_dek_rotation():
     executor = FieldEncryptionExecutor.register_with_clock(FakeClock())