diff --git a/tests/unit/sqlalchemy/test_dialect.py b/tests/unit/sqlalchemy/test_dialect.py index d247a536..25c4bec2 100644 --- a/tests/unit/sqlalchemy/test_dialect.py +++ b/tests/unit/sqlalchemy/test_dialect.py @@ -4,6 +4,7 @@ from unittest import mock import pytest +from sqlalchemy import exc from sqlalchemy.engine.url import make_url from sqlalchemy.engine.url import URL @@ -314,3 +315,37 @@ def test_trino_connection_oauth2_auth(): _, cparams = dialect.create_connect_args(url) assert isinstance(cparams['auth'], OAuth2Authentication) + + +@pytest.mark.parametrize( + "host", + [ + "https://trino.example.com", # scheme prefix (the common mistake) + "http://trino.example.com", + "trino.example.com:443", # port embedded in host + "trino.example.com/path", # path embedded in host + "user@trino.example.com", # credentials embedded in host + "::1", # unbracketed IPv6 literal + "[::1", # unbalanced bracket + "example.com]", # stray bracket + "https://[::1", # scheme plus malformed IPv6 literal + ], +) +def test_url_rejects_non_hostname_host(host): + with pytest.raises(exc.ArgumentError, match="host must be a hostname only"): + trino_url(host=host, port=443, user="user") + + +@pytest.mark.parametrize( + "host", + [ + "trino.example.com", + "localhost", + "192.168.0.1", + "[::1]", # bracketed IPv6 literal + "host_with.dots-and_underscores", + ], +) +def test_url_accepts_bare_hostname(host): + # A valid host must round-trip through make_url without error. + assert make_url(trino_url(host=host, port=443, user="user")).host is not None diff --git a/trino/sqlalchemy/util.py b/trino/sqlalchemy/util.py index 58af9206..d56ed279 100644 --- a/trino/sqlalchemy/util.py +++ b/trino/sqlalchemy/util.py @@ -6,6 +6,7 @@ from typing import Tuple from typing import Union from urllib.parse import quote_plus +from urllib.parse import urlparse from sqlalchemy import exc @@ -14,6 +15,29 @@ def _rfc_1738_quote(text): return re.sub(r"[:@/]", lambda m: "%%%X" % ord(m.group(0)), text) +def _assert_valid_host(host: str) -> None: + # Parse with a leading "//" so a bare hostname is treated as the authority + # rather than a path. `urlparse(host).scheme` catches an embedded scheme. + try: + parsed = urlparse("//" + host) + invalid = bool( + urlparse(host).scheme + or parsed.port is not None + or parsed.path + or parsed.username + or parsed.password + ) + except ValueError: + # A malformed port or IPv6 literal raises here. + invalid = True + if invalid: + raise exc.ArgumentError( + f"host must be a hostname only, without a scheme, port, path, or credentials. Got {host!r}. " + "For example, use 'example.com' instead of 'https://example.com' and pass the port via the " + "'port' argument." + ) + + def _url( host: str, port: Optional[int] = 8080, @@ -56,6 +80,8 @@ def _url( if not host: raise exc.ArgumentError("host must be specified.") + _assert_valid_host(host) + trino_url += host if not port: