Parse the host argument once and validate the endpoint

README.md

@@ -29,7 +29,12 @@ $ pip install trino
 
 Use the DBAPI interface to query Trino:
 
-if `host` is a valid url, the port and http schema will be automatically determined. For example `https://my-trino-server:9999` will assign the `http_schema` property to `https` and port to `9999`.
+`host` accepts a full URL, a hostname on its own, or a `hostname:port`. For
+example `https://my-trino-server:9999` sets `http_scheme` to `https` and `port`
+to `9999`. A scheme or port in `host` that contradicts the `http_scheme` or
+`port` argument is an error. A missing scheme is inferred from the port, and a
+missing port from the scheme. Wrap an IPv6 literal in brackets when it carries
+a port, as in `[::1]:8080`.
 
 ```python
 from trino.dbapi import connect

tests/unit/test_client.py

@@ -52,6 +52,8 @@
 from trino.client import _RetryWithExponentialBackoff
 from trino.client import ClientSession
 from trino.client import CompressedQueryDataDecoderFactory
+from trino.client import port_for_scheme
+from trino.client import scheme_for_port
 from trino.client import TrinoQuery
 from trino.client import TrinoRequest
 from trino.client import TrinoResult
@@ -369,6 +371,40 @@ def test_enabling_https_automatically_when_using_port_443(mock_get_and_post):
     assert parsed_url.scheme == constants.HTTPS
 
 
+def test_trino_request_accepts_a_bracketed_ipv6_host():
+    # Connection always passes an unbracketed literal, so this covers direct construction.
+    req = TrinoRequest(
+        host="[::1]",
+        port=8080,
+        client_session=ClientSession(user="test"),
+    )
+    assert req.statement_url == "http://[::1]:8080/v1/statement"
+
+
+def test_scheme_and_port_inference_are_inverses():
+    assert scheme_for_port(constants.DEFAULT_TLS_PORT) == constants.HTTPS
+    assert scheme_for_port(constants.DEFAULT_PORT) == constants.HTTP
+    assert scheme_for_port(None) == constants.HTTP
+    assert scheme_for_port(9999) == constants.HTTP
+
+    assert port_for_scheme(constants.HTTPS) == constants.DEFAULT_TLS_PORT
+    assert port_for_scheme(constants.HTTP) == constants.DEFAULT_PORT
+
+    for port in (constants.DEFAULT_PORT, constants.DEFAULT_TLS_PORT):
+        assert port_for_scheme(scheme_for_port(port)) == port
+
+
+@pytest.mark.parametrize("http_scheme", ["", "ftp", "gopher"])
+def test_invalid_http_scheme_is_rejected(http_scheme):
+    with pytest.raises(ValueError, match="Invalid http_scheme"):
+        TrinoRequest(
+            host="coordinator",
+            port=constants.DEFAULT_PORT,
+            client_session=ClientSession(user="test"),
+            http_scheme=http_scheme,
+        )
+
+
 def test_https_scheme(mock_get_and_post):
     _, post = mock_get_and_post
 

tests/unit/test_client_spooling.py

@@ -23,6 +23,7 @@
 from trino.client import SpooledSegment
 from trino.client import TrinoQuery
 from trino.client import TrinoRequest
+from trino.dbapi import Connection
 
 
 def _mock_trino_request():
@@ -327,3 +328,35 @@ def fake_get(uri, headers=None, **kwargs):
     segment._send_spooling_request(segment.uri)
 
     assert recorded["headers"]["X-Trino-Spooling-Token"] == "token-abc"
+
+
+def test_send_spooling_request_forwards_custom_headers_for_a_mixed_case_host():
+    # Connection lowercases the host so it matches the hostname urlsplit
+    # parses out of the segment URI. Storing it as given misses the comparison
+    # and drops the custom headers.
+    custom_headers = {"X-Auth-Gateway-Token": "user-token"}
+    connection = Connection("https://MyHost.Domain", user="test", http_headers=custom_headers)
+    request = connection._create_request()
+    segment = SpooledSegment(
+        {
+            "type": "spooled",
+            "uri": "https://MyHost.Domain/v1/spooled/download/seg1",
+            "ackUri": "https://MyHost.Domain/v1/spooled/ack/seg1",
+            "headers": {"X-Trino-Spooling-Token": ["token-abc"]},
+            "metadata": {"segmentSize": "1", "uncompressedSize": "1"},
+        },
+        request,
+        coordinator_host=request._host,
+        custom_headers=dict(request._client_session.headers),
+    )
+
+    recorded = {}
+
+    def fake_get(uri, headers=None, **kwargs):
+        recorded["headers"] = headers
+        return mock.Mock(ok=True)
+
+    segment._request._get = fake_get
+    segment._send_spooling_request(segment.uri)
+
+    assert recorded["headers"]["X-Auth-Gateway-Token"] == "user-token"

tests/unit/test_dbapi.py

@@ -277,70 +277,155 @@ def test_role_is_set_when_specified(mock_client):
     assert passed_role["roles"] == roles
 
 
-def test_hostname_parsing():
-    https_server_with_port = Connection("https://mytrinoserver.domain:9999")
-    assert https_server_with_port.host == "mytrinoserver.domain"
-    assert https_server_with_port.port == 9999
-    assert https_server_with_port.http_scheme == constants.HTTPS
-
-    https_server_without_port = Connection("https://mytrinoserver.domain")
-    assert https_server_without_port.host == "mytrinoserver.domain"
-    assert https_server_without_port.port == constants.DEFAULT_TLS_PORT
-    assert https_server_without_port.http_scheme == constants.HTTPS
-
-    http_server_with_port = Connection("http://mytrinoserver.domain:9999")
-    assert http_server_with_port.host == "mytrinoserver.domain"
-    assert http_server_with_port.port == 9999
-    assert http_server_with_port.http_scheme == constants.HTTP
-
-    http_server_without_port = Connection("http://mytrinoserver.domain")
-    assert http_server_without_port.host == "mytrinoserver.domain"
-    assert http_server_without_port.port == constants.DEFAULT_PORT
-    assert http_server_without_port.http_scheme == constants.HTTP
-
-    http_server_with_path = Connection("http://mytrinoserver.domain/some_path")
-    assert http_server_with_path.host == "mytrinoserver.domain/some_path"
-    assert http_server_with_path.port == constants.DEFAULT_PORT
-    assert http_server_with_path.http_scheme == constants.HTTP
-
-    only_hostname = Connection("mytrinoserver.domain")
-    assert only_hostname.host == "mytrinoserver.domain"
-    assert only_hostname.port == constants.DEFAULT_PORT
-    assert only_hostname.http_scheme == constants.HTTP
-
-    only_hostname_with_path = Connection("mytrinoserver.domain/some_path")
-    assert only_hostname_with_path.host == "mytrinoserver.domain/some_path"
-    assert only_hostname_with_path.port == constants.DEFAULT_PORT
-    assert only_hostname_with_path.http_scheme == constants.HTTP
+@pytest.mark.parametrize(
+    "host, expected_port, expected_scheme",
+    [
+        ("https://mytrinoserver.domain:9999", 9999, constants.HTTPS),
+        ("https://mytrinoserver.domain", constants.DEFAULT_TLS_PORT, constants.HTTPS),
+        ("http://mytrinoserver.domain:9999", 9999, constants.HTTP),
+        ("http://mytrinoserver.domain", constants.DEFAULT_PORT, constants.HTTP),
+        ("mytrinoserver.domain", constants.DEFAULT_PORT, constants.HTTP),
+        ("mytrinoserver.domain:9999", 9999, constants.HTTP),
+        ("mytrinoserver.domain:443", constants.DEFAULT_TLS_PORT, constants.HTTPS),
+    ],
+)
+def test_hostname_parsing(host, expected_port, expected_scheme):
+    connection = Connection(host)
+    assert connection.host == "mytrinoserver.domain"
+    assert connection.port == expected_port
+    assert connection.http_scheme == expected_scheme
+
+
+def test_hostname_is_lowercased():
+    assert Connection("MyTrinoServer.Domain").host == "mytrinoserver.domain"
+    assert Connection("https://MyTrinoServer.Domain").host == "mytrinoserver.domain"
+
+
+@pytest.mark.parametrize(
+    "host, expected_message",
+    [
+        ("http://mytrinoserver.domain/some_path", "a path is not allowed"),
+        ("https://mytrinoserver.domain:9999/some_path", "a path is not allowed"),
+        ("mytrinoserver.domain/some_path", "a path is not allowed"),
+        ("http://mytrinoserver.domain/", "a path is not allowed"),
+        ("mytrinoserver.domain/", "a path is not allowed"),
+        ("user@mytrinoserver.domain", "credentials are not allowed"),
+        ("https://user:password@mytrinoserver.domain", "credentials are not allowed"),
+        ("mytrinoserver.domain?key=value", "a query or fragment is not allowed"),
+        ("mytrinoserver.domain#fragment", "a query or fragment is not allowed"),
+        ("", "the hostname is empty"),
+        ("http://", "the hostname is empty"),
+        ("https://:8080", "the hostname is empty"),
+    ],
+)
+def test_invalid_host_is_rejected(host, expected_message):
+    with pytest.raises(ValueError, match=f"Invalid 'host' argument .*: {expected_message}"):
+        Connection(host)
+
+
+@pytest.mark.parametrize(
+    "host, expected_host, expected_port, expected_url",
+    [
+        ("http://[::1]:8080", "::1", 8080, "http://[::1]:8080/v1/statement"),
+        ("[::1]", "::1", constants.DEFAULT_PORT, "http://[::1]:8080/v1/statement"),
+        ("[::1]:9999", "::1", 9999, "http://[::1]:9999/v1/statement"),
+        ("::1", "::1", constants.DEFAULT_PORT, "http://[::1]:8080/v1/statement"),
+        ("https://[::1]", "::1", constants.DEFAULT_TLS_PORT, "https://[::1]:443/v1/statement"),
+        ("[2001:db8::1]:9999", "2001:db8::1", 9999, "http://[2001:db8::1]:9999/v1/statement"),
+    ],
+)
+def test_ipv6_hostname_parsing(host, expected_host, expected_port, expected_url):
+    connection = Connection(host)
+    # Stored unbracketed, bracketed only in the URL.
+    assert connection.host == expected_host
+    assert connection.port == expected_port
+    assert connection._create_request().statement_url == expected_url
 
 
 def test_description_is_none_when_cursor_is_not_executed():
     connection = Connection("sample_trino_cluster:443")
     with connection.cursor() as cursor:
-        assert hasattr(cursor, 'description')
+        assert cursor.description is None
 
 
 @pytest.mark.parametrize(
-    "host, port, http_scheme_input_argument, http_scheme_set",
+    "host, port, http_scheme, expected_http_scheme, expected_port",
     [
-        # Infer from hostname
-        ("https://mytrinoserver.domain:9999", None, None, constants.HTTPS),
-        ("http://mytrinoserver.domain:9999", None, None, constants.HTTP),
-        # Infer from port
-        ("mytrinoserver.domain", constants.DEFAULT_TLS_PORT, None, constants.HTTPS),
-        ("mytrinoserver.domain", constants.DEFAULT_PORT, None, constants.HTTP),
-        # http_scheme parameter has higher precedence than port parameter
-        ("mytrinoserver.domain", constants.DEFAULT_TLS_PORT, constants.HTTP, constants.HTTP),
-        ("mytrinoserver.domain", constants.DEFAULT_PORT, constants.HTTPS, constants.HTTPS),
-        # Set explicitly by http_scheme parameter
-        ("mytrinoserver.domain", None, constants.HTTPS, constants.HTTPS),
-        # Default
-        ("mytrinoserver.domain", None, None, constants.HTTP),
+        # Decided by a scheme in host, which the http_scheme argument may repeat
+        ("https://mytrinoserver.domain", None, None, constants.HTTPS, constants.DEFAULT_TLS_PORT),
+        ("http://mytrinoserver.domain", None, None, constants.HTTP, constants.DEFAULT_PORT),
+        ("https://mytrinoserver.domain", None, "HTTPS", constants.HTTPS, constants.DEFAULT_TLS_PORT),
+        # Decided by the http_scheme argument
+        ("mytrinoserver.domain", constants.DEFAULT_TLS_PORT, constants.HTTP, constants.HTTP,
+         constants.DEFAULT_TLS_PORT),
+        ("mytrinoserver.domain", constants.DEFAULT_PORT, constants.HTTPS, constants.HTTPS, constants.DEFAULT_PORT),
+        ("mytrinoserver.domain", None, constants.HTTPS, constants.HTTPS, constants.DEFAULT_TLS_PORT),
+        # Decided by the http_scheme argument, which is case-insensitive
+        ("mytrinoserver.domain", None, "HTTPS", constants.HTTPS, constants.DEFAULT_TLS_PORT),
+        ("mytrinoserver.domain", constants.DEFAULT_TLS_PORT, "Http", constants.HTTP, constants.DEFAULT_TLS_PORT),
+        # Decided by the port, whether it arrived in host or in port
+        ("mytrinoserver.domain", constants.DEFAULT_TLS_PORT, None, constants.HTTPS, constants.DEFAULT_TLS_PORT),
+        ("mytrinoserver.domain:443", None, None, constants.HTTPS, constants.DEFAULT_TLS_PORT),
+        ("mytrinoserver.domain", constants.DEFAULT_PORT, None, constants.HTTP, constants.DEFAULT_PORT),
+        # Decided by nothing
+        ("mytrinoserver.domain", None, None, constants.HTTP, constants.DEFAULT_PORT),
     ],
 )
-def test_setting_http_scheme(host, port, http_scheme_input_argument, http_scheme_set):
-    connection = Connection(host, port, http_scheme=http_scheme_input_argument)
-    assert connection.http_scheme == http_scheme_set
+def test_setting_http_scheme(host, port, http_scheme, expected_http_scheme, expected_port):
+    """A scheme in host and the http_scheme argument both win over the port.
+
+    A port the caller did not give is inferred back from the resolved scheme.
+    """
+    connection = Connection(host, port, http_scheme=http_scheme)
+    assert connection.http_scheme == expected_http_scheme
+    assert connection.port == expected_port
+
+
+@pytest.mark.parametrize(
+    "host, http_scheme, expected_message",
+    [
+        ("https://mytrinoserver.domain", constants.HTTP, "contradicts http_scheme"),
+        ("http://mytrinoserver.domain", constants.HTTPS, "contradicts http_scheme"),
+        ("mytrinoserver.domain", "", "Invalid http_scheme"),
+        ("mytrinoserver.domain", "ftp", "Invalid http_scheme"),
+        ("mytrinoserver.domain", "gopher", "Invalid http_scheme"),
+        ("mytrinoserver.domain", "htp", "Invalid http_scheme"),
+        ("ftp://mytrinoserver.domain", None, "Invalid scheme 'ftp' in host"),
+        ("gopher://mytrinoserver.domain", None, "Invalid scheme 'gopher' in host"),
+    ],
+)
+def test_invalid_scheme_is_rejected(host, http_scheme, expected_message):
+    with pytest.raises(ValueError, match=expected_message):
+        Connection(host, user="test", http_scheme=http_scheme)
+
+
+def test_uppercase_http_scheme_still_requires_tls_for_authentication():
+    with pytest.raises(trino.exceptions.TrinoAuthError, match="TLS/SSL is required for authentication"):
+        Connection("mytrinoserver.domain", user="test", auth=BasicAuthentication("test", "pass"), http_scheme="HTTP")
+
+
+@pytest.mark.parametrize(
+    "host, port, expected_message",
+    [
+        # Out of range, whether it arrived in host or in port
+        ("mytrinoserver.domain", -1, "Invalid port"),
+        ("mytrinoserver.domain", 65536, "Invalid port"),
+        ("mytrinoserver.domain", 99999, "Invalid port"),
+        ("mytrinoserver.domain:65536", None, "expected a port in 0-65535"),
+        ("mytrinoserver.domain:-1", None, "expected a port in 0-65535"),
+        ("mytrinoserver.domain:x", None, "expected a port in 0-65535"),
+        # Given twice and disagreeing
+        ("mytrinoserver.domain:9999", 8080, "contradicts port=8080"),
+    ],
+)
+def test_invalid_port_is_rejected(host, port, expected_message):
+    with pytest.raises(ValueError, match=expected_message):
+        Connection(host, port, user="test")
+
+
+def test_port_in_host_agreeing_with_port_argument_is_accepted():
+    connection = Connection("mytrinoserver.domain:8080", 8080, user="test")
+    assert connection.port == 8080
 
 
 @patch("trino.client.CODECS_UNAVAILABLE", {"lz4": "Not installed", "zstd": "Not installed"})

trino/client.py

@@ -379,6 +379,37 @@ def get_roles_values(headers: CaseInsensitiveDict[str], header: str) -> List[Tup
     ]
 
 
+def scheme_for_port(port: Optional[int]) -> str:
+    if port == constants.DEFAULT_TLS_PORT:
+        return constants.HTTPS
+
+    return constants.HTTP
+
+
+def port_for_scheme(scheme: str) -> int:
+    if scheme == constants.HTTPS:
+        return constants.DEFAULT_TLS_PORT
+
+    return constants.DEFAULT_PORT
+
+
+def normalize_http_scheme(http_scheme: str) -> str:
+    scheme = http_scheme.lower()
+    if scheme not in (constants.HTTP, constants.HTTPS):
+        raise ValueError(
+            f"Invalid http_scheme {http_scheme!r}, expected {constants.HTTP!r} or {constants.HTTPS!r}"
+        )
+    return scheme
+
+
+def _authority(host: str, port: int) -> str:
+    # host holds an address only, so a colon means an IPv6 literal. A URL needs
+    # brackets around one.
+    if ":" in host:
+        return f"[{host}]:{port}"
+    return f"{host}:{port}"
+
+
 @dataclass
 class TrinoStatus:
     id: str
@@ -503,17 +534,16 @@ def __init__(
         verify: bool = True,
     ) -> None:
         self._client_session = client_session
-        self._host = host
+        # Store an IPv6 literal unbracketed, to match the parsed hostnames it
+        # is compared against. _authority brackets it for request URLs.
+        self._host = host[1:-1] if host.startswith("[") and host.endswith("]") else host
         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
+            self._http_scheme = scheme_for_port(self._port)
         else:
-            self._http_scheme = http_scheme
+            self._http_scheme = normalize_http_scheme(http_scheme)
 
         if http_session is not None:
             self._http_session = http_session
@@ -661,8 +691,8 @@ def max_attempts(self, value: int) -> None:
         self._head = with_retry(self._http_session.head)
 
     def get_url(self, path: str) -> str:
-        return "{protocol}://{host}:{port}{path}".format(
-            protocol=self._http_scheme, host=self._host, port=self._port, path=path
+        return "{protocol}://{authority}{path}".format(
+            protocol=self._http_scheme, authority=_authority(self._host, self._port), path=path
         )
 
     @property

trino/constants.py

@@ -14,6 +14,8 @@
 
 DEFAULT_PORT = 8080
 DEFAULT_TLS_PORT = 443
+MIN_PORT = 0
+MAX_PORT = 65535
 DEFAULT_SOURCE = "trino-python-client"
 DEFAULT_CATALOG: Optional[str] = None
 DEFAULT_SCHEMA: Optional[str] = None