Added chunked query responses implementation and test. Added chunked parameter to client query function.

Panos · Panos · commit 1df262cd6320 · 2017-02-16T14:55:32.000Z
diff --git a/influxdb/client.py b/influxdb/client.py
@@ -293,13 +293,23 @@ def write(self, data, params=None, expected_response_code=204,
         )
         return True
 
+    def _read_chunked_response(self, response, raise_errors=True):
+        for line in response.iter_lines():
+            # import ipdb; ipdb.set_trace()
+            if isinstance(line, bytes):
+                line = line.decode('utf-8')
+            data = json.loads(line)
+            for result in data.get('results', []):
+                yield ResultSet(result, raise_errors=raise_errors)
+
     def query(self,
               query,
               params=None,
               epoch=None,
               expected_response_code=200,
               database=None,
-              raise_errors=True):
+              raise_errors=True,
+              chunked=False):
         """Send a query to InfluxDB.
 
         :param query: the actual query string
@@ -339,6 +349,10 @@ def query(self,
             expected_response_code=expected_response_code
         )
 
+        if chunked or 'chunked' in params:
+            params['chunked'] = 'true'
+            return self._read_chunked_response(response)
+
         data = response.json()
 
         results = [
diff --git a/influxdb/tests/client_test.py b/influxdb/tests/client_test.py
@@ -32,7 +32,7 @@
 import unittest
 
 from influxdb import InfluxDBClient
-
+from influxdb.resultset import ResultSet
 
 def _build_response_object(status_code=200, content=""):
     resp = requests.Response()
@@ -792,6 +792,37 @@ def test_invalid_port_fails(self):
         with self.assertRaises(ValueError):
             InfluxDBClient('host', '80/redir', 'username', 'password')
 
+    def test_chunked_response(self):
+        example_response = u'{"results":[{"statement_id":0,"series": ' \
+          '[{"name":"cpu","columns":["fieldKey","fieldType"],"values":' \
+          '[["value","integer"]]}],"partial":true}]}\n{"results":' \
+          '[{"statement_id":0,"series":[{"name":"iops","columns":' \
+          '["fieldKey","fieldType"],"values":[["value","integer"]]}],' \
+          '"partial":true}]}\n{"results":[{"statement_id":0,"series":' \
+          '[{"name":"load","columns":["fieldKey","fieldType"],"values":' \
+          '[["value","integer"]]}],"partial":true}]}\n{"results":' \
+          '[{"statement_id":0,"series":[{"name":"memory","columns":' \
+          '["fieldKey","fieldType"],"values":[["value","integer"]]}]}]}\n'
+
+        with requests_mock.Mocker() as m:
+            m.register_uri(
+                requests_mock.GET,
+                "http://localhost:8086/query",
+                text=example_response
+            )
+            response = list(self.cli.query('show series limit 4 offset 0', chunked=True))
+            self.assertTrue(len(response) == 4)
+            self.assertEqual(response[0].raw, ResultSet(
+                {"statement_id":0,
+                 "series": [{"name":"cpu","columns":["fieldKey","fieldType"],
+                             "values": [["value","integer"]]}],"partial":True}
+                ).raw)
+            self.assertEqual(response[3].raw, ResultSet(
+                {"statement_id":0,
+                 "series":[{"name":"memory","columns":
+                            ["fieldKey","fieldType"],
+                            "values":[["value","integer"]]}]}
+                ).raw)
 
 class FakeClient(InfluxDBClient):
 
