feat: Add support to write by byte array (#26) (#27)

Author: bednar

README.rst

@@ -156,7 +156,7 @@ The data should be passed as a `InfluxDB Line Protocol <https://docs.influxdata.
 The data could be written as
 """"""""""""""""""""""""""""
 
-1. ``string`` that is formatted as a InfluxDB's line protocol
+1. ``string`` or ``bytes`` that is formatted as a InfluxDB's line protocol
 2. `Data Point <https://github.com/influxdata/influxdb-client-python/blob/master/influxdb_client/client/write/point.py#L16>`__ structure
 3. Dictionary style mapping with keys: ``measurement``, ``tags``, ``fields`` and ``time``
 4. List of above items
@@ -205,12 +205,19 @@ The batching is configurable by ``write_options``\ :
                                                                 retry_interval=5_000))
 
    """
-   Write Line Protocol
+   Write Line Protocol formatted as string
    """
    _write_client.write("my-bucket", "my-org", "h2o_feet,location=coyote_creek water_level=1.0 1")
    _write_client.write("my-bucket", "my-org", ["h2o_feet,location=coyote_creek water_level=2.0 2",
                                                "h2o_feet,location=coyote_creek water_level=3.0 3"])
 
+   """
+   Write Line Protocol formatted as byte array
+   """
+   _write_client.write("my-bucket", "my-org", "h2o_feet,location=coyote_creek water_level=1.0 1".encode())
+   _write_client.write("my-bucket", "my-org", ["h2o_feet,location=coyote_creek water_level=2.0 2".encode(),
+                                               "h2o_feet,location=coyote_creek water_level=3.0 3".encode()])
+
    """
    Write Dictionary-style object
    """

influxdb_client/client/write_api.py

@@ -98,7 +98,7 @@ def __str__(self) -> str:
 
 
 def _body_reduce(batch_items):
-    return '\n'.join(map(lambda batch_item: batch_item.data, batch_items))
+    return b'\n'.join(map(lambda batch_item: batch_item.data, batch_items))
 
 
 def _create_batch(group: GroupedObservable):
