Skip to content

Commit

Permalink
Support custom requests session in presto connections (#106)
Browse files Browse the repository at this point in the history
* add connection options to presto

* requests_session

* format

* doc

* address comments
  • Loading branch information
gthomas-slack committed Oct 9, 2020
1 parent c62fe02 commit c067777
Showing 1 changed file with 7 additions and 2 deletions.
9 changes: 7 additions & 2 deletions omniduct/databases/presto.py
Original file line number Diff line number Diff line change
Expand Up @@ -49,21 +49,25 @@ def NAMESPACE_DEFAULT(self):
}

@override
def _init(self, catalog='default', schema='default', server_protocol='http', source=None):
def _init(self, catalog='default', schema='default', server_protocol='http', source=None, requests_session=None):
"""
catalog (str): The default catalog to use in database queries.
schema (str): The default schema/database to use in database queries.
server_protocol (str): The protocol over which to connect to the Presto REST
service ('http' or 'https'). (default='http')
source (str): The source of this query (by default "omniduct <version>").
If manually specified, result will be: "<source> / omniduct <version>".
requests_session (requests.Session): an optional requests.Session object for advanced usage.
Passed through to the pyhive Cursor which supports custom requests sessions for advanced usage
such as custom headers, cookie values, retry logic, etc.
"""
self.catalog = catalog
self.schema = schema
self.server_protocol = server_protocol
self.source = source
self.__presto = None
self.connection_fields += ('catalog', 'schema')
self._requests_session = requests_session

@property
def source(self):
Expand Down Expand Up @@ -115,7 +119,8 @@ def _execute(self, statement, cursor, wait, session_properties):
cursor = cursor or presto.Cursor(
host=self.host, port=self.port, username=self.username, password=self.password,
catalog=self.catalog, schema=self.schema, session_props=session_properties,
poll_interval=1, source=self.source, protocol=self.server_protocol
poll_interval=1, source=self.source, protocol=self.server_protocol,
requests_session=self._requests_session
)
cursor.execute(statement)
status = cursor.poll()
Expand Down

0 comments on commit c067777

Please sign in to comment.