Add from_DSN to InfluxDBClusterClient

cannium · cannium · commit 4555e6a560e9 · 2015-04-22T17:37:47.000+08:00
diff --git a/influxdb/client.py b/influxdb/client.py
@@ -18,7 +18,7 @@
 except NameError:
     xrange = range
 
-if version_info.major == 3:
+if version_info[0] == 3:
     from urllib.parse import urlparse
 else:
     from urlparse import urlparse
@@ -469,6 +469,9 @@ class InfluxDBClusterClient(object):
     :param hosts: A list of hosts, where a host should be in format
                   (address, port)
                   e.g. [('127.0.0.1', 8086), ('127.0.0.1', 9096)]
+    :param shuffle: If true, queries will hit servers evenly(randomly)
+    :param client_base_class: In order to support different clients,
+                              default to InfluxDBClient
     """
 
     def __init__(self,
@@ -482,11 +485,11 @@ def __init__(self,
                  use_udp=False,
                  udp_port=4444,
                  shuffle=True,
-                 client_base_class=InfluxDBClient,  # For simpler test code
+                 client_base_class=InfluxDBClient,
                  ):
         self.clients = []
         self.bad_clients = []   # Corresponding server has failures in history
-        self.shuffle = shuffle  # if true, queries will hit servers evenly
+        self.shuffle = shuffle
         for h in hosts:
             self.clients.append(client_base_class(host=h[0], port=h[1],
                                                   username=username,
@@ -505,6 +508,40 @@ def __init__(self,
                 continue
             setattr(self, method, self._make_func(orig_func))
 
+    @staticmethod
+    def from_DSN(dsn, client_base_class=InfluxDBClient,
+                 shuffle=True, **kwargs):
+        """
+        Same as InfluxDBClient.from_DSN, and supports multiple servers.
+
+        Example DSN:
+            influxdb://usr:pwd@host1:8086,usr:pwd@host2:8086/db_name
+            udp+influxdb://usr:pwd@host1:8086,usr:pwd@host2:8086/db_name
+            https+influxdb://usr:pwd@host1:8086,usr:pwd@host2:8086/db_name
+
+        :param shuffle: If true, queries will hit servers evenly(randomly)
+        :param client_base_class: In order to support different clients,
+                                  default to InfluxDBClient
+        """
+        dsn = dsn.lower()
+        conn_params = urlparse(dsn)
+        netlocs = conn_params.netloc.split(',')
+        cluster_client = InfluxDBClusterClient(
+            hosts=[],
+            client_base_class=client_base_class,
+            shuffle=shuffle,
+            **kwargs)
+        for netloc in netlocs:
+            single_dsn = '%(scheme)s://%(netloc)s%(path)s' % (
+                {'scheme': conn_params.scheme,
+                 'netloc': netloc,
+                 'path': conn_params.path}
+            )
+            cluster_client.clients.append(client_base_class.from_DSN(
+                single_dsn,
+                **kwargs))
+        return cluster_client
+
     def _make_func(self, orig_func):
 
         @wraps(orig_func)
diff --git a/tests/influxdb/client_test.py b/tests/influxdb/client_test.py
@@ -623,3 +623,55 @@ def test_recovery(self):
         self.assertEqual('Success', cluster.query(''))
         self.assertEqual(1, len(cluster.clients))
         self.assertEqual(2, len(cluster.bad_clients))
+
+    def test_dsn(self):
+        cli = InfluxDBClusterClient.from_DSN(
+            'influxdb://usr:pwd@host1:8086,usr:pwd@host2:8086/db')
+        self.assertEqual(2, len(cli.clients))
+        self.assertEqual('http://host1:8086', cli.clients[0]._baseurl)
+        self.assertEqual('usr', cli.clients[0]._username)
+        self.assertEqual('pwd', cli.clients[0]._password)
+        self.assertEqual('db', cli.clients[0]._database)
+        self.assertFalse(cli.clients[0].use_udp)
+        self.assertEqual('http://host2:8086', cli.clients[1]._baseurl)
+        self.assertEqual('usr', cli.clients[1]._username)
+        self.assertEqual('pwd', cli.clients[1]._password)
+        self.assertEqual('db', cli.clients[1]._database)
+        self.assertFalse(cli.clients[1].use_udp)
+
+        cli = InfluxDBClusterClient.from_DSN(
+            'udp+influxdb://usr:pwd@host1:8086,usr:pwd@host2:8086/db')
+        self.assertTrue(cli.clients[0].use_udp)
+        self.assertTrue(cli.clients[1].use_udp)
+
+        cli = InfluxDBClusterClient.from_DSN(
+            'https+influxdb://usr:pwd@host1:8086,usr:pwd@host2:8086/db')
+        self.assertEqual('https://host1:8086', cli.clients[0]._baseurl)
+        self.assertEqual('https://host2:8086', cli.clients[1]._baseurl)
+
+        cli = InfluxDBClusterClient.from_DSN(
+            'https+influxdb://usr:pwd@host1:8086,usr:pwd@host2:8086/db',
+            **{'ssl': False})
+        self.assertEqual('http://host1:8086', cli.clients[0]._baseurl)
+        self.assertEqual('http://host2:8086', cli.clients[1]._baseurl)
+
+    def test_dsn_single_client(self):
+        cli = InfluxDBClusterClient.from_DSN('influxdb://usr:pwd@host:8086/db')
+        self.assertEqual('http://host:8086', cli.clients[0]._baseurl)
+        self.assertEqual('usr', cli.clients[0]._username)
+        self.assertEqual('pwd', cli.clients[0]._password)
+        self.assertEqual('db', cli.clients[0]._database)
+        self.assertFalse(cli.clients[0].use_udp)
+
+        cli = InfluxDBClusterClient.from_DSN(
+            'udp+influxdb://usr:pwd@host:8086/db')
+        self.assertTrue(cli.clients[0].use_udp)
+
+        cli = InfluxDBClusterClient.from_DSN(
+            'https+influxdb://usr:pwd@host:8086/db')
+        self.assertEqual('https://host:8086', cli.clients[0]._baseurl)
+
+        cli = InfluxDBClusterClient.from_DSN(
+            'https+influxdb://usr:pwd@host:8086/db',
+            **{'ssl': False})
+        self.assertEqual('http://host:8086', cli.clients[0]._baseurl)
