Make ClientSession thread-safe

cherry-pick of trinodb/trino-python-client@79a4814

Co-authored-by: Michiel De Smet <[email protected]>

Author: mlyublena

README.md

@@ -63,7 +63,7 @@ Following example shows a use case where both Kerberos and Oauth authentication
 ```python
 import getpass
 import prestodb
-from prestodb.client import PrestoRequest, PrestoQuery
+from prestodb.client import ClientSession, PrestoRequest, PrestoQuery
 from requests_kerberos import DISABLED
 
 kerberos_auth = prestodb.auth.KerberosAuthentication(
@@ -76,7 +76,7 @@ kerberos_auth = prestodb.auth.KerberosAuthentication(
 req = PrestoRequest(
     host='GCP coordinator url',
     port=443,
-    user=getpass.getuser(),
+    client_session=ClientSession(user=getpass.getuser()),
     service_account_file='Service account json file path',
     http_scheme='https',
     auth=kerberos_auth

integration_tests/fixtures.py

@@ -23,7 +23,7 @@
 import logging
 import pytest
 import requests
-from prestodb.client import PrestoQuery, PrestoRequest
+from prestodb.client import ClientSession, PrestoQuery, PrestoRequest
 from prestodb.constants import DEFAULT_PORT
 from prestodb.exceptions import TimeoutError
 
@@ -110,7 +110,7 @@ def start_presto(image_tag=None, build=True, with_cache=True):
 
 
 def wait_for_presto_workers(host, port, timeout=30):
-    request = PrestoRequest(host=host, port=port, user="test_fixture")
+    request = PrestoRequest(host=host, port=port, client_session=ClientSession(user="test_fixture"))
     sql = "SELECT state FROM system.runtime.nodes"
     t0 = time.time()
     while True:

prestodb/client.py

@@ -37,10 +37,10 @@
 import logging
 import os
 from typing import Any, Dict, List, Optional, Text, Tuple, Union  # NOQA for mypy types
-import six.moves.urllib_parse as parse
 
 import prestodb.redirect
 import requests
+import six.moves.urllib_parse as parse
 from prestodb import constants, exceptions
 from prestodb.transaction import NO_TRANSACTION
 
@@ -62,10 +62,10 @@
 class ClientSession(object):
     def __init__(
         self,
-        catalog,
-        schema,
-        source,
         user,
+        catalog = None,
+        schema = None,
+        source = None,
         properties=None,
         headers=None,
         transaction_id=None,
@@ -80,6 +80,9 @@ def __init__(
         self._headers = headers or {}
         self.transaction_id = transaction_id
 
+    def __repr__(self):
+        return f"ClientSession({self.catalog}, {self.schema}, {self.source}, {self.user}, {self._properties}, {self._headers}, {self.transaction_id})"
+
     @property
     def properties(self):
         return self._properties
@@ -199,11 +202,7 @@ def __init__(
         self,
         host,  # type: Text
         port,  # type: int
-        user,  # type: Text
-        source=None,  # type: Text
-        catalog=None,  # type: Text
-        schema=None,  # type: Text
-        session_properties=None,  # type: Optional[Dict[Text, Any]]
+        client_session,  # type: ClientSession
         http_session=None,  # type: Any
         http_headers=None,  # type: Optional[Dict[Text, Text]]
         transaction_id=NO_TRANSACTION,  # type: Optional[Text]
@@ -216,16 +215,7 @@ def __init__(
         service_account_file=None,
     ):
         # type: (...) -> None
-        self._client_session = ClientSession(
-            catalog,
-            schema,
-            source,
-            user,
-            session_properties,
-            http_headers,
-            transaction_id,
-        )
-
+        self._client_session = client_session
         self._host = host
         self._port = port
         self._next_uri = None  # type: Optional[Text]
@@ -539,18 +529,25 @@ def execute(self):
 
         response = self._request.post(self._sql)
         status = self._request.process(response)
+        if status.next_uri is None:
+            self._finished = True
         self.query_id = status.id
-        self._stats.update({u"queryId": self.query_id})
+        self._stats.update({"queryId": self.query_id})
         self._stats.update(status.stats)
         self._warnings = getattr(status, "warnings", [])
-        if status.next_uri is None:
-            self._finished = True
         self._result = PrestoResult(self, status.rows)
+        while (
+            not self._finished and not self._cancelled
+        ):
+            self._result._rows += self.fetch()
         return self._result
 
     def fetch(self):
         # type: () -> List[List[Any]]
         """Continue fetching data for the current query_id"""
+        if self._request.next_uri is None:
+            self._finished = True
+            return []
         response = self._request.get(self._request.next_uri)
         status = self._request.process(response)
         if status.columns:

prestodb/dbapi.py

@@ -17,20 +17,21 @@
 Fetch methods returns rows as a list of lists on purpose to let the caller
 decide to convert then to a list of tuples.
 """
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
+from __future__ import absolute_import, division, print_function
 
+import binascii
+import datetime
 import logging
+import uuid
 from typing import Any, List, Optional  # NOQA for mypy types
-import datetime
 
-from prestodb import constants
-import prestodb.exceptions
 import prestodb.client
+import prestodb.exceptions
 import prestodb.redirect
-from prestodb.transaction import Transaction, IsolationLevel, NO_TRANSACTION
 
+from prestodb import constants
+from prestodb.transaction import IsolationLevel, NO_TRANSACTION, Transaction
+from prestodb.transaction import NO_TRANSACTION
 
 __all__ = ["connect", "Connection", "Cursor"]
 
@@ -75,6 +76,7 @@ def __init__(
         max_attempts=constants.DEFAULT_MAX_ATTEMPTS,
         request_timeout=constants.DEFAULT_REQUEST_TIMEOUT,
         isolation_level=IsolationLevel.AUTOCOMMIT,
+        **kwargs,
     ):
         self.host = host
         self.port = port
@@ -83,6 +85,15 @@ def __init__(
         self.catalog = catalog
         self.schema = schema
         self.session_properties = session_properties
+        self._client_session = prestodb.client.ClientSession(
+            user,
+            catalog,
+            schema,
+            source,
+            session_properties,
+            http_headers,
+            NO_TRANSACTION,
+        )
         # mypy cannot follow module import
         self._http_session = prestodb.client.PrestoRequest.http.Session()
         self.http_headers = http_headers
@@ -141,11 +152,7 @@ def _create_request(self):
         return prestodb.client.PrestoRequest(
             self.host,
             self.port,
-            self.user,
-            self.source,
-            self.catalog,
-            self.schema,
-            self.session_properties,
+            self._client_session,
             self._http_session,
             self.http_headers,
             NO_TRANSACTION,
@@ -187,6 +194,9 @@ def __init__(self, connection, request):
         self._iterator = None
         self._query = None
 
+    def __iter__(self):
+        return self._iterator
+
     @property
     def connection(self):
         return self._connection

tests/test_client.py

@@ -20,7 +20,7 @@
 import time
 
 from requests_kerberos.exceptions import KerberosExchangeError
-from prestodb.client import PrestoRequest
+from prestodb.client import ClientSession, PrestoRequest
 from prestodb.auth import KerberosAuthentication
 from prestodb import constants
 import prestodb.exceptions
@@ -270,12 +270,14 @@ def test_presto_initial_request(monkeypatch):
     req = PrestoRequest(
         host="coordinator",
         port=8080,
-        user="test",
-        source="test",
-        catalog="test",
-        schema="test",
-        http_scheme="http",
-        session_properties={},
+        client_session=ClientSession(
+            user="test",
+            source="test",
+            catalog="test",
+            schema="test",
+            properties={},
+        ),
+        http_scheme="http"
     )
 
     http_resp = PrestoRequest.http.Response()
@@ -318,16 +320,17 @@ def test_request_headers(monkeypatch):
     req = PrestoRequest(
         host="coordinator",
         port=8080,
-        user=user,
-        source=source,
-        catalog=catalog,
-        schema=schema,
+        client_session=ClientSession(
+            user=user,
+            source=source,
+            catalog=catalog,
+            schema=schema,
+            properties={"hash_partition_count": 500, "needs_url_encoding": 'foo,bar'},
+            headers={
+                accept_encoding_header: accept_encoding_value,
+                client_info_header: client_info_value,
+            }),
         http_scheme="http",
-        session_properties={"hash_partition_count": 500, "needs_url_encoding": 'foo,bar'},
-        http_headers={
-            accept_encoding_header: accept_encoding_value,
-            client_info_header: client_info_value,
-        },
         redirect_handler=None,
     )
 
@@ -353,8 +356,9 @@ def test_request_invalid_http_headers():
         PrestoRequest(
             host="coordinator",
             port=8080,
-            user="test",
-            http_headers={constants.HEADER_USER: "invalid_header"},
+            client_session=ClientSession(
+                user="test",
+                headers={constants.HEADER_USER: "invalid_header"})
         )
     assert str(value_error.value).startswith("cannot override reserved HTTP header")
 
@@ -379,7 +383,8 @@ def long_call(request, uri, headers):
         req = PrestoRequest(
             host=host,
             port=port,
-            user="test",
+            client_session = ClientSession(
+                user="test"),
             http_scheme=http_scheme,
             max_attempts=1,
             request_timeout=request_timeout,
@@ -401,12 +406,13 @@ def test_presto_fetch_request(monkeypatch):
     req = PrestoRequest(
         host="coordinator",
         port=8080,
-        user="test",
-        source="test",
-        catalog="test",
-        schema="test",
+        client_session = ClientSession(
+            user="test",
+            source="test",
+            catalog="test",
+            schema="test",
+            properties={}),
         http_scheme="http",
-        session_properties={},
     )
 
     http_resp = PrestoRequest.http.Response()
@@ -424,12 +430,13 @@ def test_presto_fetch_error(monkeypatch):
     req = PrestoRequest(
         host="coordinator",
         port=8080,
-        user="test",
-        source="test",
-        catalog="test",
-        schema="test",
+        client_session = ClientSession(
+            user="test",
+            source="test",
+            catalog="test",
+            schema="test",
+            properties={}),
         http_scheme="http",
-        session_properties={},
     )
 
     http_resp = PrestoRequest.http.Response()
@@ -465,12 +472,13 @@ def test_presto_connection_error(monkeypatch, error_code, error_type, error_mess
     req = PrestoRequest(
         host="coordinator",
         port=8080,
-        user="test",
-        source="test",
-        catalog="test",
-        schema="test",
+        client_session = ClientSession(
+            user="test",
+            source="test",
+            catalog="test",
+            schema="test",
+            properties={}),
         http_scheme="http",
-        session_properties={},
     )
 
     http_resp = PrestoRequest.http.Response()
@@ -511,7 +519,8 @@ def test_authentication_fail_retry(monkeypatch):
     req = PrestoRequest(
         host="coordinator",
         port=8080,
-        user="test",
+        client_session = ClientSession(
+            user="test"),
         http_scheme=constants.HTTPS,
         auth=kerberos_auth,
         max_attempts=attempts,
@@ -538,7 +547,7 @@ def test_503_error_retry(monkeypatch):
 
     attempts = 3
     req = PrestoRequest(
-        host="coordinator", port=8080, user="test", max_attempts=attempts
+        host="coordinator", port=8080, client_session=ClientSession(user="test"), max_attempts=attempts
     )
 
     req.post("URL")
@@ -576,7 +585,7 @@ def test_gateway_redirect(monkeypatch):
         socket, "gethostbyaddr", lambda *args: ("finalhost", ["finalhost"], "1.2.3.4")
     )
 
-    req = PrestoRequest(host="coordinator", port=8080, user="test")
+    req = PrestoRequest(host="coordinator", port=8080, client_session=ClientSession(user="test"))
     result = req.post("http://host:80/path/")
     assert gateway_response.count == 3
     assert result.ok