Refactor encryption executor (confluentinc#2011)

Author: rayokota

src/confluent_kafka/schema_registry/_async/avro.py

@@ -14,7 +14,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
+import io
 from json import loads
 from typing import Dict, Union, Optional, Callable
 
@@ -29,6 +29,8 @@
                                              AsyncSchemaRegistryClient,
                                              prefix_schema_id_serializer,
                                              dual_schema_id_deserializer)
+from confluent_kafka.schema_registry.common.schema_registry_client import \
+    RulePhase
 from confluent_kafka.serialization import (SerializationError,
                                            SerializationContext)
 from confluent_kafka.schema_registry.rule_registry import RuleRegistry
@@ -348,8 +350,14 @@ def field_transformer(rule_ctx, field_transform, msg): return (  # noqa: E731
         with _ContextStringIO() as fo:
             # write the record to the rest of the buffer
             schemaless_writer(fo, parsed_schema, value)
+            buffer = fo.getvalue()
+
+            if latest_schema is not None:
+                buffer = self._execute_rules_with_phase(
+                    ctx, subject, RulePhase.ENCODING, RuleMode.WRITE,
+                    None, latest_schema.schema, buffer, None, None)
 
-            return self._schema_id_serializer(fo.getvalue(), ctx, self._schema_id)
+            return self._schema_id_serializer(buffer, ctx, self._schema_id)
 
     async def _get_parsed_schema(self, schema: Schema) -> AvroSchema:
         parsed_schema = self._parsed_schemas.get_parsed_schema(schema)
@@ -557,6 +565,12 @@ async def __deserialize(
             if subject is not None:
                 latest_schema = await self._get_reader_schema(subject)
 
+        payload = self._execute_rules_with_phase(
+            ctx, subject, RulePhase.ENCODING, RuleMode.READ,
+            None, writer_schema_raw, payload, None, None)
+        if isinstance(payload, bytes):
+            payload = io.BytesIO(payload)
+
         if latest_schema is not None:
             migrations = await self._get_migrations(subject, writer_schema_raw, latest_schema, None)
             reader_schema_raw = latest_schema.schema

src/confluent_kafka/schema_registry/_async/json_schema.py

@@ -14,7 +14,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
+import io
 import json
 from typing import Union, Optional, Tuple, Callable
 
@@ -33,6 +33,8 @@
 from confluent_kafka.schema_registry.common.json_schema import (
     DEFAULT_SPEC, JsonSchema, _retrieve_via_httpx, transform, _ContextStringIO, JSON_TYPE
 )
+from confluent_kafka.schema_registry.common.schema_registry_client import \
+    RulePhase
 from confluent_kafka.schema_registry.rule_registry import RuleRegistry
 from confluent_kafka.schema_registry.serde import AsyncBaseSerializer, AsyncBaseDeserializer, \
     ParsedSchemaCache, SchemaId
@@ -374,8 +376,14 @@ def field_transformer(rule_ctx, field_transform, msg): return (  # noqa: E731
             if isinstance(encoded_value, str):
                 encoded_value = encoded_value.encode("utf8")
             fo.write(encoded_value)
+            buffer = fo.getvalue()
+
+            if latest_schema is not None:
+                buffer = self._execute_rules_with_phase(
+                    ctx, subject, RulePhase.ENCODING, RuleMode.WRITE,
+                    None, latest_schema.schema, buffer, None, None)
 
-            return self._schema_id_serializer(fo.getvalue(), ctx, self._schema_id)
+            return self._schema_id_serializer(buffer, ctx, self._schema_id)
 
     async def _get_parsed_schema(self, schema: Schema) -> Tuple[Optional[JsonSchema], Optional[Registry]]:
         if schema is None:
@@ -552,9 +560,9 @@ async def __init_impl(
     __init__ = __init_impl
 
     def __call__(self, data: bytes, ctx: Optional[SerializationContext] = None) -> Optional[bytes]:
-        return self.__serialize(data, ctx)
+        return self.__deserialize(data, ctx)
 
-    async def __serialize(self, data: bytes, ctx: Optional[SerializationContext] = None) -> Optional[bytes]:
+    async def __deserialize(self, data: bytes, ctx: Optional[SerializationContext] = None) -> Optional[bytes]:
         """
         Deserialize a JSON encoded record with Confluent Schema Registry framing to
         a dict, or object instance according to from_dict if from_dict is specified.
@@ -583,9 +591,6 @@ async def __serialize(self, data: bytes, ctx: Optional[SerializationContext] = N
         schema_id = SchemaId(JSON_TYPE)
         payload = self._schema_id_deserializer(data, ctx, schema_id)
 
-        # JSON documents are self-describing; no need to query schema
-        obj_dict = self._json_decode(payload.read())
-
         if self._registry is not None:
             writer_schema_raw = await self._get_writer_schema(schema_id, subject)
             writer_schema, writer_ref_registry = await self._get_parsed_schema(writer_schema_raw)
@@ -597,6 +602,15 @@ async def __serialize(self, data: bytes, ctx: Optional[SerializationContext] = N
             writer_schema_raw = None
             writer_schema, writer_ref_registry = None, None
 
+        payload = self._execute_rules_with_phase(
+            ctx, subject, RulePhase.ENCODING, RuleMode.READ,
+            None, writer_schema_raw, payload, None, None)
+        if isinstance(payload, bytes):
+            payload = io.BytesIO(payload)
+
+        # JSON documents are self-describing; no need to query schema
+        obj_dict = self._json_decode(payload.read())
+
         if latest_schema is not None:
             migrations = await self._get_migrations(subject, writer_schema_raw, latest_schema, None)
             reader_schema_raw = latest_schema.schema

src/confluent_kafka/schema_registry/_async/protobuf.py

@@ -27,6 +27,8 @@
 from confluent_kafka.schema_registry import (reference_subject_name_strategy,
                                              topic_subject_name_strategy,
                                              prefix_schema_id_serializer, dual_schema_id_deserializer)
+from confluent_kafka.schema_registry.common.schema_registry_client import \
+    RulePhase
 from confluent_kafka.schema_registry.schema_registry_client import AsyncSchemaRegistryClient
 from confluent_kafka.schema_registry.common.protobuf import _bytes, _create_index_array, \
     _init_pool, _is_builtin, _schema_to_str, _str_to_proto, transform, _ContextStringIO, PROTOBUF_TYPE
@@ -426,7 +428,14 @@ def field_transformer(rule_ctx, field_transform, msg): return (  # noqa: E731
         with _ContextStringIO() as fo:
             fo.write(message.SerializeToString())
             self._schema_id.message_indexes = self._index_array
-            return self._schema_id_serializer(fo.getvalue(), ctx, self._schema_id)
+            buffer = fo.getvalue()
+
+            if latest_schema is not None:
+                buffer = self._execute_rules_with_phase(
+                    ctx, subject, RulePhase.ENCODING, RuleMode.WRITE,
+                    None, latest_schema.schema, buffer, None, None)
+
+            return self._schema_id_serializer(buffer, ctx, self._schema_id)
 
     async def _get_parsed_schema(self, schema: Schema) -> Tuple[descriptor_pb2.FileDescriptorProto, DescriptorPool]:
         result = self._parsed_schemas.get_parsed_schema(schema)
@@ -549,9 +558,9 @@ async def __init_impl(
     __init__ = __init_impl
 
     def __call__(self, data: bytes, ctx: Optional[SerializationContext] = None) -> Optional[bytes]:
-        return self.__serialize(data, ctx)
+        return self.__deserialize(data, ctx)
 
-    async def __serialize(self, data: bytes, ctx: Optional[SerializationContext] = None) -> Optional[bytes]:
+    async def __deserialize(self, data: bytes, ctx: Optional[SerializationContext] = None) -> Optional[bytes]:
         """
         Deserialize a serialized protobuf message with Confluent Schema Registry
         framing.
@@ -596,6 +605,12 @@ async def __serialize(self, data: bytes, ctx: Optional[SerializationContext] = N
             writer_schema_raw = None
             writer_schema = None
 
+        payload = self._execute_rules_with_phase(
+            ctx, subject, RulePhase.ENCODING, RuleMode.READ,
+            None, writer_schema_raw, payload, None, None)
+        if isinstance(payload, bytes):
+            payload = io.BytesIO(payload)
+
         if latest_schema is not None:
             migrations = await self._get_migrations(subject, writer_schema_raw, latest_schema, None)
             reader_schema_raw = latest_schema.schema

src/confluent_kafka/schema_registry/_async/serde.py

@@ -20,6 +20,8 @@
 from typing import List, Optional, Set, Dict, Any
 
 from confluent_kafka.schema_registry import RegisteredSchema
+from confluent_kafka.schema_registry.common.schema_registry_client import \
+    RulePhase
 from confluent_kafka.schema_registry.common.serde import ErrorAction, \
     FieldTransformer, Migration, NoneAction, RuleAction, \
     RuleConditionError, RuleContext, RuleError, SchemaId
@@ -59,6 +61,17 @@ def _execute_rules(
         source: Optional[Schema], target: Optional[Schema],
         message: Any, inline_tags: Optional[Dict[str, Set[str]]],
         field_transformer: Optional[FieldTransformer]
+    ) -> Any:
+        return self._execute_rules_with_phase(
+            ser_ctx, subject, RulePhase.DOMAIN, rule_mode,
+            source, target, message, inline_tags, field_transformer)
+
+    def _execute_rules_with_phase(
+        self, ser_ctx: SerializationContext, subject: str,
+        rule_phase: RulePhase, rule_mode: RuleMode,
+        source: Optional[Schema], target: Optional[Schema],
+        message: Any, inline_tags: Optional[Dict[str, Set[str]]],
+        field_transformer: Optional[FieldTransformer]
     ) -> Any:
         if message is None or target is None:
             return message
@@ -73,7 +86,10 @@ def _execute_rules(
                 rules.reverse()
         else:
             if target is not None and target.rule_set is not None:
-                rules = target.rule_set.domain_rules
+                if rule_phase == RulePhase.ENCODING:
+                    rules = target.rule_set.encoding_rules
+                else:
+                    rules = target.rule_set.domain_rules
                 if rule_mode == RuleMode.READ:
                     # Execute read rules in reverse order for symmetry
                     rules = rules[:] if rules else []
@@ -197,19 +213,25 @@ async def _get_writer_schema(
         else:
             raise SerializationError("Schema ID or GUID is not set")
 
-    def _has_rules(self, rule_set: RuleSet, mode: RuleMode) -> bool:
+    def _has_rules(self, rule_set: RuleSet, phase: RulePhase, mode: RuleMode) -> bool:
         if rule_set is None:
             return False
+        if phase == RulePhase.MIGRATION:
+            rules = rule_set.migration_rules
+        elif phase == RulePhase.DOMAIN:
+            rules = rule_set.domain_rules
+        elif phase == RulePhase.ENCODING:
+            rules = rule_set.encoding_rules
         if mode in (RuleMode.UPGRADE, RuleMode.DOWNGRADE):
             return any(rule.mode == mode or rule.mode == RuleMode.UPDOWN
-                       for rule in rule_set.migration_rules or [])
+                       for rule in rules or [])
         elif mode == RuleMode.UPDOWN:
-            return any(rule.mode == mode for rule in rule_set.migration_rules or [])
+            return any(rule.mode == mode for rule in rules or [])
         elif mode in (RuleMode.WRITE, RuleMode.READ):
             return any(rule.mode == mode or rule.mode == RuleMode.WRITEREAD
-                       for rule in rule_set.domain_rules or [])
+                       for rule in rules or [])
         elif mode == RuleMode.WRITEREAD:
-            return any(rule.mode == mode for rule in rule_set.migration_rules or [])
+            return any(rule.mode == mode for rule in rules or [])
         return False
 
     async def _get_migrations(
@@ -235,7 +257,8 @@ async def _get_migrations(
             if i == 0:
                 previous = version
                 continue
-            if version.schema.rule_set is not None and self._has_rules(version.schema.rule_set, migration_mode):
+            if version.schema.rule_set is not None and self._has_rules(
+                  version.schema.rule_set, RulePhase.MIGRATION, migration_mode):
                 if migration_mode == RuleMode.UPGRADE:
                     migration = Migration(migration_mode, previous, version)
                 else:
@@ -265,7 +288,8 @@ def _execute_migrations(
         migrations: List[Migration], message: Any
     ) -> Any:
         for migration in migrations:
-            message = self._execute_rules(ser_ctx, subject, migration.rule_mode,
-                                          migration.source.schema, migration.target.schema,
-                                          message, None, None)
+            message = self._execute_rules_with_phase(
+                ser_ctx, subject, RulePhase.MIGRATION, migration.rule_mode,
+                migration.source.schema, migration.target.schema,
+                message, None, None)
         return message

src/confluent_kafka/schema_registry/_sync/avro.py

@@ -14,7 +14,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
+import io
 from json import loads
 from typing import Dict, Union, Optional, Callable
 
@@ -29,6 +29,8 @@
                                              SchemaRegistryClient,
                                              prefix_schema_id_serializer,
                                              dual_schema_id_deserializer)
+from confluent_kafka.schema_registry.common.schema_registry_client import \
+    RulePhase
 from confluent_kafka.serialization import (SerializationError,
                                            SerializationContext)
 from confluent_kafka.schema_registry.rule_registry import RuleRegistry
@@ -348,8 +350,14 @@ def field_transformer(rule_ctx, field_transform, msg): return (  # noqa: E731
         with _ContextStringIO() as fo:
             # write the record to the rest of the buffer
             schemaless_writer(fo, parsed_schema, value)
+            buffer = fo.getvalue()
+
+            if latest_schema is not None:
+                buffer = self._execute_rules_with_phase(
+                    ctx, subject, RulePhase.ENCODING, RuleMode.WRITE,
+                    None, latest_schema.schema, buffer, None, None)
 
-            return self._schema_id_serializer(fo.getvalue(), ctx, self._schema_id)
+            return self._schema_id_serializer(buffer, ctx, self._schema_id)
 
     def _get_parsed_schema(self, schema: Schema) -> AvroSchema:
         parsed_schema = self._parsed_schemas.get_parsed_schema(schema)
@@ -557,6 +565,12 @@ def __deserialize(
             if subject is not None:
                 latest_schema = self._get_reader_schema(subject)
 
+        payload = self._execute_rules_with_phase(
+            ctx, subject, RulePhase.ENCODING, RuleMode.READ,
+            None, writer_schema_raw, payload, None, None)
+        if isinstance(payload, bytes):
+            payload = io.BytesIO(payload)
+
         if latest_schema is not None:
             migrations = self._get_migrations(subject, writer_schema_raw, latest_schema, None)
             reader_schema_raw = latest_schema.schema

src/confluent_kafka/schema_registry/_sync/json_schema.py

@@ -14,7 +14,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
+import io
 import json
 from typing import Union, Optional, Tuple, Callable
 
@@ -33,6 +33,8 @@
 from confluent_kafka.schema_registry.common.json_schema import (
     DEFAULT_SPEC, JsonSchema, _retrieve_via_httpx, transform, _ContextStringIO, JSON_TYPE
 )
+from confluent_kafka.schema_registry.common.schema_registry_client import \
+    RulePhase
 from confluent_kafka.schema_registry.rule_registry import RuleRegistry
 from confluent_kafka.schema_registry.serde import BaseSerializer, BaseDeserializer, \
     ParsedSchemaCache, SchemaId
@@ -374,8 +376,14 @@ def field_transformer(rule_ctx, field_transform, msg): return (  # noqa: E731
             if isinstance(encoded_value, str):
                 encoded_value = encoded_value.encode("utf8")
             fo.write(encoded_value)
+            buffer = fo.getvalue()
+
+            if latest_schema is not None:
+                buffer = self._execute_rules_with_phase(
+                    ctx, subject, RulePhase.ENCODING, RuleMode.WRITE,
+                    None, latest_schema.schema, buffer, None, None)
 
-            return self._schema_id_serializer(fo.getvalue(), ctx, self._schema_id)
+            return self._schema_id_serializer(buffer, ctx, self._schema_id)
 
     def _get_parsed_schema(self, schema: Schema) -> Tuple[Optional[JsonSchema], Optional[Registry]]:
         if schema is None:
@@ -552,9 +560,9 @@ def __init_impl(
     __init__ = __init_impl
 
     def __call__(self, data: bytes, ctx: Optional[SerializationContext] = None) -> Optional[bytes]:
-        return self.__serialize(data, ctx)
+        return self.__deserialize(data, ctx)
 
-    def __serialize(self, data: bytes, ctx: Optional[SerializationContext] = None) -> Optional[bytes]:
+    def __deserialize(self, data: bytes, ctx: Optional[SerializationContext] = None) -> Optional[bytes]:
         """
         Deserialize a JSON encoded record with Confluent Schema Registry framing to
         a dict, or object instance according to from_dict if from_dict is specified.
@@ -583,9 +591,6 @@ def __serialize(self, data: bytes, ctx: Optional[SerializationContext] = None) -
         schema_id = SchemaId(JSON_TYPE)
         payload = self._schema_id_deserializer(data, ctx, schema_id)
 
-        # JSON documents are self-describing; no need to query schema
-        obj_dict = self._json_decode(payload.read())
-
         if self._registry is not None:
             writer_schema_raw = self._get_writer_schema(schema_id, subject)
             writer_schema, writer_ref_registry = self._get_parsed_schema(writer_schema_raw)
@@ -597,6 +602,15 @@ def __serialize(self, data: bytes, ctx: Optional[SerializationContext] = None) -
             writer_schema_raw = None
             writer_schema, writer_ref_registry = None, None
 
+        payload = self._execute_rules_with_phase(
+            ctx, subject, RulePhase.ENCODING, RuleMode.READ,
+            None, writer_schema_raw, payload, None, None)
+        if isinstance(payload, bytes):
+            payload = io.BytesIO(payload)
+
+        # JSON documents are self-describing; no need to query schema
+        obj_dict = self._json_decode(payload.read())
+
         if latest_schema is not None:
             migrations = self._get_migrations(subject, writer_schema_raw, latest_schema, None)
             reader_schema_raw = latest_schema.schema