@@ -144,7 +144,8 @@ def __init__(self, influxdb_client, write_options: WriteOptions = WriteOptions()
             self._disposable = None
 
     def write(self, bucket: str, org: str,
-              record: Union[str, List['str'], Point, List['Point'], dict, List['dict'], Observable],
+              record: Union[
+                  str, List['str'], Point, List['Point'], dict, List['dict'], bytes, List['bytes'], Observable],
               write_precision: WritePrecision = DEFAULT_WRITE_PRECISION) -> None:
         """
         Writes time-series data into influxdb.
@@ -159,27 +160,7 @@ def write(self, bucket: str, org: str,
         if self._write_options.write_type is WriteType.batching:
             return self._write_batching(bucket, org, record, write_precision)
 
-        final_string = ''
-
-        if isinstance(record, str):
-            final_string = record
-
-        if isinstance(record, Point):
-            final_string = record.to_line_protocol()
-
-        if isinstance(record, dict):
-            final_string = Point.from_dict(record, write_precision=write_precision).to_line_protocol()
-
-        if isinstance(record, list):
-            lines = []
-            for item in record:
-                if isinstance(item, str):
-                    lines.append(item)
-                if isinstance(item, Point):
-                    lines.append(item.to_line_protocol())
-                if isinstance(item, dict):
-                    lines.append(Point.from_dict(item, write_precision=write_precision).to_line_protocol())
-            final_string = '\n'.join(lines)
+        final_string = self._serialize(record, write_precision)
 
         _async_req = True if self._write_options.write_type == WriteType.asynchronous else False
 
@@ -203,13 +184,35 @@ def __del__(self):
             self._disposable = None
         pass
 
+    def _serialize(self, record, write_precision) -> bytes:
+        _result = b''
+        if isinstance(record, bytes):
+            _result = record
+
+        elif isinstance(record, str):
+            _result = record.encode("utf-8")
+
+        elif isinstance(record, Point):
+            _result = self._serialize(record.to_line_protocol(), write_precision=write_precision)
+
+        elif isinstance(record, dict):
+            _result = self._serialize(Point.from_dict(record, write_precision=write_precision),
+                                      write_precision=write_precision)
+        elif isinstance(record, list):
+            _result = b'\n'.join([self._serialize(item, write_precision=write_precision) for item in record])
+
+        return _result
+
     def _write_batching(self, bucket, org, data, precision=DEFAULT_WRITE_PRECISION):
         _key = _BatchItemKey(bucket, org, precision)
-        if isinstance(data, str):
+        if isinstance(data, bytes):
             self._subject.on_next(_BatchItem(key=_key, data=data))
 
+        elif isinstance(data, str):
+            self._write_batching(bucket, org, data.encode("utf-8"), precision)
+
         elif isinstance(data, Point):
-            self._subject.on_next(_BatchItem(key=_key, data=data.to_line_protocol()))
+            self._write_batching(bucket, org, data.to_line_protocol(), precision)
 
         elif isinstance(data, dict):
             self._write_batching(bucket, org, Point.from_dict(data, write_precision=precision), precision)
@@ -234,7 +237,7 @@ def _http(self, batch_item: _BatchItem):
         return _BatchResponse(data=batch_item)
 
     def _post_write(self, _async_req, bucket, org, body, precision):
-        return self._write_service.post_write(org=org, bucket=bucket, body=body.encode("utf-8"), precision=precision,
+        return self._write_service.post_write(org=org, bucket=bucket, body=body, precision=precision,
                                               async_req=_async_req, content_encoding="identity",
                                               content_type="text/plain; charset=utf-8")
 
@@ -272,4 +275,4 @@ def _on_error(ex):
 
     def _on_complete(self):
         self._disposable.dispose()
-        logger.info("the batching processor was dispose")
+        logger.info("the batching processor was disposed")

tests/test_WriteApi.py

@@ -157,6 +157,26 @@ def test_write_dictionary(self):
 
         self.delete_test_bucket(_bucket)
 
+    def test_write_bytes(self):
+        _bucket = self.create_test_bucket()
+        _bytes = "h2o_feet,location=coyote_creek level\\ water_level=1.0 1".encode("utf-8")
+
+        self.write_client.write(_bucket.name, self.org, _bytes)
+        self.write_client.flush()
+
+        result = self.query_api.query(
+            "from(bucket:\"" + _bucket.name + "\") |> range(start: 1970-01-01T00:00:00.000000001Z) |> last()", self.org)
+
+        self.assertEqual(len(result), 1)
+        self.assertEqual(result[0].records[0].get_measurement(), "h2o_feet")
+        self.assertEqual(result[0].records[0].get_value(), 1.0)
+        self.assertEqual(result[0].records[0].values.get("location"), "coyote_creek")
+        self.assertEqual(result[0].records[0].get_field(), "level water_level")
+        self.assertEqual(result[0].records[0].get_time(),
+                         datetime.datetime(1970, 1, 1, 0, 0, tzinfo=datetime.timezone.utc))
+
+        self.delete_test_bucket(_bucket)
+
 
 class AsynchronousWriteTest(BaseTest):
 
@@ -218,6 +238,44 @@ def test_write_dictionaries(self):
 
         self.delete_test_bucket(bucket)
 
+    def test_write_bytes(self):
+        bucket = self.create_test_bucket()
+
+        _bytes1 = "h2o_feet,location=coyote_creek level\\ water_level=1.0 1".encode("utf-8")
+        _bytes2 = "h2o_feet,location=coyote_creek level\\ water_level=2.0 2".encode("utf-8")
+
+        _bytes_list = [_bytes1, _bytes2]
+
+        self.write_client.write(bucket.name, self.org, _bytes_list, write_precision=WritePrecision.S)
+        time.sleep(1)
+
+        query = 'from(bucket:"' + bucket.name + '") |> range(start: 1970-01-01T00:00:00.000000001Z)'
+        print(query)
+
+        flux_result = self.client.query_api().query(query)
+
+        self.assertEqual(1, len(flux_result))
+
+        records = flux_result[0].records
+
+        self.assertEqual(2, len(records))
+
+        self.assertEqual("h2o_feet", records[0].get_measurement())
+        self.assertEqual(1, records[0].get_value())
+        self.assertEqual("level water_level", records[0].get_field())
+        self.assertEqual("coyote_creek", records[0].values.get('location'))
+        self.assertEqual(records[0].get_time(),
+                         datetime.datetime(1970, 1, 1, 0, 0, 1, tzinfo=datetime.timezone.utc))
+
+        self.assertEqual("h2o_feet", records[1].get_measurement())
+        self.assertEqual(2, records[1].get_value())
+        self.assertEqual("level water_level", records[1].get_field())
+        self.assertEqual("coyote_creek", records[1].values.get('location'))
+        self.assertEqual(records[1].get_time(),
+                         datetime.datetime(1970, 1, 1, 0, 0, 2, tzinfo=datetime.timezone.utc))
+
+        self.delete_test_bucket(bucket)
+
 
 if __name__ == '__main__':
     unittest.main()

tests/test_WriteApiBatching.py

@@ -277,15 +277,22 @@ def test_record_types(self):
                   "time": 14, "fields": {"level water_level": 14.0}}
         _dict2 = {"measurement": "h2o_feet", "tags": {"location": "coyote_creek"},
                   "time": 15, "fields": {"level water_level": 15.0}}
-        _dict3 = {"measurement": "h2o_feet", "tags": {"location": "coyote_creek"},
-                  "time": 16, "fields": {"level water_level": 16.0}}
-        self._write_client.write("my-bucket", "my-org", [_dict1, _dict2, _dict3])
+        self._write_client.write("my-bucket", "my-org", [_dict1, _dict2])
+
+        # Bytes item
+        _bytes = "h2o_feet,location=coyote_creek level\\ water_level=16.0 16".encode("utf-8")
+        self._write_client.write("my-bucket", "my-org", _bytes)
+
+        # Bytes list
+        _bytes1 = "h2o_feet,location=coyote_creek level\\ water_level=17.0 17".encode("utf-8")
+        _bytes2 = "h2o_feet,location=coyote_creek level\\ water_level=18.0 18".encode("utf-8")
+        self._write_client.write("my-bucket", "my-org", [_bytes1, _bytes2])
 
         time.sleep(1)
 
         _requests = httpretty.httpretty.latest_requests
 
-        self.assertEqual(8, len(_requests))
+        self.assertEqual(9, len(_requests))
 
         self.assertEqual("h2o_feet,location=coyote_creek level\\ water_level=1.0 1\n"
                          "h2o_feet,location=coyote_creek level\\ water_level=2.0 2", _requests[0].parsed_body)
@@ -303,6 +310,8 @@ def test_record_types(self):
                          "h2o_feet,location=coyote_creek level\\ water_level=14.0 14", _requests[6].parsed_body)
         self.assertEqual("h2o_feet,location=coyote_creek level\\ water_level=15.0 15\n"
                          "h2o_feet,location=coyote_creek level\\ water_level=16.0 16", _requests[7].parsed_body)
+        self.assertEqual("h2o_feet,location=coyote_creek level\\ water_level=17.0 17\n"
+                         "h2o_feet,location=coyote_creek level\\ water_level=18.0 18", _requests[8].parsed_body)
 
         pass