Feat(other): add support for custom request headers

Add support for custom request headers

ISSUES CLOSED: #2

CONTRIBUTORS.md

@@ -2,3 +2,4 @@
 
 * Jeffrey Phillips Freeman <jeffrey.freeman@syncleus.com> - Project owner from 2019 to present.
 * David M. Brown <davebshow@gmail.com> - Project founder and project owner from 2016 - 2019.
+* Guy Rozendorn <guy@rzn.co.il> 

aiogremlin/driver/aiohttp/transport.py

@@ -10,13 +10,13 @@ class AiohttpTransport(transport.AbstractBaseTransport):
         self._loop = loop
         self._connected = False
 
-    async def connect(self, url, *, ssl_context=None):
+    async def connect(self, url, *, ssl_context=None, headers=None):
         await self.close()
         connector = aiohttp.TCPConnector(
             ssl_context=ssl_context, loop=self._loop)
         self._client_session = aiohttp.ClientSession(
             loop=self._loop, connector=connector)
-        self._ws = await self._client_session.ws_connect(url)
+        self._ws = await self._client_session.ws_connect(url, headers=headers)
         self._connected = True
 
     async def write(self, message):

aiogremlin/driver/cluster.py

@@ -45,6 +45,7 @@ class Cluster:
         'ssl_certfile': '',
         'ssl_keyfile': '',
         'ssl_password': '',
+        'headers': None,
         'username': '',
         'password': '',
         'response_timeout': None,

aiogremlin/driver/connection.py

@@ -58,6 +58,7 @@ class Connection:
                    protocol=None,
                    transport_factory=None,
                    ssl_context=None,
+                   headers=None,
                    username='',
                    password='',
                    max_inflight=64,
@@ -73,6 +74,7 @@ class Connection:
             Protocol implementation
         :param transport_factory: Factory function for transports
         :param ssl.SSLContext ssl_context:
+        :param dict(str, str) headers:
         :param str username: Username for database auth
         :param str password: Password for database auth
 
@@ -89,7 +91,7 @@ class Connection:
         if not transport_factory:
             transport_factory = lambda: AiohttpTransport(loop)
         transport = transport_factory()
-        await transport.connect(url, ssl_context=ssl_context)
+        await transport.connect(url, ssl_context=ssl_context, headers=headers)
         return cls(url, transport, protocol, loop, username, password,
                    max_inflight, response_timeout, message_serializer,
                    provider)

aiogremlin/driver/pool.py

@@ -81,6 +81,7 @@ class ConnectionPool:
     :param str url: url for host Gremlin Server
     :param asyncio.BaseEventLoop loop:
     :param ssl.SSLContext ssl_context:
+    :param dict(str, str) headers:
     :param str username: Username for database auth
     :param str password: Password for database auth
     :param float response_timeout: (optional) `None` by default
@@ -92,12 +93,13 @@ class ConnectionPool:
         one time on the connection
     """
 
-    def __init__(self, url, loop, ssl_context, username, password, max_conns,
+    def __init__(self, url, loop, ssl_context, headers, username, password, max_conns,
                  min_conns, max_times_acquired, max_inflight, response_timeout,
                  message_serializer, provider):
         self._url = url
         self._loop = loop
         self._ssl_context = ssl_context
+        self._headers = headers
         self._username = username
         self._password = password
         self._max_conns = max_conns
@@ -194,7 +196,7 @@ class ConnectionPool:
     async def _get_connection(self, username, password, max_inflight,
                               response_timeout, message_serializer, provider):
         conn = await connection.Connection.open(
-            self._url, self._loop, ssl_context=self._ssl_context,
+            self._url, self._loop, ssl_context=self._ssl_context, headers=self._headers,
             username=username, password=password,
             response_timeout=response_timeout,
             message_serializer=message_serializer, provider=provider)

aiogremlin/driver/server.py

@@ -35,6 +35,7 @@ class GremlinServer:
             self._ssl_context = ssl_context
         else:
             self._ssl_context = None
+        self._headers = config['headers']
 
     @property
     def url(self):
@@ -66,8 +67,8 @@ class GremlinServer:
 
     async def initialize(self):
         conn_pool = pool.ConnectionPool(
-            self._url, self._loop, self._ssl_context, self._username,
-            self._password, self._max_conns, self._min_conns,
+            self._url, self._loop, self._ssl_context, self._headers,
+            self._username, self._password, self._max_conns, self._min_conns,
             self._max_times_acquired, self._max_inflight,
             self._response_timeout, self._message_serializer, self._provider)
         await conn_pool.init_pool()
@@ -81,6 +82,7 @@ class GremlinServer:
         :param str url: url for host Gremlin Server
         :param asyncio.BaseEventLoop loop:
         :param ssl.SSLContext ssl_context:
+        :param dict(str, str) headers:
         :param str username: Username for database auth
         :param str password: Password for database auth
         :param float response_timeout: (optional) `None` by default

setup.py

@@ -32,7 +32,7 @@ setup(
               'aiogremlin.remote'],
     python_requires='>=3.5',
     install_requires=[
-        'gremlinpython<=3.4.3',
+        'gremlinpython>=3.4.8',
         'aenum>=1.4.5',  # required gremlinpython dep
         'aiohttp>=2.2.5',
         'PyYAML>=3.12',