Add Variant, Decimal, and Timestamp CEL functions

CHANGELOG.md

@@ -15,6 +15,7 @@ v2.16.0 is a feature release with the following features, fixes and enhancements
   `dlq.auto.flush=true` or give the serde its own `RuleRegistry` (closable on
   shutdown) for durability.
 - Add support for inline validation rules (#2326)
+- Add Variant, Decimal, and Timestamp CEL functions (#2332)
 
 ### Fixes
 

pyproject.toml

@@ -26,6 +26,11 @@ Homepage = "https://github.com/confluentinc/confluent-kafka-python"
 
 [tool.mypy]
 ignore_missing_imports = true
+# The generated protobuf modules are not statically analysable: _builder injects the message
+# classes into globals() at import time. Excluding them keeps them out of the build set, which
+# is what lets follow_imports = "skip" below apply - it is ignored for files mypy was asked to
+# check directly.
+exclude = '_pb2\.py$'
 
 [[tool.mypy.overrides]]
 module = [
@@ -35,12 +40,18 @@ module = [
 ]
 disable_error_code = ["assignment", "no-redef"]
 
+# Generated protobuf modules build their message classes at import time, so the classes are
+# not statically visible - neither in the generated file itself nor to anything annotating
+# against them. follow_imports = "skip" makes the module Any, which covers both.
 [[tool.mypy.overrides]]
 module = [
     "confluent_kafka.schema_registry.confluent.meta_pb2",
+    "confluent_kafka.schema_registry.confluent.type.decimal_pb2",
+    "confluent_kafka.schema_registry.confluent.type.variant_pb2",
     "confluent_kafka.schema_registry.confluent.types.decimal_pb2",
 ]
 ignore_errors = true
+follow_imports = "skip"
 
 [tool.black]
 line-length = 120

src/confluent_kafka/admin/__init__.py

@@ -33,6 +33,7 @@
 from .._model import ElectionType as _ElectionType
 from .._model import TopicCollection as _TopicCollection
 from ..cimpl import KafkaException  # noqa: F401
+from ..cimpl import _AdminClientImpl  # noqa: F401
 from ..cimpl import (  # noqa: F401
     CONFIG_SOURCE_DEFAULT_CONFIG,
     CONFIG_SOURCE_DYNAMIC_BROKER_CONFIG,
@@ -53,7 +54,6 @@
     NewTopic,
 )
 from ..cimpl import TopicPartition as _TopicPartition
-from ..cimpl import _AdminClientImpl
 from ._acl import AclOperation  # noqa: F401
 from ._acl import AclBinding, AclBindingFilter, AclPermissionType  # noqa: F401
 from ._cluster import DescribeClusterResult  # noqa: F401

src/confluent_kafka/schema_registry/common/avro.py

@@ -10,6 +10,7 @@
 from fastavro import repository, validate
 from fastavro.schema import load_schema
 
+from confluent_kafka.schema_registry.confluent.type.variant_utils import Variant
 from confluent_kafka.schema_registry.serde import (
     VALIDATION_RULES_PROP,
     FieldTransform,
@@ -25,6 +26,33 @@
 
 from .schema_registry_client import RuleKind, Schema
 
+
+# The Avro `variant` logical type: a record {metadata: bytes, value: bytes} carrying a
+# Spark/Parquet Variant. A field with this logical type decodes to / encodes from a Variant,
+# so serde consumers and CEL rules see a first-class Variant rather than raw bytes. fastavro
+# keys logical handlers by "<avro-type>-<logicalType>" = "record-variant"; this is the Python
+# counterpart of Java's io.confluent.avro.type.VariantConversion.
+def _variant_from_avro(data, writer_schema, reader_schema=None):  # noqa: ARG001
+    return Variant(bytes(data["value"]), bytes(data["metadata"]))
+
+
+def _variant_to_avro(data, schema):  # noqa: ARG001
+    if isinstance(data, Variant):
+        # standalone_value_bytes, not .value: a navigated sub-variant's own value starts at
+        # its position, and .value is the whole shared buffer.
+        return {"metadata": data.metadata, "value": data.standalone_value_bytes()}
+    return data
+
+
+try:
+    from fastavro import read as _fastavro_read
+    from fastavro import write as _fastavro_write
+
+    _fastavro_read.LOGICAL_READERS["record-variant"] = _variant_from_avro
+    _fastavro_write.LOGICAL_WRITERS["record-variant"] = _variant_to_avro
+except Exception:  # pragma: no cover - guards against a fastavro API shape change
+    pass
+
 __all__ = [
     'AvroMessage',
     'AvroSchema',
@@ -113,10 +141,14 @@
     return load_schema("$root", repo=repo)
 
 
 def transform(
     ctx: RuleContext, schema: AvroSchema, message: AvroMessage, field_transform: FieldTransform
 ) -> AvroMessage:
-    if message is None or schema is None:
+    # Only the schema being absent stops the walk. A `None` *value* is the null branch of a
+    # `["null", T]` union and has to reach the rule: the reference binds it as CEL null so a
+    # rule can guard with `value == null`, and returning early here skipped the rule entirely -
+    # indistinguishable, to the caller, from a rule that ran and passed.
+    if schema is None:
         return message
     field_ctx = ctx.current_field()
     if field_ctx is not None:
@@ -126,7 +158,12 @@
         if subschema is None:
             return message
         submessage = transform(ctx, subschema, submessage, field_transform)
-        if isinstance(message, tuple) and len(message) == 2:
+        # The branch a transformed value belongs to follows from the value, not from the branch
+        # it arrived on - the reference keeps no branch at all and resolves it from the datum.
+        # Keep the notation while its branch still accepts the result, so two same-shaped
+        # records are never swapped; drop it otherwise and let fastavro resolve, since
+        # ("null", x) would be written as null with x silently dropped.
+        if isinstance(message, tuple) and len(message) == 2 and _branch_accepts(subschema, submessage):
             return (message[0], submessage)
         return submessage
     elif isinstance(schema, dict):
@@ -142,6 +179,11 @@
                 return message
             return {key: transform(ctx, schema["values"], value, field_transform) for key, value in message.items()}
         elif schema_type == 'record':
+            # A null record has no fields to walk. Guarded before the isinstance check below
+            # so a legitimate null does not log an "incompatible message type" warning; the
+            # reference guards the record case, and only the record case, the same way.
+            if message is None:
+                return message
             if not isinstance(message, dict):
                 log.warning("Incompatible message type for record schema")
                 return message
@@ -364,6 +406,15 @@
     return '.' not in name and not subschema.get("namespace") and branch_name.rsplit('.', 1)[-1] == name
 
 
+def _branch_accepts(subschema: AvroSchema, message: AvroMessage) -> bool:
+    """Whether ``subschema`` can hold ``message``, by the same test ``_resolve_union`` uses."""
+    try:
+        validate(message, _collapse_schema(deepcopy(subschema)))
+        return True
+    except:  # noqa: E722
+        return False
+
+
 def _resolve_union(schema: AvroSchema, message: AvroMessage) -> Tuple[Optional[AvroSchema], AvroMessage]:
     is_wrapped_union = isinstance(message, tuple) and len(message) == 2
     is_typed_union = isinstance(message, dict) and '-type' in message