Refactor encryption executor (#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