Use same transaction and HTTP session within a connection

When an HTTP gateway is used to balance queries across multiple Presto
coordinators, queries within the same transactions must go to the same
coordinator. A single transation is started for a given connection. An HTTP
cookie can support this behavior such as
`Set-Cookie: Presto-Gateway-Sticky={host}:{port};Version=1`. The cookie is
persisted by the `requests.Session` object referenced by
`PrestoRequest._http_session`. Hence all queries should be sent with the same
session data. Presto client protocol is connection-less. Here the connection is
used to persist session data and use the same connection parameters.

Author: Greg Leclercq

integration_tests/test_dbapi.py

@@ -65,11 +65,13 @@ def test_select_query(presto_connection):
 
 
 def test_select_query_result_iteration(presto_connection):
-    cur = presto_connection.cursor()
-    cur.execute('select custkey from tpch.sf1.customer LIMIT 10')
-    rows0 = cur.genall()
-    cur.execute('select custkey from tpch.sf1.customer LIMIT 10')
-    rows1 = cur.fetchall()
+    cur0 = presto_connection.cursor()
+    cur0.execute('select custkey from tpch.sf1.customer LIMIT 10')
+    rows0 = cur0.genall()
+
+    cur1 = presto_connection.cursor()
+    cur1.execute('select custkey from tpch.sf1.customer LIMIT 10')
+    rows1 = cur1.fetchall()
 
     assert len(list(rows0)) == len(rows1)
 
@@ -177,7 +179,7 @@ def test_transaction_single(presto_connection_with_transaction):
 
 def test_transaction_rollback(presto_connection_with_transaction):
     connection = presto_connection_with_transaction
-    for i in range(3):
+    for _ in range(3):
         cur = connection.cursor()
         cur.execute('SELECT * FROM tpch.sf1.customer LIMIT 1000')
         rows = cur.fetchall()

prestodb/client.py

