Remove UDP Precision Restrictions (#557)

mdhausman · xginn8 · commit ed276e2187c5 · 2018-02-13T17:48:38.000-05:00
- address issue #554 - UDP writes can now convert timestamps to the desired precision - add time_precision argument to `InfluxDBClient.send_packet`, defaults to None - add tests for udp precision - remove old udp precision failure tests
diff --git a/influxdb/client.py b/influxdb/client.py
@@ -497,11 +497,6 @@ def _write_points(self,
                 "Invalid time precision is given. "
                 "(use 'n', 'u', 'ms', 's', 'm' or 'h')")
 
-        if self._use_udp and time_precision and time_precision != 's':
-            raise ValueError(
-                "InfluxDB only supports seconds precision for udp writes"
-            )
-
         if protocol == 'json':
             data = {
                 'points': points
@@ -523,7 +518,9 @@ def _write_points(self,
             params['rp'] = retention_policy
 
         if self._use_udp:
-            self.send_packet(data, protocol=protocol)
+            self.send_packet(
+                data, protocol=protocol, time_precision=time_precision
+            )
         else:
             self.write(
                 data=data,
@@ -864,17 +861,19 @@ def get_list_privileges(self, username):
         text = "SHOW GRANTS FOR {0}".format(quote_ident(username))
         return list(self.query(text).get_points())
 
-    def send_packet(self, packet, protocol='json'):
+    def send_packet(self, packet, protocol='json', time_precision=None):
         """Send an UDP packet.
 
         :param packet: the packet to be sent
         :type packet: (if protocol is 'json') dict
                       (if protocol is 'line') list of line protocol strings
         :param protocol: protocol of input data, either 'json' or 'line'
         :type protocol: str
+        :param time_precision: Either 's', 'm', 'ms' or 'u', defaults to None
+        :type time_precision: str
         """
         if protocol == 'json':
-            data = make_lines(packet).encode('utf-8')
+            data = make_lines(packet, time_precision).encode('utf-8')
         elif protocol == 'line':
             data = ('\n'.join(packet) + '\n').encode('utf-8')
         self.udp_socket.sendto(data, (self._host, self._udp_port))
diff --git a/influxdb/tests/client_test.py b/influxdb/tests/client_test.py
@@ -259,22 +259,6 @@ def test_write_points_udp(self):
             received_data.decode()
         )
 
-    def test_write_bad_precision_udp(self):
-        """Test write bad precision in UDP for TestInfluxDBClient object."""
-        cli = InfluxDBClient(
-            'localhost', 8086, 'root', 'root',
-            'test', use_udp=True, udp_port=4444
-        )
-
-        with self.assertRaisesRegexp(
-                Exception,
-                "InfluxDB only supports seconds precision for udp writes"
-        ):
-            cli.write_points(
-                self.dummy_points,
-                time_precision='ms'
-            )
-
     @raises(Exception)
     def test_write_points_fails(self):
         """Test write points fail for TestInfluxDBClient object."""
@@ -335,6 +319,65 @@ def test_write_points_with_precision(self):
                 m.last_request.body,
             )
 
+    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)
+        port = random.randint(4000, 8000)
+        s.bind(('0.0.0.0', port))
+
+        cli = InfluxDBClient(
+            'localhost', 8086, 'root', 'root',
+            'test', use_udp=True, udp_port=port
+        )
+
+        cli.write_points(self.dummy_points, time_precision='n')
+        received_data, addr = s.recvfrom(1024)
+        self.assertEqual(
+            b'cpu_load_short,host=server01,region=us-west '
+            b'value=0.64 1257894000123456000\n',
+            received_data,
+        )
+
+        cli.write_points(self.dummy_points, time_precision='u')
+        received_data, addr = s.recvfrom(1024)
+        self.assertEqual(
+            b'cpu_load_short,host=server01,region=us-west '
+            b'value=0.64 1257894000123456\n',
+            received_data,
+        )
+
+        cli.write_points(self.dummy_points, time_precision='ms')
+        received_data, addr = s.recvfrom(1024)
+        self.assertEqual(
+            b'cpu_load_short,host=server01,region=us-west '
+            b'value=0.64 1257894000123\n',
+            received_data,
+        )
+
+        cli.write_points(self.dummy_points, time_precision='s')
+        received_data, addr = s.recvfrom(1024)
+        self.assertEqual(
+            b"cpu_load_short,host=server01,region=us-west "
+            b"value=0.64 1257894000\n",
+            received_data,
+        )
+
+        cli.write_points(self.dummy_points, time_precision='m')
+        received_data, addr = s.recvfrom(1024)
+        self.assertEqual(
+            b'cpu_load_short,host=server01,region=us-west '
+            b'value=0.64 20964900\n',
+            received_data,
+        )
+
+        cli.write_points(self.dummy_points, time_precision='h')
+        received_data, addr = s.recvfrom(1024)
+        self.assertEqual(
+            b'cpu_load_short,host=server01,region=us-west '
+            b'value=0.64 349415\n',
+            received_data,
+        )
+
     def test_write_points_bad_precision(self):
         """Test write points w/bad precision TestInfluxDBClient object."""
         cli = InfluxDBClient()
