Split SerializerError into KeySerializerError and ValueSerializerError

kwilcox · kwilcox · commit b1bf1b0b613d · 2016-12-12T15:53:57.000-05:00
diff --git a/confluent_kafka/avro/__init__.py b/confluent_kafka/avro/__init__.py
@@ -53,7 +53,9 @@ def __str__(self):
 
 
 from confluent_kafka.avro.cached_schema_registry_client import CachedSchemaRegistryClient
-from confluent_kafka.avro.serializer import SerializerError
+from confluent_kafka.avro.serializer import (SerializerError,
+                                             KeySerializerError,
+                                             ValueSerializerError)
 from confluent_kafka.avro.serializer.message_serializer import MessageSerializer
 
 
@@ -99,17 +101,19 @@ def produce(self, **kwargs):
             raise ClientError("Topic name not specified.")
         value = kwargs.pop('value', None)
         key = kwargs.pop('key', None)
+
         if value:
             if value_schema:
                 value = self._serializer.encode_record_with_schema(topic, value_schema, value)
             else:
-                raise SerializerError("Avro schema required for value")
+                raise ValueSerializerError("Avro schema required for values")
 
         if key:
             if key_schema:
                 key = self._serializer.encode_record_with_schema(topic, key_schema, key, True)
             else:
-                raise SerializerError("Avro schema required for key")
+                raise KeySerializerError("Avro schema required for key")
+
 
         super(AvroProducer, self).produce(topic, value, key, **kwargs)
 
diff --git a/confluent_kafka/avro/serializer/__init__.py b/confluent_kafka/avro/serializer/__init__.py
@@ -1,11 +1,40 @@
+#!/usr/bin/env python
+#
+# Copyright 2016 Confluent Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+#
+
+
 class SerializerError(Exception):
     """Generic error from serializer package"""
 
     def __init__(self, message):
         self.message = message
 
         def __repr__(self):
-            return 'SerializerError(error={error})'.format(error=self.message)
+            return '{klass}(error={error})'.format(
+                klass=self.__class__.__name__,
+                error=self.message
+            )
 
         def __str__(self):
             return self.message
+
+
+class KeySerializerError(SerializerError):
+    pass
+
+
+class ValueSerializerError(SerializerError):
+    pass
diff --git a/confluent_kafka/avro/serializer/message_serializer.py b/confluent_kafka/avro/serializer/message_serializer.py
@@ -29,7 +29,9 @@
 import avro.io
 
 from confluent_kafka.avro import ClientError
-from . import SerializerError
+from confluent_kafka.avro.serializer import (SerializerError,
+                                             KeySerializerError,
+                                             ValueSerializerError)
 
 log = logging.getLogger(__name__)
 
@@ -84,46 +86,52 @@ def encode_record_with_schema(self, topic, schema, record, is_key=False):
         @:param topic : Topic name
         @:param schema : Avro Schema
         @:param record : A dictionary object
+        @:param is_key : If the record is a key
         @:returns : Encoded record with schema ID as bytes
         """
         if not isinstance(record, dict):
             raise SerializerError("record must be a dictionary")
+        serialize_err = KeySerializerError if is_key else ValueSerializerError
+
         subject_suffix = ('-key' if is_key else '-value')
         # get the latest schema for the subject
         subject = topic + subject_suffix
         # register it
         schema_id = self.registry_client.register(subject, schema)
         if not schema_id:
             message = "Unable to retrieve schema id for subject %s" % (subject)
-            raise SerializerError(message)
+            raise serialize_err(message)
 
         # cache writer
         self.id_to_writers[schema_id] = avro.io.DatumWriter(schema)
 
-        return self.encode_record_with_schema_id(schema_id, record)
+        return self.encode_record_with_schema_id(schema_id, record, is_key=is_key)
 
-    def encode_record_with_schema_id(self, schema_id, record):
+    def encode_record_with_schema_id(self, schema_id, record, is_key=False):
         """
         Encode a record with a given schema id.  The record must
         be a python dictionary.
         @:param: schema_id : integer ID
         @:param: record : A dictionary object
+        @:param is_key : If the record is a key
         @:returns: decoder function
         """
         if not isinstance(record, dict):
             raise SerializerError("record must be a dictionary")
+        serialize_err = KeySerializerError if is_key else ValueSerializerError
+
         # use slow avro
         if schema_id not in self.id_to_writers:
             # get the writer + schema
 
             try:
                 schema = self.registry_client.get_by_id(schema_id)
                 if not schema:
-                    raise SerializerError("Schema does not exist")
+                    raise serialize_err("Schema does not exist")
                 self.id_to_writers[schema_id] = avro.io.DatumWriter(schema)
             except ClientError as e:
                 exc_type, exc_value, exc_traceback = sys.exc_info()
-                raise SerializerError("Error fetching schema from registry:" + repr(
+                raise serialize_err( + repr(
                     traceback.format_exception(exc_type, exc_value, exc_traceback)))
 
         # get the writer
