Skip to content

feat: support prefix/route when managed backend is behind a corporate… #134

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: master
Choose a base branch
from
Open

feat: support prefix/route when managed backend is behind a corporate… #134

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
17 changes: 17 additions & 0 deletions prestodb/client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@

import logging
import os
import re
from typing import Any, Dict, List, Optional, Text, Tuple, Union # NOQA for mypy types

import prestodb.redirect
Expand Down Expand Up @@ -164,6 +165,8 @@ class PrestoRequest(object):
:request_timeout: How long (in seconds) to wait for the server to send
data before giving up, as a float or a
``(connect timeout, read timeout)`` tuple.
:prefix: when the presto backend is configured behind a corporate proxy
on a route and presto is not managed by you

The client initiates a query by sending an HTTP POST to the
coordinator. It then gets a response back from the coordinator with:
Expand Down Expand Up @@ -213,12 +216,14 @@ def __init__(
request_timeout=constants.DEFAULT_REQUEST_TIMEOUT, # type: Union[float, Tuple[float, float]]
handle_retry=exceptions.RetryWithExponentialBackoff(),
service_account_file=None,
prefix=None,
):
# type: (...) -> None
self._client_session = client_session
self._host = host
self._port = port
self._next_uri = None # type: Optional[Text]
self._prefix = prefix

if http_session is not None:
self._http_session = http_session
Expand Down Expand Up @@ -329,11 +334,23 @@ def get_url(self, path):
@property
def statement_url(self):
# type: () -> Text
if self._prefix:
return self.get_url(f"{self._prefix}{constants.URL_STATEMENT_PATH}")

return self.get_url(constants.URL_STATEMENT_PATH)

@property
def next_uri(self):
# type: () -> Text
if self._prefix:
logger.debug(f"Updating url {self._next_uri} with prefix {self._prefix}")
pattern = r"(http[s]?://[^/]+)(/.*)"
try:
return re.sub(pattern, r"\1" + self._prefix + r"\2", self._next_uri)
except TypeError:
logger.error("Unable to update url with prefix. "
"Continuing without prefix.")
Comment on lines +347 to +352
Copy link

Choose a reason for hiding this comment

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

Instead of failing during query execution, does it make sense to raise the error during the initialization if the _prefix provided doesn't match the regex?


return self._next_uri

def post(self, sql):
Expand Down
3 changes: 3 additions & 0 deletions prestodb/dbapi.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -76,6 +76,7 @@ def __init__(
max_attempts=constants.DEFAULT_MAX_ATTEMPTS,
request_timeout=constants.DEFAULT_REQUEST_TIMEOUT,
isolation_level=IsolationLevel.AUTOCOMMIT,
prefix=None,
**kwargs,
):
self.host = host
Expand All @@ -84,6 +85,7 @@ def __init__(
self.source = source
self.catalog = catalog
self.schema = schema
self.prefix = prefix
self.session_properties = session_properties
self._client_session = prestodb.client.ClientSession(
user,
Expand Down Expand Up @@ -161,6 +163,7 @@ def _create_request(self):
self.redirect_handler,
self.max_attempts,
self.request_timeout,
prefix=self.prefix,
)

def cursor(self):
Expand Down