Check log level. Allow non-default json lib.

Author: 100LL

prestodb/__init__.py

@@ -20,4 +20,4 @@
 from . import exceptions
 from . import logging
 
-__version__ = "0.7.0"
+__version__ = "0.7.1"

prestodb/client.py

@@ -210,6 +210,7 @@ def __init__(
         request_timeout=constants.DEFAULT_REQUEST_TIMEOUT,  # type: Union[float, Tuple[float, float]]
         handle_retry=exceptions.RetryWithExponentialBackoff(),
         service_account_file=None,
+        json_lib=None,
     ):
         # type: (...) -> None
         self._client_session = ClientSession(
@@ -258,6 +259,7 @@ def __init__(
         self._handle_retry = handle_retry
         self.max_attempts = max_attempts
         self._http_scheme = http_scheme
+        self.json_lib = json_lib
 
     @property
     def transaction_id(self):
@@ -355,6 +357,8 @@ def post(self, sql):
             while http_response is not None and http_response.is_redirect:
                 location = http_response.headers["Location"]
                 url = self._redirect_handler.handle(location)
+                log_level = logger.getEffectiveLevel()
+                if log_level <= logging.logging.INFO:
                 logger.info(
                     "redirect {} from {} to {}".format(
                         http_response.status_code, location, url
@@ -406,9 +410,14 @@ def process(self, http_response):
         if not http_response.ok:
             self.raise_response_error(http_response)
 
+        log_level = logger.getEffectiveLevel()
         http_response.encoding = "utf-8"
-        response = http_response.json()
-        logger.debug("HTTP {}: {}".format(http_response.status_code, response))
+        if self.json_lib:
+            response = self.json_lib.loads(http_response.content)
+        else:
+            response = http_response.json()
+        if log_level <= logging.logging.DEBUG:
+            logger.debug("HTTP {}: {}".format(http_response.status_code, response))
         if "error" in response:
             raise self._process_error(response["error"], response.get("id"))
 
@@ -474,11 +483,13 @@ def __iter__(self):
         self._rows = None
 
         # Subsequent fetches from GET requests until next_uri is empty.
+        log_level = logger.getEffectiveLevel()
         while not self._query.is_finished():
             rows = self._query.fetch()
             for row in rows:
                 self._rownumber += 1
-                logger.debug("row {}".format(row))
+                if log_level <= logging.logging.DEBUG:
+                    logger.debug("row {}".format(row))
                 yield row
 
 
@@ -554,7 +565,9 @@ def fetch(self):
         if status.columns:
             self._columns = status.columns
         self._stats.update(status.stats)
-        logger.debug(status)
+        log_level = logger.getEffectiveLevel()
+        if log_level <= logging.logging.DEBUG:
+            logger.debug(status)
         if status.next_uri is None:
             self._finished = True
         return status.rows
@@ -565,13 +578,17 @@ def cancel(self):
         if self.query_id is None or self.is_finished():
             return
 
+        log_level = logger.getEffectiveLevel()
         self._cancelled = True
         url = self._request.get_url("/v1/query/{}".format(self.query_id))
-        logger.debug("cancelling query: %s", self.query_id)
+        if log_level <= logging.logging.DEBUG:
+            logger.debug("cancelling query: %s", self.query_id)
         response = self._request.delete(url)
-        logger.info(response)
+        if log_level <= logging.logging.INFO:
+            logger.info(response)
         if response.status_code == requests.codes.no_content:
-            logger.debug("query cancelled: %s", self.query_id)
+            if log_level <= logging.logging.DEBUG:
+                logger.debug("query cancelled: %s", self.query_id)
             return
         self._request.raise_response_error(response)
 

prestodb/dbapi.py

@@ -75,6 +75,7 @@ def __init__(
         max_attempts=constants.DEFAULT_MAX_ATTEMPTS,
         request_timeout=constants.DEFAULT_REQUEST_TIMEOUT,
         isolation_level=IsolationLevel.AUTOCOMMIT,
+        json_lib=None,
     ):
         self.host = host
         self.port = port
@@ -95,6 +96,7 @@ def __init__(
         self._isolation_level = isolation_level
         self._request = None
         self._transaction = None
+        self._json_lib = json_lib
 
     @property
     def isolation_level(self):
@@ -154,6 +156,9 @@ def _create_request(self):
             self.redirect_handler,
             self.max_attempts,
             self.request_timeout,
+            None,
+            None,
+            self._json_lib,
         )
 
     def cursor(self):