Support publishing v2j batches in the event publisher sidecar. (#590)

* Support publishing v2j event batches

* lint

* add test

Co-authored-by: Alexander Early <alexander.early@reddit.com>

Author: aearly

baseplate/sidecars/event_publisher.py

@@ -86,10 +86,21 @@ def reset(self) -> None:
         self._size = len(self._header) + len(self._end)
 
 
+class V2JBatch(V2Batch):
+    # Send a batch as a plain JSON array.  Useful when your events are not
+    # Thrift JSON
+    _header = "["
+    _end = b"]"
+
+    def serialize(self) -> SerializedBatch:
+        serialized = self._header.encode() + b",".join(self._items) + self._end
+        return SerializedBatch(item_count=len(self._items), serialized=serialized)
+
+
 class BatchPublisher:
     def __init__(self, metrics_client: metrics.Client, cfg: Any):
         self.metrics = metrics_client
-        self.url = f"https://{cfg.collector.hostname}/v{cfg.collector.version:d}"
+        self.url = f"https://{cfg.collector.hostname}/v{cfg.collector.version}"
         self.key_name = cfg.key.name
         self.key_secret = cfg.key.secret
         self.session = requests.Session()
@@ -149,7 +160,7 @@ def publish(self, payload: SerializedBatch) -> None:
         raise MaxRetriesError("could not sent batch")
 
 
-SERIALIZER_BY_VERSION = {2: V2Batch}
+SERIALIZER_BY_VERSION = {"2": V2Batch, "2j": V2JBatch}
 
 
 def publish_events() -> None:
@@ -181,7 +192,7 @@ def publish_events() -> None:
         {
             "collector": {
                 "hostname": config.String,
-                "version": config.Optional(config.Integer, default=1),
+                "version": config.Optional(config.String, default="2"),
             },
             "key": {"name": config.String, "secret": config.Base64},
             "max_queue_size": config.Optional(config.Integer, MAX_QUEUE_SIZE),

tests/unit/lib/events/publisher_tests.py

@@ -56,6 +56,25 @@ def test_v2(self):
         self.assertEqual(result.item_count, 0)
 
 
+class V2JBatchTests(unittest.TestCase):
+    def test_v2j(self):
+        batch = event_publisher.V2JBatch(max_size=50)
+        batch.add(None)
+        batch.add(b'{"source":1}')
+        batch.add(b'{"source":2}')
+
+        result = batch.serialize()
+        self.assertEqual(result.item_count, 2)
+        self.assertEqual(result.serialized, b'[{"source":1},{"source":2}]')
+
+        with self.assertRaises(event_publisher.BatchFull):
+            batch.add(b"x" * 100)
+
+        batch.reset()
+        result = batch.serialize()
+        self.assertEqual(result.item_count, 0)
+
+
 class PublisherTests(unittest.TestCase):
     @mock.patch("requests.Session", autospec=True)
     def setUp(self, Session):