Merge pull request influxdata#38 from ReAzem/master

toddboom · toddboom · commit 9006644f3f2d · 2014-06-06T16:49:14.000-04:00
Added UDP support
diff --git a/influxdb/client.py b/influxdb/client.py
@@ -3,7 +3,7 @@
 python client for influxdb
 """
 import json
-
+import socket
 import requests
 session = requests.Session()
 
@@ -12,14 +12,23 @@ class InfluxDBClientError(Exception):
     def __init__(self, message, code):
       self.message = message
       self.code = code
-      
+
 class InfluxDBClient(object):
     """
     InfluxDB Client
     """
 
-    def __init__(self, host='localhost', port=8086, username='root',
-                 password='root', database=None, ssl=False, verify_ssl=False, timeout=0):
+    def __init__(self,
+                 host='localhost',
+                 port=8086,
+                 username='root',
+                 password='root',
+                 database=None,
+                 ssl=False,
+                 verify_ssl=False,
+                 timeout=0,
+                 use_udp=False,
+                 udp_port=4444):
         """
         Initialize client
         """
@@ -32,6 +41,11 @@ def __init__(self, host='localhost', port=8086, username='root',
 
         self._verify_ssl = verify_ssl
 
+        self.use_udp = use_udp
+        self.udp_port = udp_port
+        if use_udp:
+            self.udp_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+
         self._scheme = "http"
 
         if ssl is True:
@@ -102,7 +116,7 @@ def request(self, url, method='GET', params=None, data=None,
 
         if response.status_code == status_code:
             return response
-            
+
         else:
             error = InfluxDBClientError("{0}: {1}".format(response.status_code, response.content), response.status_code)
             raise error
@@ -167,13 +181,16 @@ def write_points_with_precision(self, data, time_precision='s'):
             'time_precision': time_precision
         }
 
-        self.request(
-            url=url,
-            method='POST',
-            params=params,
-            data=data,
-            status_code=200
-            )
+        if self.use_udp:
+            self.send_packet(data)
+        else:
+            self.request(
+                url=url,
+                method='POST',
+                params=params,
+                data=data,
+                status_code=200
+                )
 
         return True
 
@@ -266,7 +283,7 @@ def query(self, query, time_precision='s', chunked=False):
         try:
             res = json.loads(response.content)
         except TypeError:
-            # must decode in python 3 
+            # must decode in python 3
             res = json.loads(response.content.decode('utf8'))
 
         return res
@@ -618,3 +635,8 @@ def update_permission(self, username, json_body):
         See also: src/api/http/api.go:l57
         """
         raise NotImplementedError()
+
+    def send_packet(self, packet):
+        data = json.dumps(packet)
+        byte = data.encode('utf-8')
+        self.udp_socket.sendto(byte, (self._host, self.udp_port))
diff --git a/tests/influxdb/client_test.py b/tests/influxdb/client_test.py
@@ -4,6 +4,7 @@
 """
 import json
 import requests
+import socket
 from nose.tools import raises
 from mock import patch
 
@@ -36,7 +37,7 @@ def request(*args, **kwargs):
                 assert isinstance(data, str)
 
                 # Data must be a JSON string
-                assert c == json.loads(data)
+                assert c == json.loads(data, strict=True)
 
                 c = data
 
@@ -91,6 +92,27 @@ def test_write_points(self):
             cli = InfluxDBClient('host', 8086, 'username', 'password', 'db')
             assert cli.write_points(data) is True
 
+    def test_write_points_udp(self):
+        data = [
+            {
+                "points": [
+                    ["1", 1, 1.0],
+                    ["2", 2, 2.0]
+                ],
+                "name": "foo",
+                "columns": ["column_one", "column_two", "column_three"]
+            }
+        ]
+
+        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+        s.bind(('0.0.0.0', 4444))
+
+        cli = InfluxDBClient('localhost', 8086, 'root', 'root', 'test', use_udp=True, udp_port=4444)
+        cli.write_points(data)
+
+        received_data, addr = s.recvfrom(1024)
+        assert data == json.loads(received_data.decode(), strict=True)
+
     @raises(Exception)
     def test_write_points_fails(self):
         with _mocked_session('post', 500) as mocked:
@@ -273,4 +295,4 @@ def test_delete_database_user(self):
     @raises(NotImplementedError)
     def test_update_permission(self):
         cli = InfluxDBClient('host', 8086, 'username', 'password', 'db')
-        cli.update_permission('admin', [])
+        cli.update_permission('admin', [])
