enable presto headers if needed

Author: Ronen Hoffer

trino/client.py

@@ -206,7 +206,8 @@ def __init__(
         max_attempts: int = MAX_ATTEMPTS,
         request_timeout: Union[float, Tuple[float, float]] = constants.DEFAULT_REQUEST_TIMEOUT,
         handle_retry=exceptions.RetryWithExponentialBackoff(),
-        verify: bool = True
+        verify: bool = True,
+        presto_headers = False,
     ) -> None:
         self._client_session = ClientSession(
             catalog,
@@ -227,6 +228,7 @@ def __init__(
         else:
             self._http_session = self.http.Session()
             self._http_session.verify = verify
+        self._headers_name = constants.get_headers(presto_headers=presto_headers)
         self._http_session.headers.update(self.http_headers)
         self._exceptions = self.HTTP_EXCEPTIONS
         self._auth = auth
@@ -252,28 +254,20 @@ def transaction_id(self, value):
 
     @property
     def http_headers(self) -> Dict[str, str]:
-        headers = {}
-
-        headers[constants.HEADER_CATALOG] = self._client_session.catalog
-        headers[constants.HEADER_SCHEMA] = self._client_session.schema
-        headers[constants.HEADER_SOURCE] = self._client_session.source
-        headers[constants.HEADER_USER] = self._client_session.user
-
-        headers[constants.HEADER_SESSION] = ",".join(
-            # ``name`` must not contain ``=``
-            "{}={}".format(name, value)
-            for name, value in self._client_session.properties.items()
-        )
+        headers = {
+            self._headers_name.catalog: self._client_session.catalog,
+            self._headers_name.schema: self._client_session.schema,
+            self._headers_name.source: self._client_session.source,
+            self._headers_name.user: self._client_session.user,
+            self._headers_name.session: ",".join("{}={}".format(name, value) for name, value in self._client_session.properties.items()),
+        }
 
         # merge custom http headers
         for key in self._client_session.headers:
             if key in headers.keys():
                 raise ValueError("cannot override reserved HTTP header {}".format(key))
         headers.update(self._client_session.headers)
-
-        transaction_id = self._client_session.transaction_id
-        headers[constants.HEADER_TRANSACTION] = transaction_id
-
+        headers[self._headers_name.transaction] = self._client_session.transaction_id
         return headers
 
     @property
@@ -386,18 +380,19 @@ def process(self, http_response) -> TrinoStatus:
         http_response.encoding = "utf-8"
         response = http_response.json()
         logger.debug("HTTP %s: %s", http_response.status_code, response)
+
         if "error" in response:
             raise self._process_error(response["error"], response.get("id"))
 
-        if constants.HEADER_CLEAR_SESSION in http_response.headers:
+        if self._headers_name.clear_session in http_response.headers:
             for prop in get_header_values(
-                http_response.headers, constants.HEADER_CLEAR_SESSION
+                http_response.headers, self._headers_name.clear_session
             ):
                 self._client_session.properties.pop(prop, None)
 
-        if constants.HEADER_SET_SESSION in http_response.headers:
+        if self._headers_name.set_session in http_response.headers:
             for key, value in get_session_property_values(
-                http_response.headers, constants.HEADER_SET_SESSION
+                http_response.headers, self._headers_name.set_session
             ):
                 self._client_session.properties[key] = value
 

trino/constants.py

@@ -26,19 +26,72 @@
 
 URL_STATEMENT_PATH = "/v1/statement"
 
-HEADER_CATALOG = "X-Trino-Catalog"
-HEADER_SCHEMA = "X-Trino-Schema"
-HEADER_SOURCE = "X-Trino-Source"
-HEADER_USER = "X-Trino-User"
-HEADER_CLIENT_INFO = "X-Trino-Client-Info"
-
-HEADER_SESSION = "X-Trino-Session"
-HEADER_SET_SESSION = "X-Trino-Set-Session"
-HEADER_CLEAR_SESSION = "X-Trino-Clear-Session"
-
-HEADER_STARTED_TRANSACTION = "X-Trino-Started-Transaction-Id"
-HEADER_TRANSACTION = "X-Trino-Transaction-Id"
-
-HEADER_PREPARED_STATEMENT = 'X-Trino-Prepared-Statement'
-HEADER_ADDED_PREPARE = 'X-Trino-Added-Prepare'
-HEADER_DEALLOCATED_PREPARE = 'X-Trino-Deallocated-Prepare'
+class Headers:
+
+    @property
+    def base(self):
+        return f'X-{self.NAME}'
+
+    @property
+    def catalog(self):
+        return f'{self.base}-Catalog'
+
+    @property
+    def schema(self):
+        return f'{self.base}-Schema'
+
+    @property
+    def source(self):
+        return f'{self.base}-Source'
+
+    @property
+    def user(self):
+        return f'{self.base}-User'
+
+    @property
+    def client_info(self):
+        return f'{self.base}-Client-Info'
+
+    @property
+    def session(self):
+        return f'{self.base}-Session'
+
+    @property
+    def set_session(self):
+        return f'{self.base}-Set-Session'
+
+    @property
+    def clear_session(self):
+        return f'{self.base}-Clear-Session'
+
+    @property
+    def started_transaction_id(self):
+        return f'{self.base}-Started-Transaction-Id'
+
+    @property
+    def transaction(self):
+        return f'{self.base}-Transaction-Id'
+
+    @property
+    def prepared_statement(self):
+        return f'{self.base}-Prepared-Statement'
+
+    @property
+    def added_prepare(self):
+        return f'{self.base}-Added-Prepare'
+
+    @property
+    def deallocated_prepare(self):
+        return f'{self.base}-Deallocated-Prepare'
+
+
+class PrestoHeaders(Headers):
+    NAME = 'Presto'
+
+
+class TrinoHeaders(Headers):
+    NAME = 'Trino'
+
+
+def get_headers(presto_headers=False):
+    return PrestoHeaders() if presto_headers else TrinoHeaders()

trino/dbapi.py

@@ -76,7 +76,8 @@ def __init__(
         max_attempts=constants.DEFAULT_MAX_ATTEMPTS,
         request_timeout=constants.DEFAULT_REQUEST_TIMEOUT,
         isolation_level=IsolationLevel.AUTOCOMMIT,
-        verify=True
+        verify=True,
+        presto_headers=False
     ):
         self.host = host
         self.port = port
@@ -98,6 +99,7 @@ def __init__(
         self._isolation_level = isolation_level
         self._request = None
         self._transaction = None
+        self._presto_headers = presto_headers
 
     @property
     def isolation_level(self):
@@ -157,6 +159,7 @@ def _create_request(self):
             self.redirect_handler,
             self.max_attempts,
             self.request_timeout,
+            presto_headers=self._presto_headers
         )
 
     def cursor(self):
@@ -189,6 +192,7 @@ def __init__(self, connection, request):
         self.arraysize = 1
         self._iterator = None
         self._query = None
+        self._headers_name = constants.get_headers(presto_headers=self._connection._presto_headers)
 
     def __iter__(self):
         return self._iterator
@@ -267,9 +271,9 @@ def _prepare_statement(self, operation, statement_name):
         # until there are no more results
         for _ in result:
             response_headers = result.response_headers
-
-            if constants.HEADER_ADDED_PREPARE in response_headers:
-                return response_headers[constants.HEADER_ADDED_PREPARE]
+            
+            if self._headers_name.added_prepare in response_headers:
+                return response_headers[self._headers_name.added_prepare]
 
         raise trino.exceptions.FailedToObtainAddedPrepareHeader
 
@@ -344,17 +348,17 @@ def _deallocate_prepare_statement(self, added_prepare_header, statement_name):
         query = trino.client.TrinoQuery(copy.deepcopy(self._request), sql=sql)
         result = query.execute(
             additional_http_headers={
-                constants.HEADER_PREPARED_STATEMENT: added_prepare_header
+                self._headers_name.prepared_statement: added_prepare_header
             }
         )
 
         # Iterate until the 'X-Trino-Deallocated-Prepare' header is found or
         # until there are no more results
         for _ in result:
             response_headers = result.response_headers
-
-            if constants.HEADER_DEALLOCATED_PREPARE in response_headers:
-                return response_headers[constants.HEADER_DEALLOCATED_PREPARE]
+            
+            if self._headers_name.deallocated_prepare in response_headers:
+                return response_headers[self._headers_name.deallocated_prepare]
 
         raise trino.exceptions.FailedToObtainDeallocatedPrepareHeader
 
@@ -382,7 +386,7 @@ def execute(self, operation, params=None):
                 )
                 result = self._query.execute(
                     additional_http_headers={
-                        constants.HEADER_PREPARED_STATEMENT: added_prepare_header
+                        self._headers_name.prepared_statement: added_prepare_header
                     }
                 )
             finally: