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]]