improved session id handling

README.md

-# aiogremlin 0.0.5 [(gizmo grew up)](https://pypi.python.org/pypi/gizmo/0.1.12)
+# aiogremlin 0.0.6 [(gizmo grew up)](https://pypi.python.org/pypi/gizmo/0.1.12)
 
 `aiogremlin` is a **Python 3** driver for the the [Tinkerpop 3 Gremlin Server](http://www.tinkerpop.com/docs/3.0.0.M7/#gremlin-server). This module is built on [Asyncio](https://docs.python.org/3/library/asyncio.html). By default it uses the [aiohttp](http://aiohttp.readthedocs.org/en/v0.15.3/index.html) websocket client , but it is easy to plug in a different implementation. `aiogremlin` is currently in **alpha** mode, but all major functionality has test coverage.
 

aiogremlin/__init__.py

@@ -5,4 +5,4 @@ from .client import (create_client, GremlinClient, GremlinResponse,
     GremlinResponseStream)
 from .exceptions import RequestError, GremlinServerError, SocketClientError
 from .protocol import GremlinWriter
-__version__ = "0.0.5"
+__version__ = "0.0.6"

aiogremlin/client.py

@@ -85,7 +85,7 @@ class GremlinClient:
 
     @asyncio.coroutine
     def submit(self, gremlin, connection=None, bindings=None, lang=None,
-               op=None, processor=None, binary=True):
+               op=None, processor=None, session=None, binary=True):
         """
         """
         lang = lang or self.lang
@@ -102,14 +102,15 @@ class GremlinClient:
             }
         }
         if processor == "session":
-            message["args"]["session"] = str(uuid.uuid4())
+            session = session or str(uuid.uuid4())
+            message["args"]["session"] = session
             client_logger.info(
                 "Session ID: {}".format(message["args"]["session"]))
         if connection is None:
             connection = yield from self.pool.connect(self.uri, loop=self.loop)
         writer = GremlinWriter(connection)
         connection = yield from writer.write(message, binary=binary)
-        return GremlinResponse(connection, loop=self._loop)
+        return GremlinResponse(connection, session=session, loop=self._loop)
 
     @asyncio.coroutine
     def execute(self, gremlin, bindings=None, lang=None,
@@ -126,14 +127,19 @@ class GremlinClient:
 
 class GremlinResponse:
 
-    def __init__(self, conn, loop=None):
+    def __init__(self, conn, session=None, loop=None):
         self._loop = loop or asyncio.get_event_loop()
+        self._session = session
         self._stream = GremlinResponseStream(conn, loop=self._loop)
 
     @property
     def stream(self):
         return self._stream
 
+    @property
+    def session(self):
+        return self._session
+
     @asyncio.coroutine
     def get(self):
         return (yield from self._run())

benchmark.py

@@ -55,7 +55,7 @@ def main(client, tests, count, concurrency, warmups, loop):
     execute = client.execute
     # warmup
     for x in range(warmups):
-        print("Warmup run {}:".format(x))
+        print("Warmup run {}:".format(x + 1))
         yield from run(client, count, concurrency, loop)
     print("Warmup successful!")
     mps_list = []

setup.py

@@ -3,7 +3,7 @@ from setuptools import setup
 
 setup(
     name="aiogremlin",
-    version="0.0.5",
+    version="0.0.6",
     url="",
     license="MIT",
     author="davebshow",

tests/tests.py

@@ -4,6 +4,7 @@
 import asyncio
 import itertools
 import unittest
+import uuid
 from aiogremlin import (GremlinClient, RequestError, GremlinServerError,
     SocketClientError, WebsocketPool, AiohttpFactory, create_client)
 
@@ -76,11 +77,24 @@ class GremlinClientTests(unittest.TestCase):
             error = True
         self.assertTrue(error)
 
-    def test_session(self):
+    def test_session_gen(self):
         execute = self.gc.execute("x + x", processor="session", bindings={"x": 4})
         results = self.loop.run_until_complete(execute)
         self.assertEqual(results[0].data[0], 8)
 
+    def test_session(self):
+        @asyncio.coroutine
+        def stream_coro():
+            session = str(uuid.uuid4())
+            resp = yield from self.gc.submit("x + x", bindings={"x": 4},
+                session=session)
+            while True:
+                f = yield from resp.stream.read()
+                if f is None:
+                    break
+            self.assertEqual(resp.session, session)
+        self.loop.run_until_complete(stream_coro())
+
 
 class WebsocketPoolTests(unittest.TestCase):