feat: Add support for stream response (#30)

Author: bednar

influxdb_client/client/flux_csv_parser.py

@@ -2,8 +2,8 @@
 import codecs
 import csv as csv_parser
 
-from dateutil.parser import parse as timestamp_parser
 import ciso8601
+from urllib3 import HTTPResponse
 
 from influxdb_client.client.flux_table import FluxTable, FluxColumn, FluxRecord
 
@@ -20,21 +20,32 @@ class FluxCsvParserException(Exception):
 
 class FluxCsvParser(object):
 
-    def __init__(self) -> None:
+    def __init__(self, response: HTTPResponse, stream: bool) -> None:
+        self._response = response
+        self.tables = []
+        self._stream = stream
         pass
 
-    def parse_flux_response(self, response, cancellable, consumer):
+    def __enter__(self):
+        self._reader = csv_parser.reader(codecs.iterdecode(self._response, 'utf-8'))
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self._response.close()
+
+    def generator(self):
+        with self as parser:
+            yield from parser._parse_flux_response()
+
+    def _parse_flux_response(self):
         table_index = 0
         start_new_table = False
         table = None
         parsing_state_error = False
-        reader = csv_parser.reader(codecs.iterdecode(response, 'utf-8'))
 
-        for csv in reader:
+        for csv in self._reader:
             # debug
             # print("parsing: ", csv)
-            if (cancellable is not None) and cancellable.canceled:
-                return
 
             # Response has HTTP status ok, but response is error.
             if len(csv) < 1:
@@ -55,7 +66,7 @@ def parse_flux_response(self, response, cancellable, consumer):
             if "#datatype" == token:
                 start_new_table = True
                 table = FluxTable()
-                consumer.accept_table(index=table_index, cancellable=cancellable, flux_table=table)
+                self._insert_table(table, table_index)
                 table_index = table_index + 1
             elif table is None:
                 raise FluxCsvParserException("Unable to parse CSV response. FluxTable definition was not found.")
@@ -85,11 +96,16 @@ def parse_flux_response(self, response, cancellable, consumer):
                     flux_columns = table.columns
                     table = FluxTable()
                     table.columns.extend(flux_columns)
-                    consumer.accept_table(table_index, cancellable, table)
+                    self._insert_table(table, table_index)
                     table_index = table_index + 1
 
                 flux_record = self.parse_record(table_index - 1, table, csv)
-                consumer.accept_record(table_index - 1, cancellable, flux_record)
+
+                if not self._stream:
+                    self.tables[table_index - 1].records.append(flux_record)
+
+                yield flux_record
+
                 # debug
                 # print(flux_record)
 
@@ -163,14 +179,6 @@ def add_column_names_and_tags(table, csv):
             column.label = csv[i]
             i += 1
 
-
-class FluxResponseConsumerTable:
-
-    def __init__(self) -> None:
-        self.tables = []
-
-    def accept_table(self, index, cancellable, flux_table):
-        self.tables.insert(index, flux_table)
-
-    def accept_record(self, index, cancellable, flux_record):
-        self.tables[index].records.append(flux_record)
+    def _insert_table(self, table, table_index):
+        if not self._stream:
+            self.tables.insert(table_index, table)

influxdb_client/client/query_api.py

@@ -1,11 +1,11 @@
 import codecs
 import csv
-from typing import List
+from typing import List, Union, Iterable
 
 from influxdb_client import Dialect
 from influxdb_client import Query, QueryService
-from influxdb_client.client.flux_csv_parser import FluxCsvParser, FluxResponseConsumerTable
-from influxdb_client.client.flux_table import FluxTable
+from influxdb_client.client.flux_csv_parser import FluxCsvParser
+from influxdb_client.client.flux_table import FluxTable, FluxRecord
 
 
 class QueryApi(object):
@@ -35,6 +35,7 @@ def query_csv(self, query: str, org=None, dialect: Dialect = default_dialect):
             org = self._influxdb_client.org
         response = self._query_api.post_query(org=org, query=self._create_query(query, dialect), async_req=False,
                                               _preload_content=False)
+
         return csv.reader(codecs.iterdecode(response, 'utf-8'))
 
     def query_raw(self, query: str, org=None, dialect=default_dialect):
@@ -50,27 +51,32 @@ def query_raw(self, query: str, org=None, dialect=default_dialect):
             org = self._influxdb_client.org
         result = self._query_api.post_query(org=org, query=self._create_query(query, dialect), async_req=False,
                                             _preload_content=False)
+
         return result
-        # return codecs.iterdecode(result, 'utf-8')
 
-    def query(self, query: str, org=None, dialect=default_dialect) -> List['FluxTable']:
+    def query(self, query: str, org=None, dialect=default_dialect, stream=False) \
+            -> Union[List['FluxTable'], Iterable['FluxRecord']]:
         """
         Synchronously executes the Flux query and return result as a List['FluxTable']
 
         :param query: the Flux query
         :param org: organization name (optional if already specified in InfluxDBClient)
         :param dialect: csv dialect format
+        :param stream: csv dialect format
         :return:
         """
         if org is None:
             org = self._influxdb_client.org
         response = self._query_api.post_query(org=org, query=self._create_query(query, dialect), async_req=False,
                                               _preload_content=False, _return_http_data_only=False)
-        consumer = FluxResponseConsumerTable()
-        parser = FluxCsvParser()
 
-        parser.parse_flux_response(response=response, cancellable=None, consumer=consumer)
-        return consumer.tables
+        _parser = FluxCsvParser(response=response, stream=stream)
+        if stream:
+            return _parser.generator()
+
+        list(_parser.generator())
+
+        return _parser.tables
 
     # private helper for c
     @staticmethod

tests/test_QueryApiStream.py

@@ -0,0 +1,54 @@
+import itertools
+import time
+import types
+
+from influxdb_client import WritePrecision
+from influxdb_client.client.write_api import SYNCHRONOUS
+from tests.base_test import BaseTest
+
+
+class QueryStreamApi(BaseTest):
+
+    def setUp(self) -> None:
+        super().setUp()
+        self.write_client = self.client.write_api(write_options=SYNCHRONOUS)
+        self.bucket = self.create_test_bucket()
+
+    def tearDown(self) -> None:
+        self.write_client.__del__()
+        super().tearDown()
+
+    def test_block(self):
+        self._prepareData()
+
+        _result = self.query_api.query(
+            f'from(bucket:"{self.bucket.name}") |> range(start: 1970-01-01T00:00:00.000000001Z)', self.org)
+
+        self.assertEqual(len(_result), 1)
+        self.assertEqual(len(_result[0].records), 100)
+
+    def test_stream(self):
+        self._prepareData()
+
+        _result = self.query_api.query(
+            f'from(bucket:"{self.bucket.name}") |> range(start: 1970-01-01T00:00:00.000000001Z)', self.org, stream=True)
+
+        self.assertTrue(isinstance(_result, types.GeneratorType))
+        _result_list = list(_result)
+
+        self.assertEqual(len(_result_list), 100)
+
+    def test_stream_break(self):
+        self._prepareData()
+
+        _result = self.query_api.query(
+            f'from(bucket:"{self.bucket.name}") |> range(start: 1970-01-01T00:00:00.000000001Z)', self.org, stream=True)
+
+        _result_list = list(itertools.islice(_result, 10))
+
+        self.assertEqual(len(_result_list), 10)
+
+    def _prepareData(self):
+        _list = [f'h2o_feet,location=coyote_creek water_level={x} {x}' for x in range(1, 101)]
+        self.write_client.write(self.bucket.name, self.org, _list, write_precision=WritePrecision.S)
+        time.sleep(1)