Merge pull request #438 from prodo56/oauth-fix-for-kafka-source

Fixes issue 431 for kafka oauth source

Author: martindurant

ci/environment-py37.yml

@@ -12,18 +12,19 @@ dependencies:
   - toolz
   - zict
   - six
-  - librdkafka=1.2.2
+  - librdkafka=1.5.3
   - dask
   - distributed
   - pandas
-  - python-confluent-kafka=1.1.0
+  - python-confluent-kafka=1.5.0
   - numpydoc
   - sphinx
   - sphinx_rtd_theme
   - codecov
   - coverage
   - networkx
   - graphviz
+  - pytest-asyncio
   - python-graphviz
   - bokeh
   - ipython

ci/environment-py38.yml

@@ -12,11 +12,11 @@ dependencies:
   - toolz
   - zict
   - six
-  - librdkafka=1.2.2
+  - librdkafka=1.5.3
   - dask
   - distributed
   - pandas
-  - python-confluent-kafka=1.1.0
+  - python-confluent-kafka=1.5.0
   - numpydoc
   - sphinx
   - sphinx_rtd_theme
@@ -25,6 +25,7 @@ dependencies:
   - networkx
   - graphviz
   - python-graphviz
+  - pytest-asyncio
   - bokeh
   - ipython
   - ipykernel

streamz/sources.py

@@ -449,10 +449,14 @@ def start(self):
             self.stopped = False
             self.consumer = ck.Consumer(self.cpars)
             self.consumer.subscribe(self.topics)
-            weakref.finalize(self, lambda consumer=self.consumer: _close_consumer(consumer))
+            weakref.finalize(
+                self, lambda consumer=self.consumer: _close_consumer(consumer)
+            )
             tp = ck.TopicPartition(self.topics[0], 0, 0)
 
-            # blocks for consumer thread to come up
+            # blocks for consumer thread to come up and invoke poll to
+            # establish connection with broker to fetch oauth token for kafka
+            self.consumer.poll(timeout=1)
             self.consumer.get_watermark_offsets(tp)
             self.loop.add_callback(self.poll_kafka)
 
@@ -479,7 +483,8 @@ def __init__(self, topic, consumer_params, poll_interval='1s',
                  max_batch_size=10000, keys=False,
                  engine=None, **kwargs):
         self.consumer_params = consumer_params
-        # Override the auto-commit config to enforce custom streamz checkpointing
+        # Override the auto-commit config to enforce custom streamz
+        # checkpointing
         self.consumer_params['enable.auto.commit'] = 'false'
         if 'auto.offset.reset' not in self.consumer_params.keys():
             consumer_params['auto.offset.reset'] = 'latest'
@@ -587,7 +592,9 @@ def start(self):
             self.stopped = False
             tp = ck.TopicPartition(self.topic, 0, 0)
 
-            # blocks for consumer thread to come up
+            # blocks for consumer thread to come up and invoke poll to establish
+            # connection with broker to fetch oauth token for kafka
+            self.consumer.poll(timeout=1)
             self.consumer.get_watermark_offsets(tp)
             self.loop.add_callback(self.poll_kafka)