Skip to content

Commit

Permalink
Make the library runnable w/ and w/o requests_kerberos
Browse files Browse the repository at this point in the history
  • Loading branch information
takuti authored and ggreg committed Dec 21, 2018
1 parent 25cf54c commit dfe0240
Show file tree
Hide file tree
Showing 2 changed files with 16 additions and 6 deletions.
5 changes: 5 additions & 0 deletions prestodb/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,11 @@
from __future__ import division
from __future__ import print_function

try:
from .auth import KerberosAuthentication
except ImportError:
pass

from . import dbapi
from . import client
from . import constants
Expand Down
17 changes: 11 additions & 6 deletions prestodb/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,6 @@
from __future__ import print_function

import os
from requests_kerberos.exceptions import KerberosExchangeError
from typing import Any, Dict, List, Optional, Text, Tuple, Union # NOQA for mypy types

import requests
Expand Down Expand Up @@ -240,6 +239,16 @@ def __init__(
raise ValueError('cannot use authentication with HTTP')
self._auth.set_http_session(self._http_session)

default_exceptions = (
PrestoRequest.http.ConnectionError, # type: ignore
PrestoRequest.http.Timeout, # type: ignore
)
try:
from requests_kerberos.exceptions import KerberosExchangeError
self._exceptions = default_exceptions + (KerberosExchangeError,)
except ImportError:
self._exceptions = default_exceptions

self._redirect_handler = redirect_handler
self._request_timeout = request_timeout
self._handle_retry = handle_retry
Expand Down Expand Up @@ -298,11 +307,7 @@ def max_attempts(self, value):

with_retry = exceptions.retry_with(
self._handle_retry,
exceptions=(
KerberosExchangeError,
PrestoRequest.http.ConnectionError, # type: ignore
PrestoRequest.http.Timeout, # type: ignore
),
exceptions=self._exceptions,
conditions=(
# need retry when there is no exception but the status code is 503
lambda response: getattr(response, 'status_code', None) == 503,
Expand Down

0 comments on commit dfe0240

Please sign in to comment.