diff --git a/tests/unit/sqlalchemy/test_dialect.py b/tests/unit/sqlalchemy/test_dialect.py index d247a536..6507fa57 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,34 @@ 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 + ], +) +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..f11db09d 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,23 @@ 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. + parsed = urlparse("//" + host) + try: + has_port = parsed.port is not None + except ValueError: + # A malformed port raises here. + has_port = True + if urlparse(host).scheme or has_port or parsed.path or parsed.username or parsed.password: + 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 +74,8 @@ def _url( if not host: raise exc.ArgumentError("host must be specified.") + _assert_valid_host(host) + trino_url += host if not port: