Call stream API from separate thread when processing queues (#6)

* Call stream API from separate thread when processing queues

* Fix flake8 warning

Author: rflprr

target_datadotworld/__init__.py

@@ -21,6 +21,6 @@
 
 import singer
 
-__version__ = '1.0.0b1'
+__version__ = '1.0.0b2'
 
 logger = copy(singer.get_logger())  # copy needed in order to set level

target_datadotworld/api_client.py

@@ -16,8 +16,9 @@
 #
 # This product includes software developed at
 # data.world, Inc.(http://data.world/).
-
+import functools
 import gzip
+from concurrent.futures import ThreadPoolExecutor
 from time import sleep
 
 import backoff
@@ -46,6 +47,7 @@ def __init__(self, api_token, **kwargs):
         self._api_url = kwargs.get('api_url', 'https://api.data.world/v0')
         self._conn_timeout = kwargs.get('connect_timeout', 3.05)
         self._read_timeout = kwargs.get('read_timeout', 600)
+        self._max_threads = kwargs.get('max_threads', 10)
 
         self._session = requests.Session()
         default_headers = {
@@ -58,6 +60,11 @@ def __init__(self, api_token, **kwargs):
         self._session.mount(self._api_url,
                             BackoffAdapter(GzipAdapter(HTTPAdapter())))
 
+        # Create a limited thread pool.
+        self._executor = ThreadPoolExecutor(
+            max_workers=self._max_threads
+        )
+
     def connection_check(self):
         """Verify network connectivity
 
@@ -100,7 +107,7 @@ def append_stream(self, owner, dataset, stream, records):
                 raise convert_requests_exception(e)
 
     async def append_stream_chunked(
-            self, owner, dataset, stream, queue, chunk_size):
+            self, owner, dataset, stream, queue, chunk_size, loop):
         """Asynchronously append records to a stream in a data.world dataset
 
         :param owner: User or organization ID of the owner of the dataset
@@ -122,20 +129,33 @@ async def append_stream_chunked(
 
             delayed_exception = None
             # noinspection PyTypeChecker
+            pending_task = None
             async for chunk in to_chunks(queue, chunk_size):
                 if delayed_exception is None:
                     try:
                         logger.info('Uploading {} records in batch #{} '
                                     'from {} stream '.format(
                                         len(chunk), counter.value, stream))
-                        # TODO Invoke append_stream in a separate thread
-                        self.append_stream(owner, dataset, stream, chunk)
+
+                        if pending_task is not None:
+                            # Force chunks to be appended sequentially
+                            await pending_task
+
+                        # Call API on separate thread
+                        pending_task = loop.run_in_executor(
+                            self._executor,
+                            functools.partial(self.append_stream,
+                                              owner, dataset, stream, chunk)
+                        )
                         counter.increment()
                     except Exception as e:
                         delayed_exception = e
                 else:
                     pass  # Must exhaust queue
 
+            if pending_task is not None:
+                await pending_task
+
             if delayed_exception is not None:
                 raise delayed_exception
 

target_datadotworld/target.py

@@ -156,7 +156,7 @@ async def process_lines(self, lines, loop=None):
                                 self.config['dataset_id'],
                                 to_stream_id(msg.stream),
                                 queue,
-                                self._batch_size), loop=loop)
+                                self._batch_size, loop=loop), loop=loop)
 
                     # Add record to queue
                     await queues[msg.stream].put(msg.record)

tests/target_datadotworld/test_api_client.py

@@ -97,7 +97,7 @@ def verify_body_and_count(req):
 
             consumer = asyncio.ensure_future(client.append_stream_chunked(
                 'owner', 'dataset', 'stream', queue,
-                chunk_size=chunk_size), loop=event_loop)
+                chunk_size=chunk_size, loop=event_loop), loop=event_loop)
             await queue.join()
             await consumer
 
@@ -120,7 +120,7 @@ async def test_append_stream_chunked_error(
         with pytest.raises(dwex.ApiError):
             consumer = asyncio.ensure_future(client.append_stream_chunked(
                 'owner', 'dataset', 'stream', queue,
-                chunk_size=chunk_size), loop=event_loop)
+                chunk_size=chunk_size, loop=event_loop), loop=event_loop)
             await queue.join()
             await consumer
 

tests/target_datadotworld/test_target.py

@@ -44,7 +44,7 @@ def create_dataset(self, owner, dataset, **kwargs):
             return {}
 
         async def append_stream_chunked(
-                self, owner, dataset, stream, queue, chunk_size):
+                self, owner, dataset, stream, queue, chunk_size, loop):
             while True:
                 item = await queue.get()
                 queue.task_done()