trying to fix influxdata#531 by returning an actual generator when using query() in chunked mode

psy0rz · psy0rz · commit f50cdbc510e8 · 2017-11-26T01:56:07.000+01:00
diff --git a/influxdb/client.py b/influxdb/client.py
@@ -201,7 +201,7 @@ def switch_user(self, username, password):
         self._password = password
 
     def request(self, url, method='GET', params=None, data=None,
-                expected_response_code=200, headers=None):
+                expected_response_code=200, headers=None, chunked=False):
         """Make a HTTP request to the InfluxDB API.
 
         :param url: the path of the HTTP request, e.g. write, query, etc.
@@ -249,7 +249,8 @@ def request(self, url, method='GET', params=None, data=None,
                     headers=headers,
                     proxies=self._proxies,
                     verify=self._verify_ssl,
-                    timeout=self._timeout
+                    timeout=self._timeout,
+                    stream=chunked
                 )
                 break
             except (requests.exceptions.ConnectionError,
@@ -315,17 +316,17 @@ def write(self, data, params=None, expected_response_code=204,
 
     @staticmethod
     def _read_chunked_response(response, raise_errors=True):
-        result_set = {}
         for line in response.iter_lines():
             if isinstance(line, bytes):
                 line = line.decode('utf-8')
             data = json.loads(line)
+            result_set = {}
             for result in data.get('results', []):
                 for _key in result:
                     if isinstance(result[_key], list):
                         result_set.setdefault(
                             _key, []).extend(result[_key])
-        return ResultSet(result_set, raise_errors=raise_errors)
+            yield(ResultSet(result_set, raise_errors=raise_errors))
 
     def query(self,
               query,
@@ -391,7 +392,8 @@ def query(self,
             method='GET',
             params=params,
             data=None,
-            expected_response_code=expected_response_code
+            expected_response_code=expected_response_code,
+            chunked=chunked
         )
 
         if chunked:
