Update client to default protocol to HTTPS if port is 443

If a user passes in a port that is 443, but they do not pass in a protocol, we can assume that the user is using the HTTPS protocol. This change makes the Python client consistent with the behaviour of the JDBC driver.
dvecerdea · Oct 7, 2021 · 7f4473c · 7f4473c
1 parent a0fb0c3
commit 7f4473c
Show file tree

Hide file tree

Showing 3 changed files with 64 additions and 3 deletions.
diff --git a/tests/test_client.py b/tests/test_client.py
@@ -15,6 +15,7 @@
 import requests
 import time
 from unittest import mock
+from urllib.parse import urlparse
 
 
 from requests_kerberos.exceptions import KerberosExchangeError
@@ -402,6 +403,58 @@ def test_request_invalid_http_headers():
     assert str(value_error.value).startswith("cannot override reserved HTTP header")
 
 
+def test_enabling_https_automatically_when_using_port_443(monkeypatch):
+    post_recorder = ArgumentsRecorder()
+    monkeypatch.setattr(TrinoRequest.http.Session, "post", post_recorder)
+
+    req = TrinoRequest(
+        host="coordinator",
+        port=constants.DEFAULT_TLS_PORT,
+        user="test",
+    )
+
+    req.post("SELECT 1")
+    parsed_url = urlparse(post_recorder.args[0])
+
+    assert parsed_url.scheme == constants.HTTPS
+
+
+def test_https_scheme(monkeypatch):
+    post_recorder = ArgumentsRecorder()
+    monkeypatch.setattr(TrinoRequest.http.Session, "post", post_recorder)
+
+    req = TrinoRequest(
+        host="coordinator",
+        port=constants.DEFAULT_TLS_PORT,
+        user="test",
+        http_scheme=constants.HTTPS,
+    )
+
+    req.post("SELECT 1")
+    parsed_url = urlparse(post_recorder.args[0])
+
+    assert parsed_url.scheme == constants.HTTPS
+    assert parsed_url.port == constants.DEFAULT_TLS_PORT
+
+
+def test_http_scheme_with_port(monkeypatch):
+    post_recorder = ArgumentsRecorder()
+    monkeypatch.setattr(TrinoRequest.http.Session, "post", post_recorder)
+
+    req = TrinoRequest(
+        host="coordinator",
+        port=constants.DEFAULT_TLS_PORT,
+        user="test",
+        http_scheme=constants.HTTP,
+    )
+
+    req.post("SELECT 1")
+    parsed_url = urlparse(post_recorder.args[0])
+
+    assert parsed_url.scheme == constants.HTTP
+    assert parsed_url.port == constants.DEFAULT_TLS_PORT
+
+
 def test_request_timeout():
     timeout = 0.1
     http_scheme = "http"

diff --git a/trino/client.py b/trino/client.py
@@ -200,7 +200,7 @@ def __init__(
         http_session: Any = None,
         http_headers: Optional[Dict[str, str]] = None,
         transaction_id: Optional[str] = NO_TRANSACTION,
-        http_scheme: str = constants.HTTP,
+        http_scheme: str = None,
         auth: Optional[Any] = constants.DEFAULT_AUTH,
         redirect_handler: Any = None,
         max_attempts: int = MAX_ATTEMPTS,
@@ -222,6 +222,14 @@ def __init__(
         self._port = port
         self._next_uri: Optional[str] = None
 
+        if http_scheme is None:
+            if self._port == constants.DEFAULT_TLS_PORT:
+                self._http_scheme = constants.HTTPS
+            else:
+                self._http_scheme = constants.HTTP
+        else:
+            self._http_scheme = http_scheme
+
         if http_session is not None:
             self._http_session = http_session
         else:
@@ -231,7 +239,7 @@ def __init__(
         self._exceptions = self.HTTP_EXCEPTIONS
         self._auth = auth
         if self._auth:
-            if http_scheme == constants.HTTP:
+            if self._http_scheme == constants.HTTP:
                 raise ValueError("cannot use authentication with HTTP")
             self._auth.set_http_session(self._http_session)
             self._exceptions += self._auth.get_exceptions()
@@ -240,7 +248,6 @@ def __init__(
         self._request_timeout = request_timeout
         self._handle_retry = handle_retry
         self.max_attempts = max_attempts
-        self._http_scheme = http_scheme
 
     @property
     def transaction_id(self):

diff --git a/trino/constants.py b/trino/constants.py
@@ -14,6 +14,7 @@
 
 
 DEFAULT_PORT = 8080
+DEFAULT_TLS_PORT = 443
 DEFAULT_SOURCE = "trino-python-client"
 DEFAULT_CATALOG: Optional[str] = None
 DEFAULT_SCHEMA: Optional[str] = None