Add flush method to AIOKafkaProducer

aiokafka/message_accumulator.py

@@ -181,13 +181,17 @@ def set_api_version(self, api_version):
         self._api_version = api_version
 
     @asyncio.coroutine
-    def close(self):
-        self._closed = True
+    def flush(self):
         # NOTE: we copy to avoid mutation during `yield from` below
         for batches in list(self._batches.values()):
             for batch in list(batches):
                 yield from batch.wait_deliver()
 
+    @asyncio.coroutine
+    def close(self):
+        self._closed = True
+        yield from self.flush()
+
     @asyncio.coroutine
     def add_message(self, tp, key, value, timeout, timestamp_ms=None):
         """ Add message to batch by topic-partition

aiokafka/producer.py

@@ -200,13 +200,17 @@ def start(self):
         self._message_accumulator.set_api_version(self.client.api_version)
         log.debug("Kafka producer started")
 
+    @asyncio.coroutine
+    def flush(self):
+        """Wait untill all batches are Delivered and futures resolved"""
+        yield from self._message_accumulator.flush()
+
     @asyncio.coroutine
     def stop(self):
         """Flush all pending data and close all connections to kafka cluster"""
         if self._closed:
             return
 
-        # Wait untill all batches are Delivered and futures resolved
         yield from self._message_accumulator.close()
 
         if self._sender_task: