Commit aaa87f83 authored by Guy Rozendorn's avatar Guy Rozendorn
Browse files

Feat(other): add support for custom request headers

Add support for custom request headers

ISSUES CLOSED: #2
parent 94a0c29c
Pipeline #356 failed with stage
in 3 minutes and 57 seconds
......@@ -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>
......@@ -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):
......
......@@ -45,6 +45,7 @@ class Cluster:
'ssl_certfile': '',
'ssl_keyfile': '',
'ssl_password': '',
'headers': None,
'username': '',
'password': '',
'response_timeout': None,
......
......@@ -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)
......
......@@ -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)
......
......@@ -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
......
......@@ -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',
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment