[bug] fixing issue with trying to send empty batch to firehose (airbn…

…b#902) (airbnb#903) * fixing issue with trying to send empty batch to firehose (airbnb#902) * logging metric if record is too large to send to SQS
sayak119 · Feb 25, 2019 · b67c1f4 · b67c1f4
1 parent 81ced71
commit b67c1f4
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 2 deletions.
diff --git a/stream_alert/classifier/clients/firehose.py b/stream_alert/classifier/clients/firehose.py
@@ -99,7 +99,7 @@ def _record_batches(cls, records):
             # Check if the max size of the batch has been reached or if the current
             # record will exceed the max batch size and start a new batch
             if ((len(current_batch) == cls.MAX_BATCH_COUNT) or
-                    (current_batch_size + line_len > cls.MAX_BATCH_SIZE)):
+                    (current_batch_size + line_len > cls.MAX_BATCH_SIZE)) and current_batch:
                 yield current_batch[:]
                 current_batch_size = 0
                 del current_batch[:]

diff --git a/stream_alert/classifier/clients/sqs.py b/stream_alert/classifier/clients/sqs.py
@@ -76,6 +76,7 @@ def _segment_records(cls, records):
             size = len(record) + (1 if idx != record_count and batch else 0)
             if size + 2 > cls.MAX_SIZE:
                 LOGGER.error('Record is too large to send to SQS:\n%s', record)
+                MetricLogger.log_metric(FUNCTION_NAME, MetricLogger.SQS_FAILED_RECORDS, 1)
                 continue
 
             if idx == record_count or size + batch_size >= cls.MAX_SIZE:

diff --git a/tests/unit/streamalert/classifier/clients/test_firehose.py b/tests/unit/streamalert/classifier/clients/test_firehose.py
@@ -105,7 +105,7 @@ def test_record_batches_rec_too_large(self, failure_mock):
         ]
 
         result = list(FirehoseClient._record_batches(records))
-        assert_equal(result, [[]])
+        assert_equal(result, [])
         failure_mock.assert_called_with(1)
 
     def test_record_batches_max_batch_count(self):