Add parameter extra_credential

Signed-off-by: Ke Zhu <[email protected]>

Author: Ke Zhu

README.md

@@ -125,6 +125,24 @@ The transaction is created when the first SQL statement is executed.
 exits the *with* context and the queries succeed, otherwise
 `trino.dbapi.Connection.rollback()` will be called.
 
+# Extra Credential
+
+Send [`extra credentials`](https://trino.io/docs/current/develop/client-protocol.html#client-request-headers):
+
+```python
+import trino
+conn = trino.dbapi.connect(
+    host='localhost',
+    port=443,
+    user='the-user',
+    extra_credential=[('a.username', 'bar'), ('a.password', 'foo')],
+)
+
+cur = conn.cursor()
+cur.execute('SELECT * FROM system.runtime.nodes')
+rows = cur.fetchall()
+```
+
 # Development
 
 ## Getting Started With Development

tests/unit/test_client.py

@@ -713,6 +713,64 @@ def test_trino_connection_error(monkeypatch, error_code, error_type, error_messa
     assert error_message in str(error)
 
 
+def test_extra_credential(mock_get_and_post):
+    _, post = mock_get_and_post
+
+    req = TrinoRequest(
+        host="coordinator",
+        port=constants.DEFAULT_TLS_PORT,
+        user="test",
+        extra_credential=[("a.username", "foo"), ("b.password", "bar")],
+    )
+
+    req.post("SELECT 1")
+    _, post_kwargs = post.call_args
+    headers = post_kwargs["headers"]
+    assert constants.HEADER_EXTRA_CREDENTIAL in headers
+    assert headers[constants.HEADER_EXTRA_CREDENTIAL] == "a.username=foo, b.password=bar"
+
+
+def test_extra_credential_key_with_illegal_chars():
+    with pytest.raises(ValueError) as e_info:
+        TrinoRequest(
+            host="coordinator",
+            port=constants.DEFAULT_TLS_PORT,
+            user="test",
+            extra_credential=[("a=b", "")],
+        )
+
+    assert str(e_info.value) == "whitespace or '=' are disallowed in extra credential 'a=b'"
+
+
+def test_extra_credential_key_non_ascii():
+    with pytest.raises(ValueError) as e_info:
+        TrinoRequest(
+            host="coordinator",
+            port=constants.DEFAULT_TLS_PORT,
+            user="test",
+            extra_credential=[("的", "")],
+        )
+
+    assert str(e_info.value) == "only ASCII characters are allowed in extra credential '的'"
+
+
+def test_extra_credential_value_encoding(mock_get_and_post):
+    _, post = mock_get_and_post
+
+    req = TrinoRequest(
+        host="coordinator",
+        port=constants.DEFAULT_TLS_PORT,
+        user="test",
+        extra_credential=[("foo", "bar 的")],
+    )
+
+    req.post("SELECT 1")
+    _, post_kwargs = post.call_args
+    headers = post_kwargs["headers"]
+    assert constants.HEADER_EXTRA_CREDENTIAL in headers
+    assert headers[constants.HEADER_EXTRA_CREDENTIAL] == "foo=bar+%E7%9A%84"
+
+
 class RetryRecorder(object):
     def __init__(self, error=None, result=None):
         self.__name__ = "RetryRecorder"

trino/client.py

@@ -35,6 +35,7 @@
 
 import copy
 import os
+import re
 from typing import Any, Dict, List, Optional, Tuple, Union
 import urllib.parse
 
@@ -55,6 +56,8 @@
 else:
     PROXIES = {}
 
+_HEADER_EXTRA_CREDENTIAL_KEY_REGEX = re.compile(r'^\S[^\s=]*$')
+
 
 class ClientSession(object):
     def __init__(
@@ -66,6 +69,7 @@ def __init__(
         properties=None,
         headers=None,
         transaction_id=None,
+        extra_credential=None,
     ):
         self.catalog = catalog
         self.schema = schema
@@ -76,6 +80,7 @@ def __init__(
         self._properties = properties
         self._headers = headers or {}
         self.transaction_id = transaction_id
+        self.extra_credential = extra_credential
 
     @property
     def properties(self):
@@ -153,6 +158,8 @@ class TrinoRequest(object):
     :param http_scheme: "http" or "https"
     :param auth: class that manages user authentication. ``None`` means no
                  authentication.
+    :param extra_credential: extra credentials. as list of ``(key, value)``
+                             tuples.
     :max_attempts: maximum number of attempts when sending HTTP requests. An
                    attempt is an HTTP request. 5 attempts means 4 retries.
     :request_timeout: How long (in seconds) to wait for the server to send
@@ -206,6 +213,7 @@ def __init__(
         transaction_id: Optional[str] = NO_TRANSACTION,
         http_scheme: str = None,
         auth: Optional[Any] = constants.DEFAULT_AUTH,
+        extra_credential: Optional[List[Tuple[str, str]]] = None,
         redirect_handler: Any = None,
         max_attempts: int = MAX_ATTEMPTS,
         request_timeout: Union[float, Tuple[float, float]] = constants.DEFAULT_REQUEST_TIMEOUT,
@@ -220,6 +228,7 @@ def __init__(
             session_properties,
             http_headers,
             transaction_id,
+            extra_credential,
         )
 
         self._host = host
@@ -285,6 +294,19 @@ def http_headers(self) -> Dict[str, str]:
         transaction_id = self._client_session.transaction_id
         headers[constants.HEADER_TRANSACTION] = transaction_id
 
+        if self._client_session.extra_credential is not None and \
+                len(self._client_session.extra_credential) > 0:
+
+            for tup in self._client_session.extra_credential:
+                self._verify_extra_credential(tup)
+
+            # HTTP 1.1 section 4.2 combine multiple extra credentials into a
+            # comma-separated value
+            # extra credential value is encoded per spec (application/x-www-form-urlencoded MIME format)
+            headers[constants.HEADER_EXTRA_CREDENTIAL] = \
+                ", ".join(
+                    [f"{tup[0]}={urllib.parse.quote_plus(tup[1])}" for tup in self._client_session.extra_credential])
+
         return headers
 
     @property
@@ -427,6 +449,20 @@ def process(self, http_response) -> TrinoStatus:
             columns=response.get("columns"),
         )
 
+    def _verify_extra_credential(self, header):
+        """
+        Verifies that key has ASCII only and non-whitespace characters.
+        """
+        key = header[0]
+
+        if not _HEADER_EXTRA_CREDENTIAL_KEY_REGEX.match(key):
+            raise ValueError(f"whitespace or '=' are disallowed in extra credential '{key}'")
+
+        try:
+            key.encode().decode('ascii')
+        except UnicodeDecodeError:
+            raise ValueError(f"only ASCII characters are allowed in extra credential '{key}'")
+
 
 class TrinoResult(object):
     """

trino/constants.py

@@ -32,6 +32,7 @@
 HEADER_SOURCE = "X-Trino-Source"
 HEADER_USER = "X-Trino-User"
 HEADER_CLIENT_INFO = "X-Trino-Client-Info"
+HEADER_EXTRA_CREDENTIAL = "X-Trino-Extra-Credential"
 
 HEADER_SESSION = "X-Trino-Session"
 HEADER_SET_SESSION = "X-Trino-Set-Session"

trino/dbapi.py

@@ -102,6 +102,7 @@ def __init__(
         http_headers=None,
         http_scheme=constants.HTTP,
         auth=constants.DEFAULT_AUTH,
+        extra_credential=None,
         redirect_handler=None,
         max_attempts=constants.DEFAULT_MAX_ATTEMPTS,
         request_timeout=constants.DEFAULT_REQUEST_TIMEOUT,
@@ -125,6 +126,7 @@ def __init__(
         self.http_headers = http_headers
         self.http_scheme = http_scheme
         self.auth = auth
+        self.extra_credential = extra_credential
         self.redirect_handler = redirect_handler
         self.max_attempts = max_attempts
         self.request_timeout = request_timeout
@@ -188,6 +190,7 @@ def _create_request(self):
             NO_TRANSACTION,
             self.http_scheme,
             self.auth,
+            self.extra_credential,
             self.redirect_handler,
             self.max_attempts,
             self.request_timeout,