diff --git a/README.md b/README.md index d9c7477a..cf42aabd 100644 --- a/README.md +++ b/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 diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index c5d33dbd..66431231 100644 --- a/tests/unit/test_client.py +++ b/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 diff --git a/tests/unit/test_client_spooling.py b/tests/unit/test_client_spooling.py index 21912d37..198305e6 100644 --- a/tests/unit/test_client_spooling.py +++ b/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" diff --git a/tests/unit/test_dbapi.py b/tests/unit/test_dbapi.py index b5d6b50c..7ad2ce9c 100644 --- a/tests/unit/test_dbapi.py +++ b/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"}) diff --git a/trino/client.py b/trino/client.py index 9d4956bf..bd49cc6f 100644 --- a/trino/client.py +++ b/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 diff --git a/trino/constants.py b/trino/constants.py index 13ceb0ea..4fa5e4be 100644 --- a/trino/constants.py +++ b/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 diff --git a/trino/dbapi.py b/trino/dbapi.py index 2519c86b..75d47230 100644 --- a/trino/dbapi.py +++ b/trino/dbapi.py @@ -18,6 +18,7 @@ decide to convert then to a list of tuples. """ import datetime +import ipaddress import math import uuid from collections import OrderedDict @@ -30,9 +31,11 @@ from typing import Dict from typing import List from typing import NamedTuple +from typing import NoReturn from typing import Optional +from typing import Tuple from typing import Union -from urllib.parse import urlparse +from urllib.parse import urlsplit from zoneinfo import ZoneInfo import trino.client @@ -133,6 +136,94 @@ def connect(*args, **kwargs): _USE_DEFAULT_ENCODING = object() +def _parse_host(host: str) -> Tuple[Optional[str], str, Optional[int]]: + """Split a host argument into its scheme, hostname and port. + + Accepts a full URL, a hostname on its own or a hostname with a port. + Returns None for a missing scheme or port. The caller fills in the + defaults and checks the port range. + + Does not check the hostname itself. An unresolvable name fails later, when + requests connects. + """ + def fail(reason: str) -> NoReturn: + raise ValueError(f"Invalid 'host' argument {host!r}: {reason}.") + + # Without "//" urlsplit reads a scheme-less argument as a path and returns + # hostname=None. The prefix forces the authority form. urlsplit does not + # validate the hostname it hands back. + parts = urlsplit(host if "://" in host else "//" + host, allow_fragments=True) + try: + hostname, port = parts.hostname, parts.port + except ValueError: + # urlsplit cannot read an unbracketed IPv6 literal. URLs require + # the brackets because ::1 is otherwise ambiguous with host:port. + try: + ipaddress.ip_address(host) + return None, host, None + except ValueError: + pass + # urlsplit raises the same error for an out-of-range port and a + # non-numeric one. Report the range either way. + fail(f"expected a port in {constants.MIN_PORT}-{constants.MAX_PORT}") + + if parts.path: + fail("a path is not allowed") + if parts.username is not None or parts.password is not None: + fail("credentials are not allowed, pass the 'user' and 'auth' arguments instead") + if parts.query or parts.fragment: + fail("a query or fragment is not allowed") + if not hostname: + fail("the hostname is empty") + return (parts.scheme or None), hostname, port + + +def _resolve_endpoint( + host: str, port: Optional[int], http_scheme: Optional[str] +) -> Tuple[str, str, int]: + """Resolve the scheme, hostname and port to connect to. + + The scheme and the port can each arrive in the host argument or in their + own argument. Passing both raises unless they agree. A missing port is + inferred from the scheme and a missing scheme is inferred from the port. + """ + host_scheme, hostname, host_port = _parse_host(host) + + if host_port is not None and port is not None and host_port != port: + raise ValueError( + f"The port in host {host!r} contradicts port={port!r}. " + "Drop one of the two, or make them agree." + ) + + given_port = host_port if host_port is not None else port + if given_port is not None and not constants.MIN_PORT <= given_port <= constants.MAX_PORT: + raise ValueError( + f"Invalid port {given_port!r}, expected {constants.MIN_PORT}-{constants.MAX_PORT}." + ) + + if host_scheme: + try: + scheme = trino.client.normalize_http_scheme(host_scheme) + except ValueError: + raise ValueError( + f"Invalid scheme {host_scheme!r} in host {host!r}, " + f"expected {constants.HTTP!r} or {constants.HTTPS!r}." + ) from None + if http_scheme is not None and trino.client.normalize_http_scheme(http_scheme) != scheme: + raise ValueError( + f"The scheme in host {host!r} contradicts http_scheme={http_scheme!r}. " + "Drop one of the two, or make them agree." + ) + elif http_scheme is not None: + scheme = trino.client.normalize_http_scheme(http_scheme) + else: + scheme = trino.client.scheme_for_port(given_port) + + if given_port is None: + given_port = trino.client.port_for_scheme(scheme) + return scheme, hostname, given_port + + class Connection: """Trino supports transactions and the ability to either commit or rollback a sequence of SQL statements. A single query i.e. the execution of a SQL @@ -168,8 +259,7 @@ def __init__( heartbeat_interval: Optional[float] = constants.DEFAULT_HEARTBEAT_INTERVAL, allow_insecure_auth: bool = False, ): - # Automatically assign http_schema, port based on hostname - parsed_host = urlparse(host, allow_fragments=False) + self.http_scheme, self.host, self.port = _resolve_endpoint(host, port, http_scheme) if encoding is _USE_DEFAULT_ENCODING: encoding = [ @@ -177,8 +267,6 @@ def __init__( for enc in trino.client.ENCODINGS if (enc.split("+")[1] if "+" in enc else None) not in trino.client.CODECS_UNAVAILABLE ] - - self.host = host if parsed_host.hostname is None else parsed_host.hostname + parsed_host.path self.user = user self.source = source self.catalog = catalog @@ -207,38 +295,16 @@ def __init__( self._http_session = http_session self.http_headers = http_headers - # Set http_scheme - if parsed_host.scheme: - self.http_scheme = parsed_host.scheme - elif http_scheme: - self.http_scheme = http_scheme - elif port == constants.DEFAULT_TLS_PORT: - self.http_scheme = constants.HTTPS - elif port == constants.DEFAULT_PORT: - self.http_scheme = constants.HTTP - else: - self.http_scheme = constants.HTTP - if auth is not None and self.http_scheme == constants.HTTP and not allow_insecure_auth: raise trino.exceptions.TrinoAuthError( "TLS/SSL is required for authentication. " - "To use HTTPS, specify 'https://' in the host URL (which takes precedence " - "over http_scheme), or, if the host URL has no scheme, pass http_scheme='https'. " + "To use HTTPS, specify 'https://' in the host URL or pass http_scheme='https'. " "If your connection is encrypted below the application layer (for example behind an mTLS " "service mesh sidecar), pass allow_insecure_auth=True and ensure " "http-server.authentication.allow-insecure-over-http=true is set on the coordinator if it " "has HTTPS enabled." ) - # Infer connection port: `hostname` takes precedence over explicit `port` argument - # If none is given, use default based on HTTP protocol - default_port = constants.DEFAULT_TLS_PORT if self.http_scheme == constants.HTTPS else constants.DEFAULT_PORT - self.port = ( - parsed_host.port if parsed_host.port is not None - else port if port is not None - else default_port - ) - self.auth = auth self.extra_credential = extra_credential self.max_attempts = max_attempts diff --git a/trino/sqlalchemy/util.py b/trino/sqlalchemy/util.py index 372f5a44..0de01f78 100644 --- a/trino/sqlalchemy/util.py +++ b/trino/sqlalchemy/util.py @@ -6,29 +6,28 @@ from typing import Tuple from typing import Union from urllib.parse import quote_plus -from urllib.parse import urlparse from sqlalchemy import exc +from trino import constants +from trino.dbapi import _parse_host + 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. + """Require a hostname on its own and reject a scheme or port. + + Stricter than `trino.dbapi._parse_host`. The caller builds a URL that + already has its own scheme and port so a scheme or port here has nowhere + to go. An IPv6 literal must be bracketed as that URL needs the brackets. + """ try: - parsed = urlparse("//" + host) - invalid = bool( - urlparse(host).scheme - or parsed.port is not None - or parsed.path - or parsed.username - or parsed.password - ) + scheme, _, port = _parse_host(host) + invalid = scheme is not None or port is not None or (":" in host and not host.startswith("[")) except ValueError: - # A malformed port or IPv6 literal raises here. invalid = True if invalid: raise exc.ArgumentError( @@ -40,7 +39,7 @@ def _assert_valid_host(host: str) -> None: def _url( host: str, - port: Optional[int] = 8080, + port: Optional[int] = constants.DEFAULT_PORT, user: Optional[str] = None, password: Optional[str] = None, catalog: Optional[str] = None,