@@ -203,6 +203,7 @@ def __init__(
         catalog=None,  # type: Text
         schema=None,  # type: Text
         session_properties=None,  # type: Optional[Dict[Text, Any]]
+        http_session=None,  # type: Any
         http_headers=None,  # type: Optional[Dict[Text, Text]]
         transaction_id=NO_TRANSACTION,  # type: Optional[Text]
         http_scheme=constants.HTTP,  # type: Text
@@ -226,8 +227,12 @@ def __init__(
         self._host = host
         self._port = port
         self._next_uri = None  # type: Optional[Text]
-        # mypy cannot follow module import
-        self._http_session = self.http.Session()  # type: ignore
+
+        if http_session is not None:
+            self._http_session = http_session
+        else:
+            # mypy cannot follow module import
+            self._http_session = self.http.Session()  # type: ignore
         self._http_session.headers.update(self.http_headers)
         self._auth = auth
         if self._auth:
@@ -413,9 +418,9 @@ def process(self, http_response):
 
         if constants.HEADER_SET_SESSION in http_response.headers:
             for key, value in get_session_property_values(
-                    response.headers,
-                    constants.HEADER_SET_SESSION,
-                ):
+                response.headers,
+                constants.HEADER_SET_SESSION,
+            ):
                 self._client_session.properties[key] = value
 
         self._next_uri = response.get('nextUri')

prestodb/dbapi.py

@@ -100,6 +100,8 @@ def __init__(
         self.catalog = catalog
         self.schema = schema
         self.session_properties = session_properties
+        # mypy cannot follow module import
+        self._http_session = prestodb.client.PrestoRequest.http.Session()
         self.http_headers = http_headers
         self.http_scheme = http_scheme
         self.auth = auth
@@ -108,6 +110,7 @@ def __init__(
         self.request_timeout = request_timeout
 
         self._isolation_level = isolation_level
+        self._request = None
         self._transaction = None
 
     @property
@@ -124,7 +127,7 @@ def __enter__(self):
     def __exit__(self, exc_type, exc_value, traceback):
         try:
             self.commit()
-        except:
+        except Exception:
             self.rollback()
         else:
             self.close()
@@ -135,14 +138,32 @@ def close(self):
         pass
 
     def start_transaction(self):
-        request = prestodb.client.PrestoRequest(
+        self._transaction = Transaction(self._create_request())
+        self._transaction.begin()
+        return self._transaction
+
+    def commit(self):
+        if self.transaction is None:
+            return
+        self._transaction.commit()
+        self._transaction = None
+
+    def rollback(self):
+        if self.transaction is None:
+            raise RuntimeError('no transaction was started')
+        self._transaction.rollback()
+        self._transaction = None
+
+    def _create_request(self):
+        return prestodb.client.PrestoRequest(
             self.host,
             self.port,
             self.user,
             self.source,
             self.catalog,
             self.schema,
             self.session_properties,
+            self._http_session,
             self.http_headers,
             NO_TRANSACTION,
             self.http_scheme,
@@ -151,25 +172,10 @@ def start_transaction(self):
             self.max_attempts,
             self.request_timeout,
         )
-        self._transaction = Transaction(request)
-        self._transaction.begin()
-        return self._transaction
-
-    def commit(self):
-        if self.transaction is None:
-            return
-        self._transaction.commit()
-        self._transaction = None
-
-    def rollback(self):
-        if self.transaction is None:
-            raise RuntimeError('no transaction was started')
-        self._transaction.rollback()
-        self._transaction = None
 
     def cursor(self):
         """Return a new :py:class:`Cursor` object using the connection."""
-        return Cursor(self)
+        return Cursor(self, self._create_request())
 
 
 class Cursor(object):
@@ -179,13 +185,14 @@ class Cursor(object):
     cursor are immediately visible by other cursors or connections.
 
     """
-    def __init__(self, connection):
+    def __init__(self, connection, request):
         if not isinstance(connection, Connection):
             raise ValueError(
                 'connection must be a Connection object: {}'.format(
                     type(connection)
                 ))
         self._connection = connection
+        self._request = request
 
         self.arraysize = 1
         self._iterator = None
@@ -195,7 +202,6 @@ def __init__(self, connection):
     def connection(self):
         return self._connection
 
-
     @property
     def description(self):
         if self._query.columns is None:
@@ -234,28 +240,8 @@ def execute(self, operation, params=None):
         if self.connection.isolation_level != IsolationLevel.AUTOCOMMIT:
             if self.connection.transaction is None:
                 self.connection.start_transaction()
-            transaction_id = self.connection.transaction.id
-        else:
-            transaction_id = 'NONE'
-
-        request = prestodb.client.PrestoRequest(
-            self.connection.host,
-            self.connection.port,
-            self.connection.user,
-            self.connection.source,
-            self.connection.catalog,
-            self.connection.schema,
-            self.connection.session_properties,
-            self.connection.http_headers,
-            transaction_id,
-            self.connection.http_scheme,
-            self.connection.auth,
-            self.connection.redirect_handler,
-            self.connection.max_attempts,
-            self.connection.request_timeout,
-        )
 
-        self._query = prestodb.client.PrestoQuery(request, sql=operation)
+        self._query = prestodb.client.PrestoQuery(self._request, sql=operation)
         result = self._query.execute()
         self._iterator = iter(result)
         return result

prestodb/transaction.py

@@ -87,22 +87,25 @@ def begin(self):
                     response.headers[constants.HEADER_STARTED_TRANSACTION]
                 )
             status = self._request.process(response)
+        self._request.transaction_id = self._id
         logger.info('transaction started: ' + self._id)
 
     def commit(self):
-        self._request.transaction_id = self._id
         query = prestodb.client.PrestoQuery(self._request, COMMIT)
         try:
             list(query.execute())
         except Exception as err:
             raise prestodb.exceptions.DatabaseError(
                 'failed to commit transaction {}: {}'.format(self._id, err))
+        self._id = NO_TRANSACTION
+        self._request.transaction_id = self._id
 
     def rollback(self):
-        self._request.transaction_id = self._id
         query = prestodb.client.PrestoQuery(self._request, ROLLBACK)
         try:
             list(query.execute())
         except Exception as err:
             raise prestodb.exceptions.DatabaseError(
                 'failed to rollback transaction {}: {}'.format(self._id, err))
+        self._id = NO_TRANSACTION
+        self._request.transaction_id = self._id