From 4555e6a560e911a0f52982dd9b7c50e79db51673 Mon Sep 17 00:00:00 2001 From: Can ZHANG Date: Wed, 22 Apr 2015 17:37:47 +0800 Subject: [PATCH] Add `from_DSN` to InfluxDBClusterClient --- influxdb/client.py | 43 +++++++++++++++++++++++++++-- tests/influxdb/client_test.py | 52 +++++++++++++++++++++++++++++++++++ 2 files changed, 92 insertions(+), 3 deletions(-) diff --git a/influxdb/client.py b/influxdb/client.py index e47bd326..f13bb116 100755 --- 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 index a55c65ee..e4eed104 100644 --- 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)