diff --git a/tests/unit/sqlalchemy/test_dialect.py b/tests/unit/sqlalchemy/test_dialect.py index d247a536..61922bb1 100644 --- a/tests/unit/sqlalchemy/test_dialect.py +++ b/tests/unit/sqlalchemy/test_dialect.py @@ -271,6 +271,86 @@ def test_trino_connection_basic_auth(): assert cparams['auth']._password == password +@pytest.mark.parametrize( + "userinfo, expected_user, expected_password", + [ + # A literal '+' in the userinfo must stay a '+'. The userinfo is not + # form-encoded, so '+' does not mean space there. + ("user+name:pass+word", "user+name", "pass+word"), + # '%2B' is the percent-encoding of '+'. SQLAlchemy decodes it to '+', and + # the dialect must not decode it a second time into a space. + ("user%2Bname:pass%2Bword", "user+name", "pass+word"), + # A percent-encoded space really is a space. + ("user%20name:pass%20word", "user name", "pass word"), + ], +) +def test_trino_connection_basic_auth_with_plus_in_credentials(userinfo, expected_user, expected_password): + # SQLAlchemy's make_url already decodes the userinfo, so the dialect must not run + # it through unquote_plus again (that would turn a literal '+' into a space). + dialect = TrinoDialect() + url = make_url(f"trino://{userinfo}@host") + _, cparams = dialect.create_connect_args(url) + + assert cparams['user'] == expected_user + assert isinstance(cparams['auth'], BasicAuthentication) + assert cparams['auth']._username == expected_user + assert cparams['auth']._password == expected_password + + +def test_trino_connection_query_params_preserve_plus(): + # Regression for the same double-decode on the other unquote_plus call sites. + # SQLAlchemy already decodes the query values, so a '+' round-tripped through the + # trino URL builder must survive. The catalog and schema are the exception: they + # come from url.database, which SQLAlchemy leaves url-encoded, so they are still + # decoded in the dialect. + dialect = TrinoDialect() + url = make_url(trino_url( + host="host", + user="us+er", + password="pa+ss", + catalog="ca+t", + schema="sc+h", + source="so+urce", + session_properties={"prop": "1+1"}, + http_headers={"x-trino": "a+b"}, + extra_credential=[("cred", "va+lue")], + client_tags=["ta+g"], + roles={"system": "ro+le"}, + )) + _, cparams = dialect.create_connect_args(url) + + assert cparams['user'] == 'us+er' + assert cparams['auth']._password == 'pa+ss' + assert cparams['catalog'] == 'ca+t' + assert cparams['schema'] == 'sc+h' + assert cparams['source'] == 'so+urce' + assert cparams['session_properties'] == {"prop": "1+1"} + assert cparams['http_headers'] == {"x-trino": "a+b"} + assert cparams['extra_credential'] == [("cred", "va+lue")] + assert cparams['client_tags'] == ["ta+g"] + assert cparams['roles'] == {"system": "ro+le"} + + +def test_trino_connection_jwt_auth_preserves_plus(): + # '%2B' in the token is decoded to '+' by SQLAlchemy and must not be decoded again. + dialect = TrinoDialect() + url = make_url("trino://host/?access_token=tok%2Ben") + _, cparams = dialect.create_connect_args(url) + + assert isinstance(cparams['auth'], JWTAuthentication) + assert cparams['auth'].token == 'tok+en' + + +def test_trino_connection_certificate_auth_preserves_plus(): + dialect = TrinoDialect() + url = make_url("trino://host/?cert=ce%2Brt&key=ke%2By") + _, cparams = dialect.create_connect_args(url) + + assert isinstance(cparams['auth'], CertificateAuthentication) + assert cparams['auth']._cert == 'ce+rt' + assert cparams['auth']._key == 'ke+y' + + def test_trino_connection_jwt_auth(): dialect = TrinoDialect() access_token = 'sample-token' diff --git a/trino/sqlalchemy/dialect.py b/trino/sqlalchemy/dialect.py index a3191e03..8407cb31 100644 --- a/trino/sqlalchemy/dialect.py +++ b/trino/sqlalchemy/dialect.py @@ -118,6 +118,8 @@ def create_connect_args(self, url: URL) -> Tuple[Sequence[Any], Mapping[str, Any if url.port: kwargs["port"] = url.port + # url.database (catalog and schema) is the one component SQLAlchemy leaves + # url-encoded, so it still needs decoding here. db_parts = (url.database or "system").split("/") if len(db_parts) == 1: kwargs["catalog"] = unquote_plus(db_parts[0]) @@ -127,50 +129,55 @@ def create_connect_args(self, url: URL) -> Tuple[Sequence[Any], Mapping[str, Any else: raise ValueError(f"Unexpected database format {url.database}") + # SQLAlchemy's make_url already decodes the username, password and query + # values, so they are used as-is. Running them through unquote_plus again + # double-decodes them and turns a literal '+' into a space. For example a + # password supplied as "pass%2Bword", which SQLAlchemy hands over as + # "pass+word", would otherwise become "pass word" and fail authentication. if url.username: - kwargs["user"] = unquote_plus(url.username) + kwargs["user"] = url.username if url.password: if not url.username: raise ValueError("Username is required when specify password in connection URL") - kwargs["auth"] = BasicAuthentication(unquote_plus(url.username), unquote_plus(url.password)) + kwargs["auth"] = BasicAuthentication(url.username, url.password) if "access_token" in url.query: - kwargs["auth"] = JWTAuthentication(unquote_plus(url.query["access_token"])) + kwargs["auth"] = JWTAuthentication(url.query["access_token"]) if "cert" in url.query and "key" in url.query: - kwargs["auth"] = CertificateAuthentication(unquote_plus(url.query['cert']), unquote_plus(url.query['key'])) + kwargs["auth"] = CertificateAuthentication(url.query['cert'], url.query['key']) if "externalAuthentication" in url.query: kwargs["auth"] = OAuth2Authentication() if "source" in url.query: - kwargs["source"] = unquote_plus(url.query["source"]) + kwargs["source"] = url.query["source"] else: kwargs["source"] = "trino-sqlalchemy" if "session_properties" in url.query: - kwargs["session_properties"] = json.loads(unquote_plus(url.query["session_properties"])) + kwargs["session_properties"] = json.loads(url.query["session_properties"]) if "http_headers" in url.query: - kwargs["http_headers"] = json.loads(unquote_plus(url.query["http_headers"])) + kwargs["http_headers"] = json.loads(url.query["http_headers"]) if "extra_credential" in url.query: kwargs["extra_credential"] = [ - tuple(extra_credential) for extra_credential in json.loads(unquote_plus(url.query["extra_credential"])) + tuple(extra_credential) for extra_credential in json.loads(url.query["extra_credential"]) ] if "client_tags" in url.query: - kwargs["client_tags"] = json.loads(unquote_plus(url.query["client_tags"])) + kwargs["client_tags"] = json.loads(url.query["client_tags"]) if "legacy_primitive_types" in url.query: - kwargs["legacy_primitive_types"] = json.loads(unquote_plus(url.query["legacy_primitive_types"])) + kwargs["legacy_primitive_types"] = json.loads(url.query["legacy_primitive_types"]) if "legacy_prepared_statements" in url.query: - kwargs["legacy_prepared_statements"] = json.loads(unquote_plus(url.query["legacy_prepared_statements"])) + kwargs["legacy_prepared_statements"] = json.loads(url.query["legacy_prepared_statements"]) if "verify" in url.query: - kwargs["verify"] = json.loads(unquote_plus(url.query["verify"])) + kwargs["verify"] = json.loads(url.query["verify"]) if "roles" in url.query: kwargs["roles"] = json.loads(url.query["roles"])