add consistency parameter to write_points (#664)

RonRothman · xginn8 · commit dc54e5eb716d · 2019-04-08T12:13:13.000-05:00
* add consistency parameter to write_points [#643]
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -12,6 +12,7 @@ and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0.
 - query() now accepts a bind_params argument for parameter binding (#678 thx @clslgrnc)
 
 ### Changed
+- Add consistency param to InfluxDBClient.write_points (#643 thx @RonRothman)
 - Update test suite to add support for Python 3.7 and InfluxDB v1.6.4 and 1.7.4 (#692 thx @clslgrnc)
 - Update classifiers tuple to list in setup.py (#697 thx @Hanaasagi)
 - Update documentation for empty `delete_series` confusion
diff --git a/influxdb/client.py b/influxdb/client.py
@@ -458,7 +458,8 @@ def write_points(self,
                      retention_policy=None,
                      tags=None,
                      batch_size=None,
-                     protocol='json'
+                     protocol='json',
+                     consistency=None
                      ):
         """Write to multiple time series names.
 
@@ -486,6 +487,9 @@ def write_points(self,
         :type batch_size: int
         :param protocol: Protocol for writing data. Either 'line' or 'json'.
         :type protocol: str
+        :param consistency: Consistency for the points.
+            One of {'any','one','quorum','all'}.
+        :type consistency: str
         :returns: True, if the operation is successful
         :rtype: bool
 
@@ -498,14 +502,16 @@ def write_points(self,
                                    time_precision=time_precision,
                                    database=database,
                                    retention_policy=retention_policy,
-                                   tags=tags, protocol=protocol)
+                                   tags=tags, protocol=protocol,
+                                   consistency=consistency)
             return True
 
         return self._write_points(points=points,
                                   time_precision=time_precision,
                                   database=database,
                                   retention_policy=retention_policy,
-                                  tags=tags, protocol=protocol)
+                                  tags=tags, protocol=protocol,
+                                  consistency=consistency)
 
     def ping(self):
         """Check connectivity to InfluxDB.
@@ -531,12 +537,16 @@ def _write_points(self,
                       database,
                       retention_policy,
                       tags,
-                      protocol='json'):
+                      protocol='json',
+                      consistency=None):
         if time_precision not in ['n', 'u', 'ms', 's', 'm', 'h', None]:
             raise ValueError(
                 "Invalid time precision is given. "
                 "(use 'n', 'u', 'ms', 's', 'm' or 'h')")
 
+        if consistency not in ['any', 'one', 'quorum', 'all', None]:
+            raise ValueError('Invalid consistency: {}'.format(consistency))
+
         if protocol == 'json':
             data = {
                 'points': points
@@ -551,6 +561,9 @@ def _write_points(self,
             'db': database or self._database
         }
 
+        if consistency is not None:
+            params['consistency'] = consistency
+
         if time_precision is not None:
             params['precision'] = time_precision
 
diff --git a/influxdb/tests/client_test.py b/influxdb/tests/client_test.py
@@ -337,6 +337,23 @@ def test_write_points_with_precision(self):
                 m.last_request.body,
             )
 
+    def test_write_points_with_consistency(self):
+        """Test write points with consistency for TestInfluxDBClient object."""
+        with requests_mock.Mocker() as m:
+            m.register_uri(
+                requests_mock.POST,
+                'http://localhost:8086/write',
+                status_code=204
+            )
+
+            cli = InfluxDBClient(database='db')
+
+            cli.write_points(self.dummy_points, consistency='any')
+            self.assertEqual(
+                m.last_request.qs,
+                {'db': ['db'], 'consistency': ['any']}
+            )
+
     def test_write_points_with_precision_udp(self):
         """Test write points with precision for TestInfluxDBClient object."""
         s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
@@ -409,6 +426,15 @@ def test_write_points_bad_precision(self):
                 time_precision='g'
             )
 
+    def test_write_points_bad_consistency(self):
+        """Test write points w/bad consistency value."""
+        cli = InfluxDBClient()
+        with self.assertRaises(ValueError):
+            cli.write_points(
+                self.dummy_points,
+                consistency='boo'
+            )
+
     @raises(Exception)
     def test_write_points_with_precision_fails(self):
         """Test write points w/precision fail for TestInfluxDBClient object."""
