From 566f65a5f2b7f2e7a007303c82f3d3b1a3523413 Mon Sep 17 00:00:00 2001
From: davebshow <davebshow@gmail.com>
Date: Sun, 12 Mar 2017 21:43:13 -0400
Subject: [PATCH] clients can be assigned to a specific host
---
aiogremlin/driver/client.py | 5 +++--
aiogremlin/driver/cluster.py | 22 ++++++++++++++++------
aiogremlin/gremlin_python/__init__.py | 2 +-
setup.py | 2 +-
tests/test_client.py | 11 +++++++++++
5 files changed, 32 insertions(+), 10 deletions(-)
diff --git a/aiogremlin/driver/client.py b/aiogremlin/driver/client.py
index 2759d68..73852b5 100644
--- a/aiogremlin/driver/client.py
+++ b/aiogremlin/driver/client.py
@@ -16,11 +16,12 @@ class Client:
:param asyncio.BaseEventLoop loop:
:param dict aliases: Optional mapping for aliases. Default is `None`
"""
- def __init__(self, cluster, loop, *, aliases=None):
+ def __init__(self, cluster, loop, *, hostname=None, aliases=None):
self._cluster = cluster
self._loop = loop
if aliases is None:
aliases = {}
+ self._hostname = hostname
self._aliases = aliases
@property
@@ -76,7 +77,7 @@ class Client:
'aliases': self._aliases})
if bindings:
message.args.update({'bindings': bindings})
- conn = await self.cluster.get_connection()
+ conn = await self.cluster.get_connection(hostname=self._hostname)
resp = await conn.write(message)
self._loop.create_task(conn.release_task(resp))
return resp
diff --git a/aiogremlin/driver/cluster.py b/aiogremlin/driver/cluster.py
index 6cc403f..6f1e02f 100644
--- a/aiogremlin/driver/cluster.py
+++ b/aiogremlin/driver/cluster.py
@@ -63,6 +63,7 @@ class Cluster:
default_config.update(config)
self._config = self._process_config_imports(default_config)
self._hosts = collections.deque()
+ self._hostmap = {}
self._closed = False
if aliases is None:
aliases = {}
@@ -100,7 +101,7 @@ class Cluster:
"""
return self._config
- async def get_connection(self):
+ async def get_connection(self, hostname=None):
"""
**coroutine** Get connection from next available host in a round robin
fashion.
@@ -109,7 +110,14 @@ class Cluster:
"""
if not self._hosts:
await self.establish_hosts()
- host = self._hosts.popleft()
+ if hostname:
+ try:
+ host = self._hostmap[hostname]
+ except KeyError:
+ raise exception.ConfigError(
+ 'Unknown host: {}'.format(hostname))
+ else:
+ host = self._hosts.popleft()
conn = await host.get_connection()
self._hosts.append(host)
return conn
@@ -121,11 +129,12 @@ class Cluster:
scheme = self._config['scheme']
hosts = self._config['hosts']
port = self._config['port']
- for host in hosts:
- url = '{}://{}:{}/gremlin'.format(scheme, host, port)
+ for hostname in hosts:
+ url = '{}://{}:{}/gremlin'.format(scheme, hostname, port)
host = await driver.GremlinServer.open(
url, self._loop, **dict(self._config))
self._hosts.append(host)
+ self._hostmap[hostname] = host
def config_from_file(self, filename):
"""
@@ -186,7 +195,7 @@ class Cluster:
config = self._process_config_imports(config)
self.config.update(config)
- async def connect(self, aliases=None):
+ async def connect(self, hostname=None, aliases=None):
"""
**coroutine** Get a connected client. Main API method.
@@ -202,7 +211,8 @@ class Cluster:
# aliases=aliases)
# self._hosts.append(host)
# else:
- client = driver.Client(self, self._loop, aliases=aliases)
+ client = driver.Client(self, self._loop, hostname=hostname,
+ aliases=aliases)
return client
async def close(self):
diff --git a/aiogremlin/gremlin_python/__init__.py b/aiogremlin/gremlin_python/__init__.py
index 41aee8b..a2a5a4b 100644
--- a/aiogremlin/gremlin_python/__init__.py
+++ b/aiogremlin/gremlin_python/__init__.py
@@ -22,4 +22,4 @@ __author__ = 'Marko A. Rodriguez (http://markorodriguez.com)'
from aiogremlin.gremlin_python.statics import *
from aiogremlin.gremlin_python.process.graph_traversal import __
from aiogremlin.gremlin_python.process.strategies import *
-from aiogremlin.gremlin_python.process.traversal import Binding
+from aiogremlin.gremlin_python.process.traversal import Binding, Cardinality
diff --git a/setup.py b/setup.py
index 3e068f0..9a8b7bc 100644
--- a/setup.py
+++ b/setup.py
@@ -3,7 +3,7 @@ from setuptools import setup
setup(
name='aiogremlin',
- version='3.2.4b1',
+ version='3.2.4b2',
url='',
license='Apache Software License',
author='davebshow',
diff --git a/tests/test_client.py b/tests/test_client.py
index bf9aacd..226e4dd 100644
--- a/tests/test_client.py
+++ b/tests/test_client.py
@@ -45,6 +45,17 @@ async def test_alias(cluster):
await cluster.close()
+@pytest.mark.asyncio
+async def test_client_auto_release(cluster):
+ client = await cluster.connect(hostname='localhost')
+ resp = await client.submit("1 + 1")
+ async for msg in resp:
+ assert msg == 2
+ assert client._hostname == 'localhost'
+ await cluster.close()
+
+
+
# @pytest.mark.asyncio
# async def test_sessioned_client(cluster):
# session = str(uuid.uuid4())
--
GitLab