def put_record_batch(self, records: List[dict]) -> int:
"""
:param max_batch_size: max batch size
:param records: The records to put
:return: number of records inserted
"""
number_of_records = 0
firehose_records = self._convert_to_firehose_record(records)
chunks = split_to_chunks(firehose_records, self.batch_size)
batches: List[Batch] = self.create_batches_from_chunks(chunks)
while batches:
current_batch = batches.pop(0)
should_retry = current_batch.retry_count < self.max_retry_count
try:
response = self._client.put_record_batch(
DeliveryStreamName=self._stream_name,
Records=current_batch.records)["RequestResponses"]
failed_items = self.get_failed_items(current_batch, response)
self.update_failed_by_error_code(response)
success_items_len = len(
current_batch.records) - len(failed_items)
number_of_records += success_items_len
if any(failed_items) and should_retry:
batches.append(
self.create_next_batch(current_batch, failed_items))
except Exception as e:
get_logger().debug(
"Error while trying to send data to firehose",
exc_info=e,
)
self.failed_by_error_code[str(type(e).__name__)] += 1
if should_retry:
next_records = current_batch.records
batches.append(
self.create_next_batch(current_batch, next_records))
return number_of_records
def test_split_to_chunks_simple_flow():
lst = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
chunks = split_to_chunks(lst, 2)
assert chunks == [[1, 2], [3, 4], [5, 6], [7, 8], [9, 10]]
def test_split_to_chunks_empty_list():
lst = []
chunks = split_to_chunks(lst, 2)
assert chunks == []
def test_split_to_chunks_smaller_then_chunk_size():
lst = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
chunks = split_to_chunks(lst, 11)
assert chunks == [lst]
def test_split_to_chunks_not_equal_chunks():
lst = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
chunks = split_to_chunks(lst, 3)
assert chunks == [[1, 2, 3], [4, 5, 6], [7, 8, 9], [10]]
def test_split_to_chunks_min_value():
lst = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
chunks = split_to_chunks(lst, 0)
assert chunks == [[1], [2], [3], [4], [5], [6], [7], [8], [9], [10]]