Skip to content
Merged
Changes from 4 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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.
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: single backticks here on requests.Session

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
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

likely we should have the be private (ie. start with an underscore)


@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