PEP8 changes

parsedmarc/kafkaclient.py

@@ -2,18 +2,26 @@
 # -*- coding: utf-8 -*-
 
 from kafka import KafkaProducer
-from collections import OrderedDict
+from kafka.errors import NoBrokersAvailable, UnknownTopicOrPartitionError
 import json
 
+
 class KafkaError(RuntimeError):
     """Raised when a Kafka error occurs"""
 
-class KafkaClient(object):
-  def __init__(self, kafka_hosts):
-    self.producer = KafkaProducer(value_serializer=lambda v: json.dumps(v).encode('utf-8'),
-                                    bootstrap_servers=kafka_hosts)
 
-  def save_aggregate_reports_to_kafka(self, aggregate_reports, aggregate_topic):
+class KafkaClient(object):
+    def __init__(self, kafka_hosts):
+        try:
+            def serializer(v): lambda v: json.dumps(v).encode('utf-8')
+            self.producer = KafkaProducer(
+                          value_serializer=serializer,
+                          bootstrap_servers=kafka_hosts)
+        except NoBrokersAvailable:
+            raise KafkaError("No Kafka brokers availabe")
+
+    def save_aggregate_reports_to_kafka(self, aggregate_reports,
+                                        aggregate_topic):
         """
         Saves aggregate DMARC reports to Kafka
 
@@ -28,20 +36,18 @@ class KafkaClient(object):
         if len(aggregate_reports) < 1:
             return
 
-        for record in aggregate_reports['records']:
-            buffer = OrderedDict([('xml_schema', aggregate_reports['xml_schema']),
-                                ('report_metadata',aggregate_reports['report_metadata']),
-                                ('records', record)])
-            self.producer.send(aggregate_topic, buffer)
+        try:
+            self.producer.send(aggregate_topic, aggregate_reports)
+        except UnknownTopicOrPartitionError:
+                raise KafkaError("Unknown topic or partition on broker")
         self.producer.flush()
 
-
-  def save_forensic_reports_to_kafka(self, forensic_reports, forensic_topic):
+    def save_forensic_reports_to_kafka(self, forensic_reports, forensic_topic):
             """
             Saves forensic DMARC reports to Kafka
 
             Args:
-                forensic_reports (list):  A list of forensic report dictionaries
+                forensic_reports (list):  A list of forensic report dicts
                 to save to kafka
 
             """
@@ -51,6 +57,8 @@ class KafkaClient(object):
             if len(forensic_reports) < 1:
                 return
 
-            for report in forensic_reports:
-                self.producer.send(forensic_topic, json.dumps(report))
-                self.producer.flush()
+            try:
+                self.producer.send(forensic_topic, forensic_reports)
+            except UnknownTopicOrPartitionError:
+                raise KafkaError("Unknown topic or partition on broker")
+            self.producer.flush()