Fix chunked query to return chunk resultsets

hrbonz · hrbonz · commit a6bc05d56d47 · 2019-09-08T21:28:04.000+08:00
When querying large data sets, it's vital to get a chunked responses to manage memory usage. Wrapping the query response in a generator and streaming the request provides the desired result. It also fixes `InfluxDBClient.query()` behavior for chunked queries that is currently not working according to [specs](https://github.com/influxdata/influxdb-python/blob/master/influxdb/client.py#L410) close #585 close #531
diff --git a/influxdb/client.py b/influxdb/client.py
@@ -231,7 +231,7 @@ def switch_user(self, username, password):
         self._username = username
         self._password = password
 
-    def request(self, url, method='GET', params=None, data=None,
+    def request(self, url, method='GET', params=None, data=None, stream=False,
                 expected_response_code=200, headers=None):
         """Make a HTTP request to the InfluxDB API.
 
@@ -243,6 +243,8 @@ def request(self, url, method='GET', params=None, data=None,
         :type params: dict
         :param data: the data of the request, defaults to None
         :type data: str
+        :param stream: True if a query uses chunked responses
+        :type stream: bool
         :param expected_response_code: the expected response code of
             the request, defaults to 200
         :type expected_response_code: int
@@ -277,6 +279,7 @@ def request(self, url, method='GET', params=None, data=None,
                     auth=(self._username, self._password),
                     params=params,
                     data=data,
+                    stream=stream,
                     headers=headers,
                     proxies=self._proxies,
                     verify=self._verify_ssl,
@@ -346,17 +349,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,
@@ -447,6 +450,7 @@ def query(self,
             method=method,
             params=params,
             data=None,
+            stream=chunked,
             expected_response_code=expected_response_code
         )
 
diff --git a/influxdb/tests/client_test.py b/influxdb/tests/client_test.py
@@ -1239,8 +1239,8 @@ def test_chunked_response(self):
             )
             response = self.cli.query('show series limit 4 offset 0',
                                       chunked=True, chunk_size=4)
-            self.assertTrue(len(response) == 4)
-            self.assertEqual(response.__repr__(), ResultSet(
+            self.assertTrue(len(list(response)) == 4)
+            self.assertEqual(list(response.__repr__()), ResultSet(
                 {'series': [{'values': [['value', 'integer']],
                              'name': 'cpu',
                              'columns': ['fieldKey', 'fieldType']},
