From dee8b3922373573714ed6cb754f6a2b940b867fe Mon Sep 17 00:00:00 2001 From: "Mateusz \"Serafin\" Gajewski" Date: Mon, 31 Aug 2026 16:37:51 +0200 Subject: [PATCH] Replace requests with httpx2 and add an asynchronous client Migrate the HTTP stack from requests to httpx2, with HTTP/2 negotiated via ALPN on TLS connections and transparent HTTP/1.1 fallback. - Extract the transport-agnostic protocol logic (headers, response processing, retry decisions, session state, segment decoding) into trino/_protocol.py, shared by the sync and async clients. - Rework trino.auth: Authentication now contributes get_http_auth() and get_client_arguments() since verify/cert/trust_env are constructor-only in httpx; OAuth2 is an httpx2.Auth generator flow working for both sync and async clients; Kerberos/GSSAPI use an in-repo SPNEGO flow on python-gssapi (trino/_spnego.py) instead of requests_kerberos and requests_gssapi. - Add trino.aio, an asynchronous DBAPI-like interface on httpx2.AsyncClient with spooled protocol support (async segment downloads, heartbeat task, acknowledgment tasks drained on close). - Replace httpretty with a hand-written recording MockTrinoServer over httpx2.MockTransport serving sync and async unit tests alike. --- CHANGES.md | 35 + CLAUDE.md | 11 +- README.md | 51 +- setup.cfg | 6 +- setup.py | 22 +- tests/integration/test_dbapi_integration.py | 19 +- tests/unit/aio/__init__.py | 0 tests/unit/aio/test_async_client.py | 219 ++++ tests/unit/aio/test_async_dbapi.py | 234 +++++ tests/unit/conftest.py | 27 +- tests/unit/mock_http.py | 125 +++ tests/unit/oauth_test_utils.py | 123 ++- tests/unit/test_auth_gssapi.py | 18 +- tests/unit/test_client.py | 795 ++++++-------- tests/unit/test_client_spooling.py | 10 +- tests/unit/test_dbapi.py | 189 ++-- trino/_protocol.py | 1013 ++++++++++++++++++ trino/_spnego.py | 199 ++++ trino/aio/__init__.py | 33 + trino/aio/client.py | 709 +++++++++++++ trino/aio/dbapi.py | 463 ++++++++ trino/auth.py | 331 +++--- trino/client.py | 1042 +++---------------- trino/dbapi.py | 86 +- trino/exceptions.py | 9 + trino/transaction.py | 2 +- 26 files changed, 4001 insertions(+), 1770 deletions(-) create mode 100644 tests/unit/aio/__init__.py create mode 100644 tests/unit/aio/test_async_client.py create mode 100644 tests/unit/aio/test_async_dbapi.py create mode 100644 tests/unit/mock_http.py create mode 100644 trino/_protocol.py create mode 100644 trino/_spnego.py create mode 100644 trino/aio/__init__.py create mode 100644 trino/aio/client.py create mode 100644 trino/aio/dbapi.py diff --git a/CHANGES.md b/CHANGES.md index eb71eca2..a15c635b 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -6,6 +6,41 @@ list](https://github.com/trinodb/trino-python-client/tags), the [README](https://github.com/trinodb/trino-python-client/blob/master/README.md) and the [PyPI page](https://pypi.org/project/trino/). +## Release 0.340.0 + +* **Breaking:** Replace the `requests` HTTP stack with + [httpx2](https://github.com/pydantic/httpx2). HTTP/2 is negotiated via ALPN + on TLS connections when the server supports it, with transparent HTTP/1.1 + fallback. Note the following interface changes: + * The `http_session` parameter of `trino.dbapi.Connection` and + `trino.client.TrinoRequest` now takes an `httpx2.Client` instead of a + `requests.Session`. TLS options (`verify`, `cert`) and `trust_env` can + only be configured when the client is constructed, so combining a custom + `http_session` with an authentication that needs them (for example + `CertificateAuthentication`) now raises `TrinoConnectionError` instead of + silently mutating the session. + * The `trino.auth.Authentication` interface changed: `set_http_session` was + replaced by `get_http_auth()` (returning an `httpx2.Auth`) and + `get_client_arguments()` (returning `verify`/`cert`/`trust_env` + constructor arguments). Third-party implementations must be updated. + * Transport errors are now `httpx2` exceptions (`httpx2.HTTPError` and + subclasses) instead of `requests.exceptions.*`. + * `trino.client.PROXIES` uses the httpx mounts shape + (`{"all://": "socks5://host:port"}`) and is applied when the HTTP client + is constructed, not per request. SOCKS support requires the new + `trino[socks]` extra. +* **Breaking:** Kerberos and GSSAPI authentication implement the SPNEGO token + exchange in the client on top of `python-gssapi` instead of + `requests_kerberos`/`requests_gssapi`. The public constructor signatures of + `KerberosAuthentication` and `GSSAPIAuthentication` are unchanged and both + the `trino[kerberos]` and `trino[gssapi]` extras now install the same + dependencies. A GSSAPI implementation such as MIT Kerberos is required; + Windows SSPI is not supported. +* Add `trino.aio`, an asynchronous DBAPI-like interface built on + `httpx2.AsyncClient` with support for all authentication mechanisms and the + spooled protocol. Transactions are not supported by the asynchronous client + yet. + ## Release 0.339.0 * Add support for returning column comments from SQLAlchemy `get_columns`. diff --git a/CLAUDE.md b/CLAUDE.md index 6f17a41d..7729248f 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -33,7 +33,10 @@ tox ### Package layout (`trino/`) -- **`client.py`** - Core HTTP protocol implementation. `TrinoRequest` manages HTTP requests to the coordinator, `TrinoQuery` manages query lifecycle (submit -> poll -> fetch results), `ClientSession` holds connection state (headers, session properties, transaction ID). Implements the protocol described at https://github.com/trinodb/trino/wiki/HTTP-Protocol. +- **`client.py`** - Core synchronous HTTP protocol implementation on `httpx2.Client`. `TrinoRequest` manages HTTP requests to the coordinator, `TrinoQuery` manages query lifecycle (submit -> poll -> fetch results), `ClientSession` holds connection state (headers, session properties, transaction ID). Implements the protocol described at https://github.com/trinodb/trino/wiki/HTTP-Protocol. +- **`_protocol.py`** - Transport-agnostic protocol logic shared by the sync and async clients: header construction/parsing, response processing, retry decisions, session state, segment decoding. +- **`aio/`** - Asynchronous client on `httpx2.AsyncClient`: `AsyncTrinoRequest`/`AsyncTrinoQuery` in `aio/client.py` and a DBAPI-like async `Connection`/`Cursor` in `aio/dbapi.py` (entry point `trino.aio.connect()`). Not PEP 249; transactions unsupported. +- **`_spnego.py`** - In-repo SPNEGO (Kerberos/GSSAPI) `httpx2.Auth` flow on python-gssapi, backing `KerberosAuthentication` and `GSSAPIAuthentication`. - **`dbapi.py`** - PEP 249 DBAPI 2.0 interface. `Connection` and `Cursor` classes that wrap `client.py`. Entry point is `trino.dbapi.connect()`. - **`sqlalchemy/`** - SQLAlchemy dialect (`TrinoDialect`), compiler, and type mapping. Registered as `trino://` via `sqlalchemy.dialects` entry point. Compatible with SQLAlchemy 1.3.x, 1.4.x, and 2.0.x. - **`auth.py`** - Authentication implementations (Basic, JWT, OAuth2, Kerberos/GSSAPI, Certificate). @@ -46,14 +49,14 @@ tox ### Tests (`tests/`) -- `tests/unit/` - No external dependencies needed. Uses `httpretty` for HTTP mocking. +- `tests/unit/` - No external dependencies needed. HTTP is mocked with the hand-written recording router in `tests/unit/mock_http.py` (`MockTrinoServer`, served through `httpx2.MockTransport`; works for both sync and async clients). Async tests live in `tests/unit/aio/` and run under pytest-asyncio (`asyncio_mode = auto`). - `tests/integration/` - Requires Docker. Tests automatically pull `trinodb/trino` image and start a container (or reuse one on port 8080). - Env vars: `TRINO_VERSION` (image tag, default `latest`), `TRINO_RUNNING_HOST`/`TRINO_RUNNING_PORT` (use an existing server instead of starting one). - When `TRINO_VERSION` is `latest` or >= 466, a LocalStack container with an S3 `spooling` bucket is also started to test the spooled client protocol; older versions use the `etc/*-pre-466*` configs. ### Key dependencies -- `requests` - HTTP transport +- `httpx2` - HTTP transport (sync and async, HTTP/2 via ALPN with HTTP/1.1 fallback) - `orjson` (CPython) / `json` (PyPy) - JSON parsing, selected at import time in `client.py` - `lz4`, `zstandard` - Response decompression - `python-dateutil`, `pytz`, `tzlocal` - Timezone/datetime handling @@ -63,5 +66,5 @@ tox - **Max line length**: 120 characters (flake8) - **Import ordering**: managed by `reorder-python-imports` (Python 3.9+ style) - **Type checking**: mypy with strict settings, though `tests/*`, `trino/client`, `trino/dbapi`, and `trino/sqlalchemy.*` have `ignore_errors = true` -- **No mocking libraries**: write mocks by hand instead of using `unittest.mock` or similar. The project uses `httpretty` for HTTP-level stubbing only. +- **No mocking libraries**: write mocks by hand instead of using `unittest.mock` or similar. HTTP-level stubbing goes through `tests/unit/mock_http.py` (`httpx2.MockTransport`-based). - **Pre-commit hooks**: flake8, mypy, reorder-python-imports, trailing whitespace, EOF newlines, YAML syntax, case-conflict checks diff --git a/README.md b/README.md index d9c7477a..e64410f9 100644 --- a/README.md +++ b/README.md @@ -78,6 +78,38 @@ the query's current stats dictionary (the same dictionary returned by `Cursor.stats`), so mutating it has no effect on the client. Any exception raised by the callback propagates to the caller of `execute()`/`fetch()`. +### Asynchronous client + +`trino.aio` provides an asynchronous, DBAPI-like interface built on +`httpx2.AsyncClient`. It mirrors `trino.dbapi` with coroutine methods and +accepts the same connection arguments (it is not a PEP 249 implementation, as +PEP 249 is a synchronous specification): + +```python +import asyncio +import trino.aio + + +async def main(): + async with trino.aio.connect( + host="", + port=, + user="", + catalog="", + schema="", + ) as conn: + cur = conn.cursor() + await cur.execute("SELECT * FROM system.runtime.nodes") + rows = await cur.fetchall() + # or: async for row in cur: ... + + +asyncio.run(main()) +``` + +All authentication mechanisms and the spooled protocol are supported. +Transactions are not supported by the asynchronous client yet. + ### SQLAlchemy **Prerequisite** @@ -357,8 +389,10 @@ Make sure that the GSSAPI support is installed using `pip install trino[gssapi]` The `GSSAPIAuthentication` class can be used to connect to a Trino cluster configured with the [`Kerberos` authentication type](https://trino.io/docs/current/security/kerberos.html): -It follows the interface for `KerberosAuthentication`, but is using -[requests-gssapi](https://github.com/pythongssapi/requests-gssapi), instead of [requests-kerberos](https://github.com/requests/requests-kerberos) under the hood. +It follows the interface for `KerberosAuthentication`. Both classes implement +the SPNEGO token exchange directly on top of +[python-gssapi](https://github.com/pythongssapi/python-gssapi), which requires +a GSSAPI implementation such as MIT Kerberos (Windows SSPI is not supported). - DBAPI @@ -479,7 +513,7 @@ conn = connect( ### Self-signed certificates To use self-signed certificates, specify a path to the certificate in `verify` parameter. -More details can be found in [the Python requests library documentation](https://requests.readthedocs.io/en/latest/user/advanced/#ssl-cert-verification). +More details can be found in [the httpx SSL documentation](https://www.python-httpx.org/advanced/ssl/). ```python from trino.dbapi import connect @@ -568,16 +602,15 @@ The transaction is created when the first SQL statement is executed. exits the *with* context and the queries succeed, otherwise `trino.dbapi.Connection.rollback()` will be called. -## Custom requests Session +## Custom HTTP client -You can create a custom [requests.Session object](https://requests.readthedocs.io/en/latest/user/advanced/#session-objects) and pass it to the `http_session` parameter. This can be used for things like setting additional HTTP headers, client certificates, etc. +You can create a custom [httpx2.Client](https://www.python-httpx.org/advanced/clients/) and pass it to the `http_session` parameter. This can be used for things like setting additional HTTP headers, client certificates, proxies, etc. Note that TLS options (`verify`, `cert`) and `trust_env` can only be set when the client is constructed. ```python -import requests +import httpx2 from trino.dbapi import connect -s = requests.Session() -s.cert = '/path/client.cert' +s = httpx2.Client(cert='/path/client.cert', http2=True) conn = connect( http_session=s, @@ -585,6 +618,8 @@ conn = connect( ) ``` +The asynchronous client accepts an `httpx2.AsyncClient` the same way. + ## Legacy Primitive types By default, the client will convert the results of the query to the diff --git a/setup.cfg b/setup.cfg index c7cc318a..18f27610 100644 --- a/setup.cfg +++ b/setup.cfg @@ -20,5 +20,9 @@ no_implicit_optional = true warn_unused_ignores = true disable_error_code = import-untyped -[mypy-tests.*,trino.client,trino.dbapi,trino.sqlalchemy.*] +[mypy-tests.*,trino.client,trino.dbapi,trino.aio.*,trino.sqlalchemy.*] ignore_errors = true + +[tool:pytest] +asyncio_mode = auto +asyncio_default_fixture_loop_scope = function diff --git a/setup.py b/setup.py index 76d67490..c5c4a08e 100755 --- a/setup.py +++ b/setup.py @@ -26,22 +26,22 @@ with open(os.path.join(here, "README.md"), "r", "utf-8") as f: readme = f.read() -kerberos_require = ["requests_kerberos"] -gssapi_require = ["" - "requests_gssapi", - # PyPy compatibility issue https://github.com/jborean93/pykrb5/issues/49 - "krb5 == 0.5.1"] +# Kerberos and GSSAPI authentication use the in-repo SPNEGO flow (trino/_spnego.py) +# built on python-gssapi; both extras pull the same dependencies. +kerberos_require = ["gssapi", + # PyPy compatibility issue https://github.com/jborean93/pykrb5/issues/49 + "krb5 == 0.5.1"] +gssapi_require = kerberos_require sqlalchemy_require = ["sqlalchemy >= 1.3"] +socks_require = ["httpx2[socks]"] external_authentication_token_cache_require = ["keyring"] # We don't add localstorage_require to all_require as users must explicitly opt in to use keyring. -all_require = kerberos_require + sqlalchemy_require +all_require = sqlalchemy_require tests_require = all_require + [ - # httpretty >= 1.1 duplicates requests in `httpretty.latest_requests` - # https://github.com/gabrielfalcao/HTTPretty/issues/425 - "httpretty < 1.1", "pytest", + "pytest-asyncio", "pytest-runner", "pre-commit", "black", @@ -83,12 +83,11 @@ ], python_requires=">=3.9", install_requires=[ + "httpx2[http2]", "lz4", "orjson >= 3.11.0 ; platform_python_implementation != 'PyPy'", "python-dateutil", "pytz", - # CVE-2024-47081 - "requests>=2.32.4", "tzlocal", "zstandard", ], @@ -97,6 +96,7 @@ "kerberos": kerberos_require, "gssapi": gssapi_require, "sqlalchemy": sqlalchemy_require, + "socks": socks_require, "tests": tests_require, "external-authentication-token-cache": external_authentication_token_cache_require, }, diff --git a/tests/integration/test_dbapi_integration.py b/tests/integration/test_dbapi_integration.py index 49556077..21f1fc16 100644 --- a/tests/integration/test_dbapi_integration.py +++ b/tests/integration/test_dbapi_integration.py @@ -22,8 +22,8 @@ from typing import Tuple from zoneinfo import ZoneInfo +import httpx2 import pytest -import requests from tzlocal import get_localzone_name # type: ignore import trino @@ -1580,18 +1580,19 @@ def retrieve_client_tags_from_query(run_trino, client_tags): api_url = "http://" + trino_connection.host + ":" + str(trino_connection.port) if trino_version() >= 483: - session = requests.Session() - resp = session.post(api_url + "/ui/auth/login", json={ - "username": "admin", "password": "" - }) - assert resp.ok, f"POST request to /ui/auth/login failed: {resp.status_code} {resp.reason}" - query_info = session.get(api_url + "/ui/api/query/" + cur._query.query_id).json() + with httpx2.Client() as session: + resp = session.post(api_url + "/ui/auth/login", json={ + "username": "admin", "password": "" + }) + assert resp.is_success, \ + f"POST request to /ui/auth/login failed: {resp.status_code} {resp.reason_phrase}" + query_info = session.get(api_url + "/ui/api/query/" + cur._query.query_id).json() else: - query_info = requests.post(api_url + "/ui/login", data={ + query_info = httpx2.post(api_url + "/ui/login", data={ "username": "admin", "password": "", "redirectPath": api_url + '/ui/api/query/' + cur._query.query_id - }).json() + }, follow_redirects=True).json() query_client_tags = query_info['session']['clientTags'] return query_client_tags diff --git a/tests/unit/aio/__init__.py b/tests/unit/aio/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/unit/aio/test_async_client.py b/tests/unit/aio/test_async_client.py new file mode 100644 index 00000000..ff868ea0 --- /dev/null +++ b/tests/unit/aio/test_async_client.py @@ -0,0 +1,219 @@ +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import asyncio +import uuid + +import httpx2 +import pytest + +import trino.auth +import trino.exceptions +from tests.unit.mock_http import MockTrinoServer +from tests.unit.oauth_test_utils import _get_token_requests +from tests.unit.oauth_test_utils import _post_statement_requests +from tests.unit.oauth_test_utils import GetTokenCallback +from tests.unit.oauth_test_utils import PostStatementCallback +from tests.unit.oauth_test_utils import REDIRECT_RESOURCE +from tests.unit.oauth_test_utils import RedirectHandler +from tests.unit.oauth_test_utils import TOKEN_PATH +from tests.unit.oauth_test_utils import TOKEN_RESOURCE +from trino import constants +from trino.aio.client import _AsyncRequestHeartbeat +from trino.aio.client import AsyncTrinoRequest +from trino.client import ClientSession + + +async def test_request_headers_are_sent(trino_server): + req = AsyncTrinoRequest( + host="coordinator", + port=8080, + client_session=ClientSession( + user="test_user", + client_tags=["tag1", "tag2"], + ), + http_session=trino_server.async_client(), + ) + + await req.post("SELECT 1") + headers = trino_server.last_request.headers + assert headers[constants.HEADER_USER] == "test_user" + assert headers[constants.HEADER_CLIENT_TAGS] == "tag1,tag2" + assert headers[constants.HEADER_CONTENT_TYPE] == constants.CONTENT_TYPE_TEXT_UTF8 + # None-valued headers (like the unset transaction id) are not sent + assert constants.HEADER_TRANSACTION not in headers + + await req.get(req.statement_url) + headers = trino_server.last_request.headers + assert headers[constants.HEADER_USER] == "test_user" + assert constants.HEADER_CONTENT_TYPE not in headers + + await req.aclose() + + +async def test_5XX_error_retry(): + calls = [] + + def handler(request): + calls.append(request) + return httpx2.Response(502) + + attempts = 3 + req = AsyncTrinoRequest( + host="coordinator", + port=8080, + client_session=ClientSession(user="test"), + http_session=httpx2.AsyncClient(transport=httpx2.MockTransport(handler)), + max_attempts=attempts, + ) + + await req.post("SELECT 1") + assert len(calls) == attempts + + await req.get(req.statement_url) + assert len(calls) == 2 * attempts + + await req.aclose() + + +async def test_transport_error_retry(): + calls = [] + + def handler(request): + calls.append(request) + raise httpx2.ConnectError("connection refused", request=request) + + attempts = 3 + req = AsyncTrinoRequest( + host="coordinator", + port=8080, + client_session=ClientSession(user="test"), + http_session=httpx2.AsyncClient(transport=httpx2.MockTransport(handler)), + max_attempts=attempts, + ) + + with pytest.raises(httpx2.ConnectError): + await req.post("SELECT 1") + assert len(calls) == attempts + + await req.aclose() + + +async def test_request_timeout_is_propagated_to_transport(): + seen = [] + + def record_timeout(request): + seen.append(request.extensions.get("timeout")) + return httpx2.Response(200, text="ok") + + req = AsyncTrinoRequest( + host="coordinator", + port=8080, + client_session=ClientSession(user="test"), + http_session=httpx2.AsyncClient(transport=httpx2.MockTransport(record_timeout)), + max_attempts=1, + request_timeout=(0.1, 0.2), + ) + + await req.post("SELECT 1") + assert seen == [{"connect": 0.1, "read": 0.2, "write": None, "pool": None}] + + await req.aclose() + + +async def test_oauth2_authentication_flow(sample_post_response_data): + token = str(uuid.uuid4()) + challenge_id = str(uuid.uuid4()) + + redirect_server = f"{REDIRECT_RESOURCE}/{challenge_id}" + token_server = f"{TOKEN_RESOURCE}/{challenge_id}" + + server = MockTrinoServer() + + post_statement_callback = PostStatementCallback(redirect_server, token_server, [token], sample_post_response_data) + server.register("POST", constants.URL_STATEMENT_PATH, post_statement_callback) + + get_token_callback = GetTokenCallback(token_server, token, attempts=2) + server.register("GET", f"/{TOKEN_PATH}/{challenge_id}", get_token_callback) + + redirect_handler = RedirectHandler() + auth = trino.auth.OAuth2Authentication(redirect_auth_url_handler=redirect_handler) + + request = AsyncTrinoRequest( + host="coordinator", + port=constants.DEFAULT_TLS_PORT, + client_session=ClientSession( + user="test", + ), + http_scheme=constants.HTTPS, + http_session=server.async_client(), + auth=auth) + response = await request.post("select 1") + + assert response.request.headers['Authorization'] == f"Bearer {token}" + assert redirect_handler.redirect_server == redirect_server + assert get_token_callback.attempts == 0 + assert len(_post_statement_requests(server)) == 2 + assert len(_get_token_requests(server, challenge_id)) == 2 + + # The cached token is reused on the next request without a new challenge. + response = await request.post("select 2") + assert response.status_code == 200 + assert len(_post_statement_requests(server)) == 3 + + await request.aclose() + + +async def test_heartbeat_sends_head_requests(): + calls = [] + + def handler(request): + calls.append(request) + return httpx2.Response(200, text="") + + req = AsyncTrinoRequest( + host="coordinator", + port=8080, + client_session=ClientSession(user="test"), + http_session=httpx2.AsyncClient(transport=httpx2.MockTransport(handler)), + ) + req._next_uri = "http://coordinator:8080/v1/statement/x/1" + + async with _AsyncRequestHeartbeat(req, interval=0.01): + await asyncio.sleep(0.1) + + assert len(calls) >= 1 + assert all(request.method == "HEAD" for request in calls) + + await req.aclose() + + +async def test_heartbeat_stops_when_unsupported(): + calls = [] + + def handler(request): + calls.append(request) + return httpx2.Response(405, text="") + + req = AsyncTrinoRequest( + host="coordinator", + port=8080, + client_session=ClientSession(user="test"), + http_session=httpx2.AsyncClient(transport=httpx2.MockTransport(handler)), + ) + req._next_uri = "http://coordinator:8080/v1/statement/x/1" + + async with _AsyncRequestHeartbeat(req, interval=0.01): + await asyncio.sleep(0.1) + + assert len(calls) == 1 + + await req.aclose() diff --git a/tests/unit/aio/test_async_dbapi.py b/tests/unit/aio/test_async_dbapi.py new file mode 100644 index 00000000..64b4459b --- /dev/null +++ b/tests/unit/aio/test_async_dbapi.py @@ -0,0 +1,234 @@ +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import base64 + +import httpx2 +import pytest + +import trino.aio +import trino.exceptions +from tests.unit.mock_http import MockTrinoServer +from trino import constants +from trino.transaction import IsolationLevel + +_QUERY_ID = "20210817_140827_00000_arvdv" +SERVER_ADDRESS = "https://coordinator" + +_BIGINT_COLUMNS = [{ + "name": "x", + "type": "bigint", + "typeSignature": {"rawType": "bigint", "arguments": [], "typeArguments": []}, +}] + + +def _statement_path(token): + return f"{constants.URL_STATEMENT_PATH}/{_QUERY_ID}/{token}" + + +def _statement_uri(token): + return f"{SERVER_ADDRESS}{_statement_path(token)}" + + +def _select_query_server(rows_pages): + """A mock coordinator serving a SELECT query returning the given pages of rows.""" + server = MockTrinoServer() + server.register("POST", constants.URL_STATEMENT_PATH, json={ + "id": _QUERY_ID, + "nextUri": _statement_uri(1), + "infoUri": f"{SERVER_ADDRESS}/query.html?{_QUERY_ID}", + "stats": {"state": "QUEUED"}, + }) + for index, page in enumerate(rows_pages, start=1): + response = { + "id": _QUERY_ID, + "infoUri": f"{SERVER_ADDRESS}/query.html?{_QUERY_ID}", + "columns": _BIGINT_COLUMNS, + "data": page, + "stats": {"state": "RUNNING"}, + } + if index < len(rows_pages): + response["nextUri"] = _statement_uri(index + 1) + else: + response["stats"] = {"state": "FINISHED"} + server.register("GET", _statement_path(index), json=response) + return server + + +def _connect(server, **kwargs): + return trino.aio.connect( + "coordinator", + user="test", + http_scheme=constants.HTTPS, + http_session=server.async_client(), + **kwargs, + ) + + +async def test_fetchall(): + server = _select_query_server([[[1]], [[2], [3]]]) + async with _connect(server) as conn: + cur = conn.cursor() + await cur.execute("SELECT x FROM some_table") + rows = await cur.fetchall() + assert rows == [[1], [2], [3]] + assert cur.query_id == _QUERY_ID + assert cur.description[0].name == "x" + assert cur.description[0].type_code == "bigint" + + +async def test_async_iteration(): + server = _select_query_server([[[1]], [[2], [3]]]) + async with _connect(server) as conn: + cur = conn.cursor() + await cur.execute("SELECT x FROM some_table") + rows = [row async for row in cur] + assert rows == [[1], [2], [3]] + + +async def test_fetchone_and_fetchmany(): + server = _select_query_server([[[1], [2], [3]]]) + async with _connect(server) as conn: + cur = conn.cursor() + await cur.execute("SELECT x FROM some_table") + assert await cur.fetchone() == [1] + assert await cur.fetchmany(2) == [[2], [3]] + assert await cur.fetchmany(2) == [] + + +async def test_cursor_close_cancels_unfinished_query(): + server = MockTrinoServer() + server.register("POST", constants.URL_STATEMENT_PATH, json={ + "id": _QUERY_ID, + "nextUri": _statement_uri(1), + "infoUri": f"{SERVER_ADDRESS}/query.html?{_QUERY_ID}", + "stats": {"state": "QUEUED"}, + }) + server.register("GET", _statement_path(1), json={ + "id": _QUERY_ID, + "nextUri": _statement_uri(2), + "infoUri": f"{SERVER_ADDRESS}/query.html?{_QUERY_ID}", + "columns": _BIGINT_COLUMNS, + "data": [[1]], + "stats": {"state": "RUNNING"}, + }) + server.register("DELETE", _statement_path(2), status=204, text="") + + async with _connect(server) as conn: + cur = conn.cursor() + await cur.execute("SELECT x FROM some_table") + assert cur._query.finished is False + await cur.close() + + assert len(server.requests(method="DELETE")) == 1 + + +async def test_spooled_protocol_with_inline_segment(): + encoded_rows = base64.b64encode(b"[[1], [2]]").decode("utf8") + server = MockTrinoServer() + server.register("POST", constants.URL_STATEMENT_PATH, json={ + "id": _QUERY_ID, + "nextUri": _statement_uri(1), + "infoUri": f"{SERVER_ADDRESS}/query.html?{_QUERY_ID}", + "stats": {"state": "QUEUED"}, + }) + server.register("GET", _statement_path(1), json={ + "id": _QUERY_ID, + "infoUri": f"{SERVER_ADDRESS}/query.html?{_QUERY_ID}", + "columns": _BIGINT_COLUMNS, + "data": { + "encoding": "json", + "segments": [{ + "type": "inline", + "metadata": {"segmentSize": "10"}, + "data": encoded_rows, + }], + }, + "stats": {"state": "FINISHED"}, + }) + + async with _connect(server, encoding="json") as conn: + cur = conn.cursor() + await cur.execute("SELECT x FROM some_table") + rows = await cur.fetchall() + assert rows == [[1], [2]] + + +async def test_spooled_protocol_downloads_segment_and_acknowledges(monkeypatch): + segment_rows = b"[[1], [2], [3]]" + + server = MockTrinoServer() + server.register("POST", constants.URL_STATEMENT_PATH, json={ + "id": _QUERY_ID, + "nextUri": _statement_uri(1), + "infoUri": f"{SERVER_ADDRESS}/query.html?{_QUERY_ID}", + "stats": {"state": "QUEUED"}, + }) + server.register("GET", _statement_path(1), json={ + "id": _QUERY_ID, + "infoUri": f"{SERVER_ADDRESS}/query.html?{_QUERY_ID}", + "columns": _BIGINT_COLUMNS, + "data": { + "encoding": "json", + "segments": [{ + "type": "spooled", + "metadata": {"segmentSize": str(len(segment_rows))}, + "uri": "https://storage.example.com/segments/1", + "ackUri": "https://coordinator/v1/spooled/ack/1", + "headers": {"X-Trino-Spooling-Token": ["token-abc"]}, + }], + }, + "stats": {"state": "FINISHED"}, + }) + + def segment_handler(request): + assert request.headers["X-Trino-Spooling-Token"] == "token-abc" + return httpx2.Response(200, content=segment_rows) + + server.register("GET", "/segments/1", segment_handler) + server.register("GET", "/v1/spooled/ack/1", status=204, text="") + + # Spooled segment downloads run over a fresh unauthenticated client built + # by create_http_client; route those clients through the mock transport. + original = trino.aio.AsyncTrinoRequest.create_http_client.__func__ + + def create_with_mock_transport(cls, **kwargs): + kwargs.setdefault("transport", server.transport()) + return original(cls, **kwargs) + + monkeypatch.setattr( + trino.aio.AsyncTrinoRequest, "create_http_client", classmethod(create_with_mock_transport)) + + async with _connect(server, encoding="json") as conn: + cur = conn.cursor() + await cur.execute("SELECT x FROM some_table") + rows = await cur.fetchall() + assert rows == [[1], [2], [3]] + + # conn.close() drained the acknowledgment task. + ack_requests = server.requests(method="GET", path="/v1/spooled/ack/1") + assert len(ack_requests) == 1 + + +async def test_transactions_not_supported(): + with pytest.raises(trino.exceptions.NotSupportedError): + trino.aio.connect("coordinator", user="test", isolation_level=IsolationLevel.SERIALIZABLE) + + conn = trino.aio.connect("coordinator", user="test") + with pytest.raises(trino.exceptions.NotSupportedError): + conn.start_transaction() + await conn.close() + + +async def test_http_session_is_defaulted_when_not_specified(): + async with trino.aio.connect("coordinator", user="test") as conn: + assert isinstance(conn._http_session, httpx2.AsyncClient) + assert conn._create_request()._http_session is conn._http_session diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 39c4eaa0..4143cf9f 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -9,11 +9,11 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -from unittest.mock import MagicMock -from unittest.mock import patch - import pytest +from tests.unit.mock_http import MATCH_ALL +from tests.unit.mock_http import MockTrinoServer + @pytest.fixture(scope="session") def sample_post_response_data(): @@ -285,15 +285,18 @@ def sample_get_error_response_data(): @pytest.fixture -def mock_get_and_post(): - post = MagicMock() - get = MagicMock() - - with patch("trino.client.TrinoRequest.http") as mock_requests: - mock_requests.Session.return_value.get = get - mock_requests.Session.return_value.post = post - - yield get, post +def trino_server(): + """ + A recording mock coordinator answering every request with an empty JSON + body. Tests build clients with ``trino_server.client()`` (or through + ``TrinoRequest.create_http_client(transport=trino_server.transport())``) + and assert on the requests that actually went on the wire. + """ + server = MockTrinoServer() + server.register("POST", MATCH_ALL, json={}) + server.register("GET", MATCH_ALL, json={}) + server.register("DELETE", MATCH_ALL, json={}) + yield server def sqlalchemy_version() -> str: diff --git a/tests/unit/mock_http.py b/tests/unit/mock_http.py new file mode 100644 index 00000000..bb6165af --- /dev/null +++ b/tests/unit/mock_http.py @@ -0,0 +1,125 @@ +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Hand-written HTTP mocking for the unit tests, replacing httpretty. + +``MockTrinoServer`` is a recording router served through ``httpx2.MockTransport``. +Handlers receive the ``httpx2.Request`` and return an ``httpx2.Response``; every +request is recorded so tests can assert on what actually went on the wire. One +server instance can serve synchronous and asynchronous clients alike. +""" +from __future__ import annotations + +import json as _json +import re +import threading +from re import Pattern +from typing import Any +from typing import Callable +from typing import Dict +from typing import List +from typing import Optional +from typing import Union + +import httpx2 + +Handler = Callable[[httpx2.Request], httpx2.Response] + + +class MockTrinoServer: + def __init__(self) -> None: + self._routes: List[tuple[str, Union[str, Pattern[str]], Handler]] = [] + self._requests: List[httpx2.Request] = [] + # The handler can be hit concurrently from many threads or tasks. + self._lock = threading.Lock() + + def register( + self, + method: str, + path: Union[str, Pattern[str]], + handler: Optional[Handler] = None, + *, + json: Any = None, + text: Optional[str] = None, + status: int = 200, + headers: Optional[Dict[str, str]] = None, + ) -> None: + """ + Register a route. ``path`` is matched exactly (string) or by regex + ``search`` (compiled pattern) against the request path. Instead of a + ``handler`` callable, a static ``json``/``text`` body with ``status`` + and ``headers`` can be given. Like httpretty, the most recently + registered matching route wins. + """ + if handler is None: + static_json = json + static_text = text + static_status = status + static_headers = headers or {} + + def handler(request: httpx2.Request) -> httpx2.Response: + if static_json is not None: + return httpx2.Response(static_status, json=static_json, headers=static_headers) + return httpx2.Response(static_status, text=static_text or "", headers=static_headers) + + self._routes.append((method.upper(), path, handler)) + + def _matches(self, route_path: Union[str, Pattern[str]], request_path: str) -> bool: + if isinstance(route_path, str): + return route_path == request_path + return route_path.search(request_path) is not None + + def handle(self, request: httpx2.Request) -> httpx2.Response: + with self._lock: + self._requests.append(request) + for method, path, handler in reversed(self._routes): + if request.method == method and self._matches(path, request.url.path): + return handler(request) + return httpx2.Response(404, text=f"no route for {request.method} {request.url.path}") + + def transport(self) -> httpx2.MockTransport: + return httpx2.MockTransport(self.handle) + + def client(self, **kwargs: Any) -> httpx2.Client: + kwargs.setdefault("follow_redirects", True) + return httpx2.Client(transport=self.transport(), **kwargs) + + def async_client(self, **kwargs: Any) -> httpx2.AsyncClient: + kwargs.setdefault("follow_redirects", True) + return httpx2.AsyncClient(transport=self.transport(), **kwargs) + + def requests(self, method: Optional[str] = None, path: Optional[str] = None) -> List[httpx2.Request]: + """Recorded requests, optionally filtered by method and exact path.""" + with self._lock: + recorded = list(self._requests) + if method is not None: + recorded = [r for r in recorded if r.method == method.upper()] + if path is not None: + recorded = [r for r in recorded if r.url.path == path] + return recorded + + @property + def last_request(self) -> httpx2.Request: + with self._lock: + return self._requests[-1] + + def reset_requests(self) -> None: + with self._lock: + self._requests.clear() + + +def json_response(status: int, body: Any, headers: Optional[Dict[str, str]] = None) -> httpx2.Response: + return httpx2.Response(status, text=_json.dumps(body), headers=headers or {}) + + +# Compiled catch-all pattern useful for fixtures. +MATCH_ALL = re.compile(r".*") diff --git a/tests/unit/oauth_test_utils.py b/tests/unit/oauth_test_utils.py index 956fee78..1f0d8da5 100644 --- a/tests/unit/oauth_test_utils.py +++ b/tests/unit/oauth_test_utils.py @@ -11,11 +11,13 @@ # limitations under the License. import json import re +import threading import uuid from collections import namedtuple -import httpretty +import httpx2 +from tests.unit.mock_http import MockTrinoServer from trino import constants SERVER_ADDRESS = "https://coordinator" @@ -48,30 +50,32 @@ def __init__(self, redirect_server, token_server, tokens, sample_post_response_d self.tokens = tokens self.sample_post_response_data = sample_post_response_data - def __call__(self, request, uri, response_headers): + def __call__(self, request: httpx2.Request) -> httpx2.Response: authorization = request.headers.get("Authorization") if authorization and authorization.replace("Bearer ", "") in self.tokens: - return [200, response_headers, json.dumps(self.sample_post_response_data)] + return httpx2.Response(200, text=json.dumps(self.sample_post_response_data)) elif self.redirect_server is None and self.token_server is not None: - return [401, - { - 'Www-Authenticate': ( - 'Bearer realm="Trino", token_type="JWT", ' - f'Bearer x_token_server="{self.token_server}"' - ), - 'Basic realm': '"Trino"' - }, - ""] - return [401, - { + return httpx2.Response( + 401, + headers={ 'Www-Authenticate': ( 'Bearer realm="Trino", token_type="JWT", ' - f'Bearer x_redirect_server="{self.redirect_server}", ' - f'x_token_server="{self.token_server}"' + f'Bearer x_token_server="{self.token_server}"' ), 'Basic realm': '"Trino"' }, - ""] + text="") + return httpx2.Response( + 401, + headers={ + 'Www-Authenticate': ( + 'Bearer realm="Trino", token_type="JWT", ' + f'Bearer x_redirect_server="{self.redirect_server}", ' + f'x_token_server="{self.token_server}"' + ), + 'Basic realm': '"Trino"' + }, + text="") class GetTokenCallback: @@ -80,73 +84,68 @@ def __init__(self, token_server, token, attempts=1): self.token = token self.attempts = attempts - def __call__(self, request, uri, response_headers): + def __call__(self, request: httpx2.Request) -> httpx2.Response: self.attempts -= 1 if self.attempts < 0: - return [404, response_headers, "{}"] + return httpx2.Response(404, text="{}") if self.attempts == 0: - return [200, response_headers, f'{{"token": "{self.token}"}}'] - return [200, response_headers, f'{{"nextUri": "{self.token_server}"}}'] + return httpx2.Response(200, text=f'{{"token": "{self.token}"}}') + return httpx2.Response(200, text=f'{{"nextUri": "{self.token_server}"}}') -def _get_token_requests(challenge_id): - return list(filter( - lambda r: r.method == "GET" and r.path == f"/{TOKEN_PATH}/{challenge_id}", - httpretty.latest_requests())) +def _get_token_requests(server: MockTrinoServer, challenge_id): + return server.requests(method="GET", path=f"/{TOKEN_PATH}/{challenge_id}") -def _post_statement_requests(): - return list(filter( - lambda r: r.method == "POST" and r.path == constants.URL_STATEMENT_PATH, - httpretty.latest_requests())) +def _post_statement_requests(server: MockTrinoServer): + return server.requests(method="POST", path=constants.URL_STATEMENT_PATH) class MultithreadedTokenServer: Challenge = namedtuple('Challenge', ['token', 'attempts']) - def __init__(self, sample_post_response_data, attempts=1): + def __init__(self, server: MockTrinoServer, sample_post_response_data, attempts=1): self.tokens = set() self.challenges = {} self.sample_post_response_data = sample_post_response_data self.attempts = attempts + # The callbacks can run concurrently from multiple threads. + self._lock = threading.Lock() # bind post statement - httpretty.register_uri( - method=httpretty.POST, - uri=f"{SERVER_ADDRESS}{constants.URL_STATEMENT_PATH}", - body=self.post_statement_callback) + server.register("POST", constants.URL_STATEMENT_PATH, self.post_statement_callback) # bind get token - httpretty.register_uri( - method=httpretty.GET, - uri=re.compile(rf"{TOKEN_RESOURCE}/.*"), - body=self.get_token_callback) + server.register("GET", re.compile(rf"^/{TOKEN_PATH}/.*"), self.get_token_callback) - # noinspection PyUnusedLocal - def post_statement_callback(self, request, uri, response_headers): + def post_statement_callback(self, request: httpx2.Request) -> httpx2.Response: authorization = request.headers.get("Authorization") - if authorization and authorization.replace("Bearer ", "") in self.tokens: - return [200, response_headers, json.dumps(self.sample_post_response_data)] - - challenge_id = str(uuid.uuid4()) - token = str(uuid.uuid4()) - self.tokens.add(token) - self.challenges[challenge_id] = MultithreadedTokenServer.Challenge(token, self.attempts) - redirect_server = f"{REDIRECT_RESOURCE}/{challenge_id}" - token_server = f"{TOKEN_RESOURCE}/{challenge_id}" - return [401, {'Www-Authenticate': f'Bearer x_redirect_server="{redirect_server}", ' - f'x_token_server="{token_server}"', - 'Basic realm': '"Trino"'}, ""] - - # noinspection PyUnusedLocal - def get_token_callback(self, request, uri, response_headers): - challenge_id = uri.replace(f"{TOKEN_RESOURCE}/", "") - challenge = self.challenges[challenge_id] - challenge = challenge._replace(attempts=challenge.attempts - 1) - self.challenges[challenge_id] = challenge + with self._lock: + if authorization and authorization.replace("Bearer ", "") in self.tokens: + return httpx2.Response(200, text=json.dumps(self.sample_post_response_data)) + + challenge_id = str(uuid.uuid4()) + token = str(uuid.uuid4()) + self.tokens.add(token) + self.challenges[challenge_id] = MultithreadedTokenServer.Challenge(token, self.attempts) + redirect_server = f"{REDIRECT_RESOURCE}/{challenge_id}" + token_server = f"{TOKEN_RESOURCE}/{challenge_id}" + return httpx2.Response( + 401, + headers={'Www-Authenticate': f'Bearer x_redirect_server="{redirect_server}", ' + f'x_token_server="{token_server}"', + 'Basic realm': '"Trino"'}, + text="") + + def get_token_callback(self, request: httpx2.Request) -> httpx2.Response: + challenge_id = request.url.path.replace(f"/{TOKEN_PATH}/", "") + with self._lock: + challenge = self.challenges[challenge_id] + challenge = challenge._replace(attempts=challenge.attempts - 1) + self.challenges[challenge_id] = challenge if challenge.attempts < 0: - return [404, response_headers, "{}"] + return httpx2.Response(404, text="{}") if challenge.attempts == 0: - return [200, response_headers, f'{{"token": "{challenge.token}"}}'] - return [200, response_headers, f'{{"nextUri": "{uri}"}}'] + return httpx2.Response(200, text=f'{{"token": "{challenge.token}"}}') + return httpx2.Response(200, text=f'{{"nextUri": "{str(request.url)}"}}') diff --git a/tests/unit/test_auth_gssapi.py b/tests/unit/test_auth_gssapi.py index f4015f50..d459e559 100644 --- a/tests/unit/test_auth_gssapi.py +++ b/tests/unit/test_auth_gssapi.py @@ -13,10 +13,10 @@ from typing import Any import pytest -import requests gssapi = pytest.importorskip("gssapi", exc_type=ImportError) +from trino._spnego import SPNEGOAuth # noqa: E402 from trino.auth import GSSAPIAuthentication # noqa: E402 @@ -78,10 +78,16 @@ def test_authentication_gssapi_init_arguments( ): auth = GSSAPIAuthentication(**options) - session = requests.Session() - with expected_exception: - auth.set_http_session(session) + http_auth = auth.get_http_auth() + + assert isinstance(http_auth, SPNEGOAuth) + assert http_auth.target_name == expected_hostname + assert http_auth.creds == expected_credentials + - assert session.auth.target_name == expected_hostname - assert session.auth.creds == expected_credentials +def test_authentication_gssapi_client_arguments(): + auth = GSSAPIAuthentication(ca_bundle="/tmp/ca.crt") + arguments = auth.get_client_arguments() + assert arguments["trust_env"] is False + assert arguments["verify"] == "/tmp/ca.crt" diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index c5d33dbd..3c592f45 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -11,28 +11,24 @@ # limitations under the License. import base64 import threading -import time import urllib import uuid from typing import Dict from typing import Optional -from unittest import mock from unittest import TestCase -from urllib.parse import urlparse from zoneinfo import ZoneInfoNotFoundError -import httpretty +import httpx2 import keyring try: import orjson as json except ImportError: import json import pytest -import requests -from httpretty import httprettified from tzlocal import get_localzone_name # type: ignore import trino.exceptions +from tests.unit.mock_http import MockTrinoServer from tests.unit.oauth_test_utils import _get_token_requests from tests.unit.oauth_test_utils import _post_statement_requests from tests.unit.oauth_test_utils import GetTokenCallback @@ -41,12 +37,14 @@ from tests.unit.oauth_test_utils import REDIRECT_RESOURCE from tests.unit.oauth_test_utils import RedirectHandler from tests.unit.oauth_test_utils import RedirectHandlerWithException -from tests.unit.oauth_test_utils import SERVER_ADDRESS +from tests.unit.oauth_test_utils import TOKEN_PATH from tests.unit.oauth_test_utils import TOKEN_RESOURCE from trino import __version__ from trino import constants from trino.auth import _OAuth2KeyRingTokenCache from trino.auth import _OAuth2TokenBearer +from trino.auth import GSSAPIAuthentication +from trino.auth import KerberosAuthentication from trino.client import _DelayExponential from trino.client import _retry_with from trino.client import _RetryWithExponentialBackoff @@ -55,34 +53,28 @@ from trino.client import TrinoQuery from trino.client import TrinoRequest from trino.client import TrinoResult +from trino.exceptions import SPNEGOExchangeError try: - from requests_kerberos.exceptions import KerberosExchangeError - from trino.auth import KerberosAuthentication + import gssapi # noqa: F401 + _GSSAPI_AVAILABLE = True except ImportError: - KerberosAuthentication = None - KerberosExchangeError = None + _GSSAPI_AVAILABLE = False -try: - from requests_gssapi.exceptions import SPNEGOExchangeError - from trino.auth import GSSAPIAuthentication -except ImportError: - GSSAPIAuthentication = None - SPNEGOExchangeError = None - -requires_kerberos = pytest.mark.skipif( - KerberosAuthentication is None, - reason="requests_kerberos is not installed", -) requires_gssapi = pytest.mark.skipif( - GSSAPIAuthentication is None, - reason="gssapi is not available (CPython-only)", + not _GSSAPI_AVAILABLE, + reason="gssapi is not installed", ) -@mock.patch("trino.client.TrinoRequest.http") -def test_trino_initial_request(mock_requests, sample_post_response_data): - mock_requests.Response.return_value.text = json.dumps(sample_post_response_data) +def _json_response(data, status_code=200, headers=None): + body = json.dumps(data) + if isinstance(body, str): + body = body.encode("utf-8") + return httpx2.Response(status_code, content=body, headers=headers or {}) + + +def test_trino_initial_request(sample_post_response_data): req = TrinoRequest( host="coordinator", port=8080, @@ -96,17 +88,14 @@ def test_trino_initial_request(mock_requests, sample_post_response_data): http_scheme="http", ) - http_resp = TrinoRequest.http.Response() - http_resp.status_code = 200 + http_resp = _json_response(sample_post_response_data) status = req.process(http_resp) assert status.next_uri == sample_post_response_data["nextUri"] assert status.id == sample_post_response_data["id"] -def test_request_headers(mock_get_and_post): - get, post = mock_get_and_post - +def test_request_headers(trino_server): catalog = "test_catalog" schema = "test_schema" user = "test_user" @@ -144,6 +133,7 @@ def test_request_headers(mock_get_and_post): } ), http_scheme="http", + http_session=trino_server.client(), ) def assert_headers(headers): @@ -153,7 +143,8 @@ def assert_headers(headers): assert headers[constants.HEADER_ORIGINAL_USER] == user assert headers[constants.HEADER_USER] == authorization_user assert headers[constants.HEADER_SESSION] == "" - assert headers[constants.HEADER_TRANSACTION] is None + # None-valued headers (like the unset transaction id) are not sent + assert constants.HEADER_TRANSACTION not in headers assert headers[constants.HEADER_TIMEZONE] == timezone assert headers[constants.HEADER_CLIENT_CAPABILITIES] == constants.CLIENT_CAPABILITIES assert headers[accept_encoding_header] == accept_encoding_value @@ -166,27 +157,19 @@ def assert_headers(headers): ) assert headers["User-Agent"] == f"{constants.CLIENT_NAME}/{__version__}" assert headers[constants.HEADER_ENCODING] == encoding - assert len(headers.keys()) == 14 req.post("URL") - _, post_kwargs = post.call_args - post_headers = post_kwargs["headers"] + post_headers = trino_server.last_request.headers assert post_headers[constants.HEADER_CONTENT_TYPE] == constants.CONTENT_TYPE_TEXT_UTF8 - # Content-Type is specific to POST requests. Strip it before reusing the - # shared assertions (which count the common headers present on both). - del post_headers[constants.HEADER_CONTENT_TYPE] assert_headers(post_headers) - req.get("URL") - _, get_kwargs = get.call_args - get_headers = get_kwargs["headers"] + req.get(req.statement_url) + get_headers = trino_server.last_request.headers assert constants.HEADER_CONTENT_TYPE not in get_headers assert_headers(get_headers) -def test_post_sets_content_type_charset(mock_get_and_post): - _, post = mock_get_and_post - +def test_post_sets_content_type_charset(trino_server): req = TrinoRequest( host="coordinator", port=8080, @@ -198,16 +181,15 @@ def test_post_sets_content_type_charset(mock_get_and_post): properties={}, ), http_scheme="http", + http_session=trino_server.client(), ) req.post("SELECT 1") - _, post_kwargs = post.call_args - assert post_kwargs["headers"][constants.HEADER_CONTENT_TYPE] == constants.CONTENT_TYPE_TEXT_UTF8 - + headers = trino_server.last_request.headers + assert headers[constants.HEADER_CONTENT_TYPE] == constants.CONTENT_TYPE_TEXT_UTF8 -def test_post_content_type_can_be_overridden(mock_get_and_post): - _, post = mock_get_and_post +def test_post_content_type_can_be_overridden(trino_server): req = TrinoRequest( host="coordinator", port=8080, @@ -219,16 +201,15 @@ def test_post_content_type_can_be_overridden(mock_get_and_post): properties={}, ), http_scheme="http", + http_session=trino_server.client(), ) req.post("SELECT 1", additional_http_headers={constants.HEADER_CONTENT_TYPE: "application/xyz"}) - _, post_kwargs = post.call_args - assert post_kwargs["headers"][constants.HEADER_CONTENT_TYPE] == "application/xyz" + headers = trino_server.last_request.headers + assert headers[constants.HEADER_CONTENT_TYPE] == "application/xyz" -def test_request_session_properties_headers(mock_get_and_post): - get, post = mock_get_and_post - +def test_request_session_properties_headers(trino_server): req = TrinoRequest( host="coordinator", port=8080, @@ -239,28 +220,25 @@ def test_request_session_properties_headers(mock_get_and_post): "b": "2", "c": "more=v1,v2" } - ) + ), + http_session=trino_server.client(), ) def assert_headers(headers): assert headers[constants.HEADER_SESSION] == "a=1,b=2,c=more%3Dv1%2Cv2" req.post("URL") - _, post_kwargs = post.call_args - assert_headers(post_kwargs["headers"]) + assert_headers(trino_server.last_request.headers) - req.get("URL") - _, get_kwargs = get.call_args - assert_headers(get_kwargs["headers"]) + req.get(req.statement_url) + assert_headers(trino_server.last_request.headers) -def test_additional_request_post_headers(mock_get_and_post): +def test_additional_request_post_headers(trino_server): """ Tests that the `TrinoRequest.post` function can take addtional headers and that it combines them with the existing ones to perform the request. """ - _, post = mock_get_and_post - req = TrinoRequest( host="coordinator", port=8080, @@ -272,6 +250,7 @@ def test_additional_request_post_headers(mock_get_and_post): properties={}, ), http_scheme="http", + http_session=trino_server.client(), ) sql = 'select 1' @@ -287,8 +266,12 @@ def test_additional_request_post_headers(mock_get_and_post): req.post(sql, additional_headers) # Validate that the post call was performed including the addtional headers - _, post_kwargs = post.call_args - assert post_kwargs['headers'] == combined_headers + sent_headers = trino_server.last_request.headers + for key, value in combined_headers.items(): + if value is None: + assert key not in sent_headers + else: + assert sent_headers[key] == value def test_request_invalid_http_headers(): @@ -304,9 +287,7 @@ def test_request_invalid_http_headers(): assert str(value_error.value).startswith("cannot override reserved HTTP header") -def test_request_client_tags_headers(mock_get_and_post): - get, post = mock_get_and_post - +def test_request_client_tags_headers(trino_server): req = TrinoRequest( host="coordinator", port=8080, @@ -314,64 +295,54 @@ def test_request_client_tags_headers(mock_get_and_post): user="test_user", client_tags=["tag1", "tag2"] ), + http_session=trino_server.client(), ) def assert_headers(headers): assert headers[constants.HEADER_CLIENT_TAGS] == "tag1,tag2" req.post("URL") - _, post_kwargs = post.call_args - assert_headers(post_kwargs["headers"]) + assert_headers(trino_server.last_request.headers) - req.get("URL") - _, get_kwargs = get.call_args - assert_headers(get_kwargs["headers"]) + req.get(req.statement_url) + assert_headers(trino_server.last_request.headers) -def test_request_client_tags_headers_no_client_tags(mock_get_and_post): - get, post = mock_get_and_post - +def test_request_client_tags_headers_no_client_tags(trino_server): req = TrinoRequest( host="coordinator", port=8080, client_session=ClientSession( user="test_user", - ) + ), + http_session=trino_server.client(), ) def assert_headers(headers): assert constants.HEADER_CLIENT_TAGS not in headers req.post("URL") - _, post_kwargs = post.call_args - assert_headers(post_kwargs["headers"]) - - req.get("URL") - _, get_kwargs = get.call_args - assert_headers(get_kwargs["headers"]) + assert_headers(trino_server.last_request.headers) + req.get(req.statement_url) + assert_headers(trino_server.last_request.headers) -def test_enabling_https_automatically_when_using_port_443(mock_get_and_post): - _, post = mock_get_and_post +def test_enabling_https_automatically_when_using_port_443(trino_server): req = TrinoRequest( host="coordinator", port=constants.DEFAULT_TLS_PORT, client_session=ClientSession( user="test", ), + http_session=trino_server.client(), ) req.post("SELECT 1") - post_args, _ = post.call_args - parsed_url = urlparse(post_args[0]) + assert trino_server.last_request.url.scheme == constants.HTTPS - assert parsed_url.scheme == constants.HTTPS - - -def test_https_scheme(mock_get_and_post): - _, post = mock_get_and_post +def test_https_scheme(trino_server): req = TrinoRequest( host="coordinator", port=constants.DEFAULT_TLS_PORT, @@ -379,19 +350,17 @@ def test_https_scheme(mock_get_and_post): user="test", ), http_scheme=constants.HTTPS, + http_session=trino_server.client(), ) req.post("SELECT 1") - post_args, _ = post.call_args - parsed_url = urlparse(post_args[0]) - - assert parsed_url.scheme == constants.HTTPS - assert parsed_url.port == constants.DEFAULT_TLS_PORT + url = trino_server.last_request.url + assert url.scheme == constants.HTTPS + # httpx normalizes the default port away + assert url.port in (None, constants.DEFAULT_TLS_PORT) -def test_http_scheme_with_port(mock_get_and_post): - _, post = mock_get_and_post - +def test_http_scheme_with_port(trino_server): req = TrinoRequest( host="coordinator", port=constants.DEFAULT_TLS_PORT, @@ -399,56 +368,97 @@ def test_http_scheme_with_port(mock_get_and_post): user="test", ), http_scheme=constants.HTTP, + http_session=trino_server.client(), ) req.post("SELECT 1") - post_args, _ = post.call_args - parsed_url = urlparse(post_args[0]) + url = trino_server.last_request.url + assert url.scheme == constants.HTTP + assert url.port == constants.DEFAULT_TLS_PORT - assert parsed_url.scheme == constants.HTTP - assert parsed_url.port == constants.DEFAULT_TLS_PORT +def test_request_timeout_is_propagated_to_transport(): + seen = [] -def test_request_timeout(): - timeout = 0.1 - http_scheme = "http" - host = "coordinator" - port = 8080 - url = http_scheme + "://" + host + ":" + str(port) + constants.URL_STATEMENT_PATH + def record_timeout(request): + seen.append(request.extensions.get("timeout")) + return httpx2.Response(200, text="ok") - def long_call(request, uri, headers): - time.sleep(timeout * 2) - return (200, headers, "delayed success") + url = "http://coordinator:8080" + constants.URL_STATEMENT_PATH + + for request_timeout, expected in [ + (0.1, {"connect": 0.1, "read": 0.1, "write": 0.1, "pool": 0.1}), + ((0.1, 0.2), {"connect": 0.1, "read": 0.2, "write": None, "pool": None}), + ]: + seen.clear() + req = TrinoRequest( + host="coordinator", + port=8080, + client_session=ClientSession( + user="test", + ), + http_scheme="http", + http_session=httpx2.Client(transport=httpx2.MockTransport(record_timeout)), + max_attempts=1, + request_timeout=request_timeout, + ) + + req.get(url) + req.post("select 1") + assert seen == [expected, expected] - httpretty.enable() - for method in [httpretty.POST, httpretty.GET]: - httpretty.register_uri(method, url, body=long_call) + +def test_request_timeout(): + def raise_timeout(request): + raise httpx2.ReadTimeout("timed out", request=request) + + url = "http://coordinator:8080" + constants.URL_STATEMENT_PATH # timeout without retry - for request_timeout in [timeout, (timeout, timeout)]: + for request_timeout in [0.1, (0.1, 0.1)]: req = TrinoRequest( - host=host, - port=port, + host="coordinator", + port=8080, client_session=ClientSession( user="test", ), - http_scheme=http_scheme, + http_scheme="http", + http_session=httpx2.Client(transport=httpx2.MockTransport(raise_timeout)), max_attempts=1, request_timeout=request_timeout, ) - with pytest.raises(requests.exceptions.Timeout): + with pytest.raises(httpx2.TimeoutException): req.get(url) - with pytest.raises(requests.exceptions.Timeout): + with pytest.raises(httpx2.TimeoutException): req.post("select 1") - httpretty.disable() - httpretty.reset() + +def test_redirects_are_followed(trino_server): + query_path = constants.URL_STATEMENT_PATH + "/redirected" + trino_server.register( + "GET", constants.URL_STATEMENT_PATH, + status=301, headers={"Location": f"https://coordinator{query_path}"}, text="") + trino_server.register("GET", query_path, json={"redirected": True}) + + req = TrinoRequest( + host="coordinator", + port=constants.DEFAULT_TLS_PORT, + client_session=ClientSession( + user="test", + ), + http_scheme=constants.HTTPS, + http_session=trino_server.client(), + ) + + response = req.get(req.statement_url) + assert response.status_code == 200 + assert response.url.path == query_path + assert len(trino_server.requests(method="GET")) == 2 @pytest.mark.parametrize("attempts", [1, 3, 5]) -@httprettified def test_oauth2_authentication_flow(attempts, sample_post_response_data): token = str(uuid.uuid4()) challenge_id = str(uuid.uuid4()) @@ -456,20 +466,15 @@ def test_oauth2_authentication_flow(attempts, sample_post_response_data): redirect_server = f"{REDIRECT_RESOURCE}/{challenge_id}" token_server = f"{TOKEN_RESOURCE}/{challenge_id}" - post_statement_callback = PostStatementCallback(redirect_server, token_server, [token], sample_post_response_data) + server = MockTrinoServer() # bind post statement - httpretty.register_uri( - method=httpretty.POST, - uri=f"{SERVER_ADDRESS}{constants.URL_STATEMENT_PATH}", - body=post_statement_callback) + post_statement_callback = PostStatementCallback(redirect_server, token_server, [token], sample_post_response_data) + server.register("POST", constants.URL_STATEMENT_PATH, post_statement_callback) # bind get token get_token_callback = GetTokenCallback(token_server, token, attempts) - httpretty.register_uri( - method=httpretty.GET, - uri=token_server, - body=get_token_callback) + server.register("GET", f"/{TOKEN_PATH}/{challenge_id}", get_token_callback) redirect_handler = RedirectHandler() @@ -480,37 +485,30 @@ def test_oauth2_authentication_flow(attempts, sample_post_response_data): user="test", ), http_scheme=constants.HTTPS, + http_session=server.client(), auth=trino.auth.OAuth2Authentication(redirect_auth_url_handler=redirect_handler)) response = request.post("select 1") assert response.request.headers['Authorization'] == f"Bearer {token}" assert redirect_handler.redirect_server == redirect_server assert get_token_callback.attempts == 0 - assert len(_post_statement_requests()) == 2 - assert len(_get_token_requests(challenge_id)) == attempts + assert len(_post_statement_requests(server)) == 2 + assert len(_get_token_requests(server, challenge_id)) == attempts -@httprettified def test_oauth2_refresh_token_flow(sample_post_response_data): token = str(uuid.uuid4()) challenge_id = str(uuid.uuid4()) token_server = f"{TOKEN_RESOURCE}/{challenge_id}" - post_statement_callback = PostStatementCallback(None, token_server, [token], sample_post_response_data) + server = MockTrinoServer() - # bind post statement - httpretty.register_uri( - method=httpretty.POST, - uri=f"{SERVER_ADDRESS}{constants.URL_STATEMENT_PATH}", - body=post_statement_callback) + post_statement_callback = PostStatementCallback(None, token_server, [token], sample_post_response_data) + server.register("POST", constants.URL_STATEMENT_PATH, post_statement_callback) - # bind get token get_token_callback = GetTokenCallback(token_server, token) - httpretty.register_uri( - method=httpretty.GET, - uri=token_server, - body=get_token_callback) + server.register("GET", f"/{TOKEN_PATH}/{challenge_id}", get_token_callback) redirect_handler = RedirectHandlerWithException( trino.exceptions.TrinoAuthError( @@ -523,17 +521,17 @@ def test_oauth2_refresh_token_flow(sample_post_response_data): user="test", ), http_scheme=constants.HTTPS, + http_session=server.client(), auth=trino.auth.OAuth2Authentication(redirect_auth_url_handler=redirect_handler)) response = request.post("select 1") assert response.request.headers['Authorization'] == f"Bearer {token}" assert get_token_callback.attempts == 0 - assert len(_post_statement_requests()) == 2 + assert len(_post_statement_requests(server)) == 2 @pytest.mark.parametrize("attempts", [6, 10]) -@httprettified def test_oauth2_exceed_max_attempts(attempts, sample_post_response_data): token = str(uuid.uuid4()) challenge_id = str(uuid.uuid4()) @@ -541,20 +539,13 @@ def test_oauth2_exceed_max_attempts(attempts, sample_post_response_data): redirect_server = f"{REDIRECT_RESOURCE}/{challenge_id}" token_server = f"{TOKEN_RESOURCE}/{challenge_id}" - post_statement_callback = PostStatementCallback(redirect_server, token_server, [token], sample_post_response_data) + server = MockTrinoServer() - # bind post statement - httpretty.register_uri( - method=httpretty.POST, - uri=f"{SERVER_ADDRESS}{constants.URL_STATEMENT_PATH}", - body=post_statement_callback) + post_statement_callback = PostStatementCallback(redirect_server, token_server, [token], sample_post_response_data) + server.register("POST", constants.URL_STATEMENT_PATH, post_statement_callback) - # bind get token get_token_callback = GetTokenCallback(token_server, token, attempts) - httpretty.register_uri( - method=httpretty.GET, - uri=f"{TOKEN_RESOURCE}/{challenge_id}", - body=get_token_callback) + server.register("GET", f"/{TOKEN_PATH}/{challenge_id}", get_token_callback) redirect_handler = RedirectHandler() @@ -565,6 +556,7 @@ def test_oauth2_exceed_max_attempts(attempts, sample_post_response_data): user="test", ), http_scheme=constants.HTTPS, + http_session=server.client(), auth=trino.auth.OAuth2Authentication(redirect_auth_url_handler=redirect_handler)) with pytest.raises(trino.exceptions.TrinoAuthError) as exp: request.post("select 1") @@ -572,8 +564,8 @@ def test_oauth2_exceed_max_attempts(attempts, sample_post_response_data): assert str(exp.value) == "Exceeded max attempts while getting the token" assert redirect_handler.redirect_server == redirect_server assert get_token_callback.attempts == attempts - _OAuth2TokenBearer.MAX_OAUTH_ATTEMPTS - assert len(_post_statement_requests()) == 1 - assert len(_get_token_requests(challenge_id)) == _OAuth2TokenBearer.MAX_OAUTH_ATTEMPTS + assert len(_post_statement_requests(server)) == 1 + assert len(_get_token_requests(server, challenge_id)) == _OAuth2TokenBearer.MAX_OAUTH_ATTEMPTS @pytest.mark.parametrize("header,error", [ @@ -582,14 +574,11 @@ def test_oauth2_exceed_max_attempts(attempts, sample_post_response_data): ('x_redirect_server="redirect_server", x_token_server="token_server"', 'Error: header info didn\'t match x_redirect_server="redirect_server", x_token_server="token_server"'), # noqa: E501 ('Bearer x_redirect_server="redirect_server"', 'Error: header info didn\'t have x_token_server'), ]) -@httprettified def test_oauth2_authentication_missing_headers(header, error): - # bind post statement - httpretty.register_uri( - method=httpretty.POST, - uri=f"{SERVER_ADDRESS}{constants.URL_STATEMENT_PATH}", - adding_headers={'WWW-Authenticate': header}, - status=401) + server = MockTrinoServer() + server.register( + "POST", constants.URL_STATEMENT_PATH, + status=401, headers={'WWW-Authenticate': header}, text="") request = TrinoRequest( host="coordinator", @@ -598,6 +587,7 @@ def test_oauth2_authentication_missing_headers(header, error): user="test", ), http_scheme=constants.HTTPS, + http_session=server.client(), auth=trino.auth.OAuth2Authentication(redirect_auth_url_handler=RedirectHandler())) with pytest.raises(trino.exceptions.TrinoAuthError) as exp: @@ -616,7 +606,6 @@ def test_oauth2_authentication_missing_headers(header, error): 'x_token_server="{token_server}"' 'Bearer x_redirect_server="{redirect_server}",x_token_server="{token_server}",additional_challenge', ]) -@httprettified def test_oauth2_header_parsing(header, sample_post_response_data): token = str(uuid.uuid4()) challenge_id = str(uuid.uuid4()) @@ -624,26 +613,21 @@ def test_oauth2_header_parsing(header, sample_post_response_data): redirect_server = f"{REDIRECT_RESOURCE}/{challenge_id}?role=test" token_server = f"{TOKEN_RESOURCE}/{challenge_id}" - # noinspection PyUnusedLocal - def post_statement(request, uri, response_headers): + def post_statement(request): authorization = request.headers.get("Authorization") if authorization and authorization.replace("Bearer ", "") in token: - return [200, response_headers, json.dumps(sample_post_response_data)] - return [401, {'Www-Authenticate': header.format(redirect_server=redirect_server, token_server=token_server), - 'Basic realm': '"Trino"'}, ""] + return _json_response(sample_post_response_data) + return httpx2.Response( + 401, + headers={'Www-Authenticate': header.format(redirect_server=redirect_server, token_server=token_server), + 'Basic realm': '"Trino"'}, + text="") - # bind post statement - httpretty.register_uri( - method=httpretty.POST, - uri=f"{SERVER_ADDRESS}{constants.URL_STATEMENT_PATH}", - body=post_statement) + server = MockTrinoServer() + server.register("POST", constants.URL_STATEMENT_PATH, post_statement) - # bind get token get_token_callback = GetTokenCallback(token_server, token) - httpretty.register_uri( - method=httpretty.GET, - uri=token_server, - body=get_token_callback) + server.register("GET", f"/{TOKEN_PATH}/{challenge_id}", get_token_callback) redirect_handler = RedirectHandler() @@ -654,18 +638,18 @@ def post_statement(request, uri, response_headers): user="test", ), http_scheme=constants.HTTPS, + http_session=server.client(), auth=trino.auth.OAuth2Authentication(redirect_auth_url_handler=redirect_handler) ).post("select 1") assert response.request.headers['Authorization'] == f"Bearer {token}" assert redirect_handler.redirect_server == redirect_server assert get_token_callback.attempts == 0 - assert len(_post_statement_requests()) == 2 - assert len(_get_token_requests(challenge_id)) == 1 + assert len(_post_statement_requests(server)) == 2 + assert len(_get_token_requests(server, challenge_id)) == 1 @pytest.mark.parametrize("http_status", [400, 401, 500]) -@httprettified def test_oauth2_authentication_fail_token_server(http_status, sample_post_response_data): token = str(uuid.uuid4()) challenge_id = str(uuid.uuid4()) @@ -673,19 +657,12 @@ def test_oauth2_authentication_fail_token_server(http_status, sample_post_respon redirect_server = f"{REDIRECT_RESOURCE}/{challenge_id}" token_server = f"{TOKEN_RESOURCE}/{challenge_id}" - post_statement_callback = PostStatementCallback(redirect_server, token_server, [token], sample_post_response_data) + server = MockTrinoServer() - # bind post statement - httpretty.register_uri( - method=httpretty.POST, - uri=f"{SERVER_ADDRESS}{constants.URL_STATEMENT_PATH}", - body=post_statement_callback) + post_statement_callback = PostStatementCallback(redirect_server, token_server, [token], sample_post_response_data) + server.register("POST", constants.URL_STATEMENT_PATH, post_statement_callback) - httpretty.register_uri( - method=httpretty.GET, - uri=f"{TOKEN_RESOURCE}/{challenge_id}", - status=http_status, - body="error") + server.register("GET", f"/{TOKEN_PATH}/{challenge_id}", status=http_status, text="error") redirect_handler = RedirectHandler() @@ -696,6 +673,7 @@ def test_oauth2_authentication_fail_token_server(http_status, sample_post_respon user="test", ), http_scheme=constants.HTTPS, + http_session=server.client(), auth=trino.auth.OAuth2Authentication(redirect_auth_url_handler=redirect_handler)) with pytest.raises(trino.exceptions.TrinoAuthError) as exp: @@ -703,18 +681,21 @@ def test_oauth2_authentication_fail_token_server(http_status, sample_post_respon assert redirect_handler.redirect_server == redirect_server assert str(exp.value) == f"Error while getting the token response status code: {http_status}, body: error" - assert len(_post_statement_requests()) == 1 - assert len(_get_token_requests(challenge_id)) == 1 + assert len(_post_statement_requests(server)) == 1 + assert len(_get_token_requests(server, challenge_id)) == 1 -@httprettified def test_multithreaded_oauth2_authentication_flow(sample_post_response_data): redirect_handler = RedirectHandler() auth = trino.auth.OAuth2Authentication(redirect_auth_url_handler=redirect_handler) - token_server = MultithreadedTokenServer(sample_post_response_data) + server = MockTrinoServer() + token_server = MultithreadedTokenServer(server, sample_post_response_data) class RunningThread(threading.Thread): + # Serialize the posts so that exactly one challenge is issued: the + # first post completes its whole OAuth2 flow before any other thread + # sends an unauthenticated request. lock = threading.Lock() def __init__(self): @@ -729,10 +710,9 @@ def run(self) -> None: user="test", ), http_scheme=constants.HTTPS, + http_session=server.client(), auth=auth) for i in range(10): - # apparently HTTPretty in the current version is not thread-safe - # https://github.com/gabrielfalcao/HTTPretty/issues/209 with RunningThread.lock: response = request.post("select 1") self.token = response.request.headers["Authorization"].replace("Bearer ", "") @@ -755,15 +735,12 @@ def run(self) -> None: for challenge_id, challenge in token_server.challenges.items(): assert f"{REDIRECT_RESOURCE}/{challenge_id}" in redirect_handler.redirect_server assert challenge.attempts == 0 - assert len(_get_token_requests(challenge_id)) == 1 + assert len(_get_token_requests(server, challenge_id)) == 1 # 3 threads * (10 POST /statement each + 1 replied request by authentication) - assert len(_post_statement_requests()) == 31 - + assert len(_post_statement_requests(server)) == 31 -@mock.patch("trino.client.TrinoRequest.http") -def test_trino_fetch_request(mock_requests, sample_get_response_data): - mock_requests.Response.return_value.text = json.dumps(sample_get_response_data) +def test_trino_fetch_request(sample_get_response_data): req = TrinoRequest( host="coordinator", port=8080, @@ -777,8 +754,7 @@ def test_trino_fetch_request(mock_requests, sample_get_response_data): http_scheme="http", ) - http_resp = TrinoRequest.http.Response() - http_resp.status_code = 200 + http_resp = _json_response(sample_get_response_data) status = req.process(http_resp) assert status.next_uri == sample_get_response_data["nextUri"] @@ -786,10 +762,7 @@ def test_trino_fetch_request(mock_requests, sample_get_response_data): assert status.rows == sample_get_response_data["data"] -@mock.patch("trino.client.TrinoRequest.http") -def test_trino_fetch_request_data_none(mock_requests, sample_get_response_data_none): - mock_requests.Response.return_value.text = json.dumps(sample_get_response_data_none) - +def test_trino_fetch_request_data_none(sample_get_response_data_none): req = TrinoRequest( host="coordinator", port=8080, @@ -803,8 +776,7 @@ def test_trino_fetch_request_data_none(mock_requests, sample_get_response_data_n http_scheme="http", ) - http_resp = TrinoRequest.http.Response() - http_resp.status_code = 200 + http_resp = _json_response(sample_get_response_data_none) status = req.process(http_resp) assert status.next_uri == sample_get_response_data_none["nextUri"] @@ -812,10 +784,7 @@ def test_trino_fetch_request_data_none(mock_requests, sample_get_response_data_n assert status.rows == [] -@mock.patch("trino.client.TrinoRequest.http") -def test_trino_fetch_error(mock_requests, sample_get_error_response_data): - mock_requests.Response.return_value.text = json.dumps(sample_get_error_response_data) - +def test_trino_fetch_error(sample_get_error_response_data): req = TrinoRequest( host="coordinator", port=8080, @@ -829,8 +798,7 @@ def test_trino_fetch_error(mock_requests, sample_get_error_response_data): http_scheme="http", ) - http_resp = TrinoRequest.http.Response() - http_resp.status_code = 200 + http_resp = _json_response(sample_get_error_response_data) with pytest.raises(trino.exceptions.TrinoUserError) as exception_info: req.process(http_resp) error = exception_info.value @@ -857,9 +825,7 @@ def test_trino_fetch_error(mock_requests, sample_get_error_response_data): (404, trino.exceptions.HttpError, "error 404"), ], ) -def test_trino_connection_error(monkeypatch, error_code, error_type, error_message): - monkeypatch.setattr(TrinoRequest.http.Response, "json", lambda x: {}) - +def test_trino_connection_error(error_code, error_type, error_message): req = TrinoRequest( host="coordinator", port=8080, @@ -873,8 +839,7 @@ def test_trino_connection_error(monkeypatch, error_code, error_type, error_messa http_scheme="http", ) - http_resp = TrinoRequest.http.Response() - http_resp.status_code = error_code + http_resp = httpx2.Response(error_code) with pytest.raises(error_type) as error: req.process(http_resp) assert error_message in str(error) @@ -894,17 +859,13 @@ def test_trino_process_empty_200_response_error(): http_scheme="http", ) - http_resp = TrinoRequest.http.Response() - http_resp.status_code = 200 - http_resp._content = b"" + http_resp = httpx2.Response(200, content=b"") with pytest.raises(trino.exceptions.TrinoConnectionError) as error: req.process(http_resp) assert "received empty response from server (status 200)" in str(error.value) -def test_extra_credential(mock_get_and_post): - _, post = mock_get_and_post - +def test_extra_credential(trino_server): req = TrinoRequest( host="coordinator", port=constants.DEFAULT_TLS_PORT, @@ -912,11 +873,11 @@ def test_extra_credential(mock_get_and_post): user="test", extra_credential=[("a.username", "foo"), ("b.password", "bar")], ), + http_session=trino_server.client(), ) req.post("SELECT 1") - _, post_kwargs = post.call_args - headers = post_kwargs["headers"] + headers = trino_server.last_request.headers assert constants.HEADER_EXTRA_CREDENTIAL in headers assert headers[constants.HEADER_EXTRA_CREDENTIAL] == "a.username=foo, b.password=bar" @@ -949,9 +910,7 @@ def test_extra_credential_key_non_ascii(): assert str(e_info.value) == "only ASCII characters are allowed in extra credential 'çš„'" -def test_extra_credential_value_encoding(mock_get_and_post): - _, post = mock_get_and_post - +def test_extra_credential_value_encoding(trino_server): req = TrinoRequest( host="coordinator", port=constants.DEFAULT_TLS_PORT, @@ -959,18 +918,16 @@ def test_extra_credential_value_encoding(mock_get_and_post): user="test", extra_credential=[("foo", "bar çš„")], ), + http_session=trino_server.client(), ) req.post("SELECT 1") - _, post_kwargs = post.call_args - headers = post_kwargs["headers"] + headers = trino_server.last_request.headers assert constants.HEADER_EXTRA_CREDENTIAL in headers assert headers[constants.HEADER_EXTRA_CREDENTIAL] == "foo=bar+%E7%9A%84" -def test_extra_credential_value_object(mock_get_and_post): - _, post = mock_get_and_post - +def test_extra_credential_value_object(trino_server): class TestCredential: value = "initial" @@ -985,59 +942,52 @@ def __str__(self): client_session=ClientSession( user="test", extra_credential=[("foo", credential)] - ) + ), + http_session=trino_server.client(), ) req.post("SELECT 1") - _, post_kwargs = post.call_args - headers = post_kwargs["headers"] + headers = trino_server.last_request.headers assert constants.HEADER_EXTRA_CREDENTIAL in headers assert headers[constants.HEADER_EXTRA_CREDENTIAL] == "foo=initial" # Make a second request, assert that credential has changed credential.value = "changed" req.post("SELECT 1") - _, post_kwargs = post.call_args - headers = post_kwargs["headers"] + headers = trino_server.last_request.headers assert constants.HEADER_EXTRA_CREDENTIAL in headers assert headers[constants.HEADER_EXTRA_CREDENTIAL] == "foo=changed" class RetryRecorder: + """Transport handler counting how often it was hit, raising a canned error + or returning a freshly built response on every call.""" + def __init__(self, error=None, result=None): - self.__name__ = "RetryRecorder" self._retry_count = 0 self._error = error self._result = result - def __call__(self, *args, **kwargs): + def __call__(self, request: httpx2.Request) -> httpx2.Response: self._retry_count += 1 if self._error is not None: raise self._error - if self._result is not None: - return self._result + return self._result() @property def retry_count(self): return self._retry_count -@pytest.mark.parametrize( - "auth_class, retry_exception_class", - [ - pytest.param(KerberosAuthentication, KerberosExchangeError, marks=requires_kerberos), - pytest.param(GSSAPIAuthentication, SPNEGOExchangeError, marks=requires_gssapi), - ] -) -def test_authentication_fail_retry(auth_class, retry_exception_class, monkeypatch): - post_retry = RetryRecorder(error=retry_exception_class()) - monkeypatch.setattr(TrinoRequest.http.Session, "post", post_retry) - - get_retry = RetryRecorder(error=retry_exception_class()) - monkeypatch.setattr(TrinoRequest.http.Session, "get", get_retry) +@pytest.mark.parametrize("auth_class", [ + pytest.param(KerberosAuthentication, marks=requires_gssapi), + pytest.param(GSSAPIAuthentication, marks=requires_gssapi), +]) +def test_authentication_fail_retry(auth_class): + recorder = RetryRecorder(error=SPNEGOExchangeError("token exchange failed")) + auth = auth_class() attempts = 3 - kerberos_auth = auth_class() req = TrinoRequest( host="coordinator", port=8080, @@ -1045,17 +995,19 @@ def test_authentication_fail_retry(auth_class, retry_exception_class, monkeypatc user="test", ), http_scheme=constants.HTTPS, - auth=kerberos_auth, + http_session=TrinoRequest.create_http_client( + auth=auth, transport=httpx2.MockTransport(recorder)), + auth=auth, max_attempts=attempts, ) - with pytest.raises(retry_exception_class): + with pytest.raises(SPNEGOExchangeError): req.post("URL") - assert post_retry.retry_count == attempts + assert recorder.retry_count == attempts - with pytest.raises(retry_exception_class): - req.get("URL") - assert post_retry.retry_count == attempts + with pytest.raises(SPNEGOExchangeError): + req.get(req.statement_url) + assert recorder.retry_count == 2 * attempts @pytest.mark.parametrize("status_code, attempts", [ @@ -1063,15 +1015,8 @@ def test_authentication_fail_retry(auth_class, retry_exception_class, monkeypatc (503, 3), (504, 3), ]) -def test_5XX_error_retry(status_code, attempts, monkeypatch): - http_resp = TrinoRequest.http.Response() - http_resp.status_code = status_code - - post_retry = RetryRecorder(result=http_resp) - monkeypatch.setattr(TrinoRequest.http.Session, "post", post_retry) - - get_retry = RetryRecorder(result=http_resp) - monkeypatch.setattr(TrinoRequest.http.Session, "get", get_retry) +def test_5XX_error_retry(status_code, attempts): + recorder = RetryRecorder(result=lambda: httpx2.Response(status_code)) req = TrinoRequest( host="coordinator", @@ -1079,26 +1024,19 @@ def test_5XX_error_retry(status_code, attempts, monkeypatch): client_session=ClientSession( user="test", ), + http_session=httpx2.Client(transport=httpx2.MockTransport(recorder)), max_attempts=attempts ) req.post("URL") - assert post_retry.retry_count == attempts + assert recorder.retry_count == attempts - req.get("URL") - assert post_retry.retry_count == attempts + req.get(req.statement_url) + assert recorder.retry_count == 2 * attempts -def test_429_error_retry(monkeypatch): - http_resp = TrinoRequest.http.Response() - http_resp.status_code = 429 - http_resp.headers["Retry-After"] = 1 - - post_retry = RetryRecorder(result=http_resp) - monkeypatch.setattr(TrinoRequest.http.Session, "post", post_retry) - - get_retry = RetryRecorder(result=http_resp) - monkeypatch.setattr(TrinoRequest.http.Session, "get", get_retry) +def test_429_error_retry(): + recorder = RetryRecorder(result=lambda: httpx2.Response(429, headers={"Retry-After": "1"})) req = TrinoRequest( host="coordinator", @@ -1106,54 +1044,41 @@ def test_429_error_retry(monkeypatch): client_session=ClientSession( user="test", ), + http_session=httpx2.Client(transport=httpx2.MockTransport(recorder)), max_attempts=3 ) req.post("URL") - assert post_retry.retry_count == 3 - - req.get("URL") - assert post_retry.retry_count == 3 + assert recorder.retry_count == 3 + req.get(req.statement_url) + assert recorder.retry_count == 6 -def test_empty_200_response_retry(monkeypatch): - http_resp = TrinoRequest.http.Response() - http_resp.status_code = 200 - http_resp._content = b"" - post_retry = RetryRecorder(result=http_resp) - monkeypatch.setattr(TrinoRequest.http.Session, "post", post_retry) - - get_retry = RetryRecorder(result=http_resp) - monkeypatch.setattr(TrinoRequest.http.Session, "get", get_retry) +def test_empty_200_response_retry(): + recorder = RetryRecorder(result=lambda: httpx2.Response(200, content=b"")) attempts = 3 req = TrinoRequest( host="coordinator", port=8080, client_session=ClientSession(user="test"), + http_session=httpx2.Client(transport=httpx2.MockTransport(recorder)), max_attempts=attempts, ) req.post("SELECT 1") - assert post_retry.retry_count == attempts + assert recorder.retry_count == attempts - req.get("URL") - assert get_retry.retry_count == attempts + req.get(req.statement_url) + assert recorder.retry_count == 2 * attempts @pytest.mark.parametrize("status_code", [ 501 ]) -def test_error_no_retry(status_code, monkeypatch): - http_resp = TrinoRequest.http.Response() - http_resp.status_code = status_code - - post_retry = RetryRecorder(result=http_resp) - monkeypatch.setattr(TrinoRequest.http.Session, "post", post_retry) - - get_retry = RetryRecorder(result=http_resp) - monkeypatch.setattr(TrinoRequest.http.Session, "get", get_retry) +def test_error_no_retry(status_code): + recorder = RetryRecorder(result=lambda: httpx2.Response(status_code)) req = TrinoRequest( host="coordinator", @@ -1161,32 +1086,15 @@ def test_error_no_retry(status_code, monkeypatch): client_session=ClientSession( user="test", ), + http_session=httpx2.Client(transport=httpx2.MockTransport(recorder)), max_attempts=3, ) req.post("URL") - assert post_retry.retry_count == 1 - - req.get("URL") - assert post_retry.retry_count == 1 + assert recorder.retry_count == 1 - -class FakeGatewayResponse: - def __init__(self, http_response, redirect_count=1): - self.__name__ = "FakeGatewayResponse" - self.http_response = http_response - self.redirect_count = redirect_count - self.count = 0 - - def __call__(self, *args, **kwargs): - self.count += 1 - if self.count == self.redirect_count: - return self.http_response - http_response = TrinoRequest.http.Response() - http_response.status_code = 301 - http_response.headers["Location"] = "http://1.2.3.4:8080/new-path/" - assert http_response.is_redirect - return http_response + req.get(req.statement_url) + assert recorder.retry_count == 2 def test_trino_query_response_headers(sample_get_response_data): @@ -1195,20 +1103,6 @@ def test_trino_query_response_headers(sample_get_response_data): that are pass the the provided request instance post function call and it returns a `TrinoResult` instance. """ - - class MockResponse(mock.Mock): - # Fake response class - @property - def headers(self): - return { - 'X-Trino-Fake-1': 'one', - 'X-Trino-Fake-2': 'two', - } - - @property - def text(self): - return json.dumps(sample_get_response_data) - req = TrinoRequest( host="coordinator", port=8080, @@ -1228,20 +1122,33 @@ def text(self): constants.HEADER_CLIENT_CAPABILITIES: 'PARAMETRIC_DATETIME,SESSION_AUTHORIZATION' } - # Patch the post function to avoid making the requests, as well as to - # validate that the function was called with the right arguments. - with mock.patch.object(req, 'post', return_value=MockResponse()) as mock_post: - query = TrinoQuery( - request=req, - query=sql - ) - result = query.execute(additional_http_headers=additional_headers) + # Replace the post function with a hand-written fake to avoid making + # requests, and to validate it was called with the right arguments. + calls = [] + fake_response = _json_response( + sample_get_response_data, + headers={ + 'X-Trino-Fake-1': 'one', + 'X-Trino-Fake-2': 'two', + }) - # Validate the the post function was called with the right argguments - mock_post.assert_called_once_with(sql, additional_headers) + def fake_post(sql, additional_http_headers=None): + calls.append((sql, additional_http_headers)) + return fake_response - # Validate the result is an instance of TrinoResult - assert isinstance(result, TrinoResult) + req.post = fake_post + + query = TrinoQuery( + request=req, + query=sql + ) + result = query.execute(additional_http_headers=additional_headers) + + # Validate the the post function was called with the right argguments + assert calls == [(sql, additional_headers)] + + # Validate the result is an instance of TrinoResult + assert isinstance(result, TrinoResult) def test_stats_callback_cannot_mutate_query_stats(): @@ -1344,8 +1251,7 @@ def assert_headers_with_roles(headers: Dict[str, str], roles: Optional[str]): assert headers[constants.HEADER_USER] == "test_user" -def test_request_headers_role_hive_all(mock_get_and_post): - get, post = mock_get_and_post +def test_request_headers_role_hive_all(trino_server): req = TrinoRequest( host="coordinator", port=8080, @@ -1353,20 +1259,17 @@ def test_request_headers_role_hive_all(mock_get_and_post): user="test_user", roles={"hive": "ALL"} ), + http_session=trino_server.client(), ) req.post("URL") - _, post_kwargs = post.call_args - assert_headers_with_roles(post_kwargs["headers"], "hive=ALL") - - req.get("URL") - _, get_kwargs = get.call_args - assert_headers_with_roles(post_kwargs["headers"], "hive=ALL") + assert_headers_with_roles(trino_server.last_request.headers, "hive=ALL") + req.get(req.statement_url) + assert_headers_with_roles(trino_server.last_request.headers, "hive=ALL") -def test_request_headers_role_admin(mock_get_and_post): - get, post = mock_get_and_post +def test_request_headers_role_admin(trino_server): req = TrinoRequest( host="coordinator", port=8080, @@ -1374,21 +1277,18 @@ def test_request_headers_role_admin(mock_get_and_post): user="test_user", roles={"system": "admin"} ), + http_session=trino_server.client(), ) roles = "system=" + urllib.parse.quote("ROLE{admin}") req.post("URL") - _, post_kwargs = post.call_args - assert_headers_with_roles(post_kwargs["headers"], roles) + assert_headers_with_roles(trino_server.last_request.headers, roles) - req.get("URL") - _, get_kwargs = get.call_args - assert_headers_with_roles(post_kwargs["headers"], roles) + req.get(req.statement_url) + assert_headers_with_roles(trino_server.last_request.headers, roles) -def test_request_headers_role_empty(mock_get_and_post): - get, post = mock_get_and_post - +def test_request_headers_role_empty(trino_server): req = TrinoRequest( host="coordinator", port=8080, @@ -1396,24 +1296,21 @@ def test_request_headers_role_empty(mock_get_and_post): user="test_user", roles=None, ), + http_session=trino_server.client(), ) req.post("URL") - _, post_kwargs = post.call_args - assert_headers_with_roles(post_kwargs["headers"], None) + assert_headers_with_roles(trino_server.last_request.headers, None) - req.get("URL") - _, get_kwargs = get.call_args - assert_headers_with_roles(post_kwargs["headers"], None) + req.get(req.statement_url) + assert_headers_with_roles(trino_server.last_request.headers, None) def assert_headers_timezone(headers: Dict[str, str], timezone: str): assert headers[constants.HEADER_TIMEZONE] == timezone -def test_request_headers_with_timezone(mock_get_and_post): - get, post = mock_get_and_post - +def test_request_headers_with_timezone(trino_server): req = TrinoRequest( host="coordinator", port=8080, @@ -1421,39 +1318,35 @@ def test_request_headers_with_timezone(mock_get_and_post): user="test_user", timezone="Europe/Brussels" ), + http_session=trino_server.client(), ) req.post("URL") - _, post_kwargs = post.call_args - assert_headers_timezone(post_kwargs["headers"], "Europe/Brussels") - - req.get("URL") - _, get_kwargs = get.call_args - assert_headers_timezone(post_kwargs["headers"], "Europe/Brussels") + assert_headers_timezone(trino_server.last_request.headers, "Europe/Brussels") + req.get(req.statement_url) + assert_headers_timezone(trino_server.last_request.headers, "Europe/Brussels") -def test_request_headers_without_timezone(mock_get_and_post): - get, post = mock_get_and_post +def test_request_headers_without_timezone(trino_server): req = TrinoRequest( host="coordinator", port=8080, client_session=ClientSession( user="test_user", ), + http_session=trino_server.client(), ) localzone = get_localzone_name() req.post("URL") - _, post_kwargs = post.call_args - assert_headers_timezone(post_kwargs["headers"], localzone) + assert_headers_timezone(trino_server.last_request.headers, localzone) - req.get("URL") - _, get_kwargs = get.call_args - assert_headers_timezone(post_kwargs["headers"], localzone) + req.get(req.statement_url) + assert_headers_timezone(trino_server.last_request.headers, localzone) -def test_request_with_invalid_timezone(mock_get_and_post): +def test_request_with_invalid_timezone(): with pytest.raises(ZoneInfoNotFoundError) as zinfo_error: TrinoRequest( host="coordinator", @@ -1539,52 +1432,50 @@ def delete_password(self, servicename, username): os.remove(file_path) -def test_trino_request_headers_encoding_default_behavior(): +def test_trino_request_headers_encoding_default_behavior(monkeypatch): session = ClientSession(user="test", encoding=None) # Case 1: Both available -> No header - with mock.patch("trino.client.CODECS_UNAVAILABLE", {}): - req = TrinoRequest("host", 8080, session) - assert constants.HEADER_ENCODING not in req.http_headers + monkeypatch.setattr("trino._protocol.CODECS_UNAVAILABLE", {}) + req = TrinoRequest("host", 8080, session) + assert constants.HEADER_ENCODING not in req.http_headers # Case 2: Zstd missing -> Header set with json+lz4,json - with mock.patch("trino.client.CODECS_UNAVAILABLE", {"zstd": "Not installed"}): - req = TrinoRequest("host", 8080, session) - assert req.http_headers[constants.HEADER_ENCODING] == "json+lz4,json" + monkeypatch.setattr("trino._protocol.CODECS_UNAVAILABLE", {"zstd": "Not installed"}) + req = TrinoRequest("host", 8080, session) + assert req.http_headers[constants.HEADER_ENCODING] == "json+lz4,json" # Case 3: Lz4 missing -> Header set with json+zstd,json - with mock.patch("trino.client.CODECS_UNAVAILABLE", {"lz4": "Not installed"}): - req = TrinoRequest("host", 8080, session) - assert req.http_headers[constants.HEADER_ENCODING] == "json+zstd,json" + monkeypatch.setattr("trino._protocol.CODECS_UNAVAILABLE", {"lz4": "Not installed"}) + req = TrinoRequest("host", 8080, session) + assert req.http_headers[constants.HEADER_ENCODING] == "json+zstd,json" # Case 4: Both missing -> Header set with json - with mock.patch("trino.client.CODECS_UNAVAILABLE", {"lz4": "Not installed", "zstd": "Not installed"}): - req = TrinoRequest("host", 8080, session) - assert req.http_headers[constants.HEADER_ENCODING] == "json" + monkeypatch.setattr("trino._protocol.CODECS_UNAVAILABLE", {"lz4": "Not installed", "zstd": "Not installed"}) + req = TrinoRequest("host", 8080, session) + assert req.http_headers[constants.HEADER_ENCODING] == "json" -def test_decoder_factory_raises_with_message_on_missing_zstd(): - mapper = mock.Mock() - factory = CompressedQueryDataDecoderFactory(mapper) +def test_decoder_factory_raises_with_message_on_missing_zstd(monkeypatch): + factory = CompressedQueryDataDecoderFactory(object()) error_message = "No module named 'zstandard'" - with mock.patch("trino.client.CODECS_UNAVAILABLE", {"zstd": error_message}): - with pytest.raises( - ValueError, - match=f"zstd is not installed so json\\+zstd encoding is not supported: {error_message}" - ): - factory.create("json+zstd") + monkeypatch.setattr("trino._protocol.CODECS_UNAVAILABLE", {"zstd": error_message}) + with pytest.raises( + ValueError, + match=f"zstd is not installed so json\\+zstd encoding is not supported: {error_message}" + ): + factory.create("json+zstd") -def test_decoder_factory_raises_with_message_on_missing_lz4(): - mapper = mock.Mock() - factory = CompressedQueryDataDecoderFactory(mapper) +def test_decoder_factory_raises_with_message_on_missing_lz4(monkeypatch): + factory = CompressedQueryDataDecoderFactory(object()) error_message = "No module named 'lz4.block'" - with mock.patch("trino.client.CODECS_UNAVAILABLE", {"lz4": error_message}): - with pytest.raises( - ValueError, - match=f"lz4 is not installed so json\\+lz4 encoding is not supported: {error_message}" - ): - factory.create("json+lz4") + monkeypatch.setattr("trino._protocol.CODECS_UNAVAILABLE", {"lz4": error_message}) + with pytest.raises( + ValueError, + match=f"lz4 is not installed so json\\+lz4 encoding is not supported: {error_message}" + ): + factory.create("json+lz4") class _FinishedQuery: @@ -1679,17 +1570,17 @@ def fetch(self): next(it) -@httprettified def test_execute_drains_spooled_update_query_with_trailing_page(): """An update statement whose count row arrives over the spooling protocol still has a trailing page to drain. The rows are a lazy iterator at that point, so the drain loop must chain onto it rather than use list concatenation. """ query_id = "20210817_140827_00000_arvdv" - statement_path = f"{SERVER_ADDRESS}{constants.URL_STATEMENT_PATH}" + server_address = "https://coordinator" + statement_path = constants.URL_STATEMENT_PATH def statement_uri(token): - return f"{statement_path}/{query_id}/{token}" + return f"{server_address}{statement_path}/{query_id}/{token}" columns = [{ "name": "rows", @@ -1702,13 +1593,13 @@ def statement_uri(token): post_response = { "id": query_id, "nextUri": statement_uri(1), - "infoUri": f"{SERVER_ADDRESS}/query.html?{query_id}", + "infoUri": f"{server_address}/query.html?{query_id}", "stats": {"state": "QUEUED"}, } spooled_update_response = { "id": query_id, "nextUri": statement_uri(2), - "infoUri": f"{SERVER_ADDRESS}/query.html?{query_id}", + "infoUri": f"{server_address}/query.html?{query_id}", "updateType": "INSERT", "updateCount": 3, "columns": columns, @@ -1725,22 +1616,24 @@ def statement_uri(token): # Trailing page with no data. It only moves the query to a terminal state. final_response = { "id": query_id, - "infoUri": f"{SERVER_ADDRESS}/query.html?{query_id}", + "infoUri": f"{server_address}/query.html?{query_id}", "updateType": "INSERT", "updateCount": 3, "columns": columns, "stats": {"state": "FINISHED"}, } - httpretty.register_uri(method=httpretty.POST, uri=statement_path, body=json.dumps(post_response)) - httpretty.register_uri(method=httpretty.GET, uri=statement_uri(1), body=json.dumps(spooled_update_response)) - httpretty.register_uri(method=httpretty.GET, uri=statement_uri(2), body=json.dumps(final_response)) + server = MockTrinoServer() + server.register("POST", statement_path, json=post_response) + server.register("GET", f"{statement_path}/{query_id}/1", json=spooled_update_response) + server.register("GET", f"{statement_path}/{query_id}/2", json=final_response) request = TrinoRequest( host="coordinator", port=constants.DEFAULT_TLS_PORT, client_session=ClientSession(user="test", encoding="json"), http_scheme=constants.HTTPS, + http_session=server.client(), ) query = TrinoQuery(request, query="INSERT INTO some_table VALUES (1), (2), (3)") diff --git a/tests/unit/test_client_spooling.py b/tests/unit/test_client_spooling.py index 21912d37..df4d48e2 100644 --- a/tests/unit/test_client_spooling.py +++ b/tests/unit/test_client_spooling.py @@ -37,7 +37,7 @@ def _mock_trino_request(): def _head_response(status_code): - return mock.Mock(status_code=status_code, ok=(200 <= status_code < 300)) + return mock.Mock(status_code=status_code, is_error=(400 <= status_code < 600)) @pytest.fixture @@ -165,7 +165,7 @@ def _spooled_fetch_response(): """Minimal spooled protocol GET response JSON.""" resp = mock.Mock() resp.status_code = 200 - resp.ok = True + resp.is_error = False resp.headers = {} resp.text = json.dumps({ "id": "q1", @@ -285,7 +285,7 @@ def test_send_spooling_request_forwards_custom_headers_to_coordinator(): def fake_get(uri, headers=None, **kwargs): recorded["headers"] = headers - return mock.Mock(ok=True) + return mock.Mock(is_error=False) segment._request._get = fake_get segment._send_spooling_request(segment.uri) @@ -302,7 +302,7 @@ def test_send_spooling_request_does_not_forward_custom_headers_to_external_stora def fake_get(uri, headers=None, **kwargs): recorded["headers"] = headers - return mock.Mock(ok=True) + return mock.Mock(is_error=False) segment._request._get = fake_get external_uri = "https://s3.amazonaws.com/bucket/seg1?X-Amz-Signature=abc" @@ -321,7 +321,7 @@ def test_send_spooling_request_segment_header_takes_precedence_over_custom_heade def fake_get(uri, headers=None, **kwargs): recorded["headers"] = headers - return mock.Mock(ok=True) + return mock.Mock(is_error=False) segment._request._get = fake_get segment._send_spooling_request(segment.uri) diff --git a/tests/unit/test_dbapi.py b/tests/unit/test_dbapi.py index 47689336..cca09278 100644 --- a/tests/unit/test_dbapi.py +++ b/tests/unit/test_dbapi.py @@ -9,17 +9,15 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -import json import threading import uuid from unittest.mock import patch -import httpretty +import httpx2 import pytest -from httpretty import httprettified -from requests import Session import trino.exceptions +from tests.unit.mock_http import MockTrinoServer from tests.unit.oauth_test_utils import _get_token_requests from tests.unit.oauth_test_utils import _post_statement_requests from tests.unit.oauth_test_utils import GetTokenCallback @@ -27,6 +25,7 @@ from tests.unit.oauth_test_utils import REDIRECT_RESOURCE from tests.unit.oauth_test_utils import RedirectHandler from tests.unit.oauth_test_utils import SERVER_ADDRESS +from tests.unit.oauth_test_utils import TOKEN_PATH from tests.unit.oauth_test_utils import TOKEN_RESOURCE from trino import constants from trino.auth import BasicAuthentication @@ -36,34 +35,36 @@ from trino.dbapi import Connection from trino.dbapi import Cursor +_QUERY_ID = "20210817_140827_00000_arvdv" -@patch("trino.dbapi.trino.client") -def test_http_session_is_correctly_passed_in(mock_client): - # GIVEN - test_session = Session() - test_session.proxies = {"http": "some.http.proxy:81", "https": "some.http.proxy:81"} - # WHEN - with connect("sample_trino_cluster:443", http_session=test_session) as conn: - conn.cursor().execute("SOME FAKE QUERY") +def _finished_statement_response(): + """A POST /v1/statement response without a nextUri: the query is + immediately in a terminal state, so no follow-up requests are needed.""" + return { + "id": _QUERY_ID, + "infoUri": f"{SERVER_ADDRESS}/query.html?{_QUERY_ID}", + "stats": {"state": "FINISHED"}, + } - # THEN - request_args, _ = mock_client.TrinoRequest.call_args - assert test_session in request_args +def test_http_session_is_correctly_passed_in(trino_server): + trino_server.register("POST", constants.URL_STATEMENT_PATH, json=_finished_statement_response()) + test_session = trino_server.client() -@patch("trino.dbapi.trino.client") -def test_http_session_is_defaulted_when_not_specified(mock_client): - # WHEN - with connect("sample_trino_cluster:443") as conn: + with connect("https://sample_trino_cluster:443", user="test", http_session=test_session) as conn: conn.cursor().execute("SOME FAKE QUERY") + assert conn._http_session is test_session - # THEN - request_args, _ = mock_client.TrinoRequest.call_args - assert mock_client.TrinoRequest.http.Session.return_value in request_args + assert len(trino_server.requests(method="POST", path=constants.URL_STATEMENT_PATH)) == 1 + + +def test_http_session_is_defaulted_when_not_specified(): + with connect("sample_trino_cluster:443") as conn: + assert isinstance(conn._http_session, httpx2.Client) + assert conn._create_request()._http_session is conn._http_session -@httprettified def test_token_retrieved_once_per_auth_instance(sample_post_response_data, sample_get_response_data): token = str(uuid.uuid4()) challenge_id = str(uuid.uuid4()) @@ -71,27 +72,20 @@ def test_token_retrieved_once_per_auth_instance(sample_post_response_data, sampl redirect_server = f"{REDIRECT_RESOURCE}/{challenge_id}" token_server = f"{TOKEN_RESOURCE}/{challenge_id}" + server = MockTrinoServer() + post_statement_callback = PostStatementCallback(redirect_server, token_server, [token], sample_post_response_data) get_statement_callback = PostStatementCallback(redirect_server, token_server, [token], sample_get_response_data) # bind post statement to submit query - httpretty.register_uri( - method=httpretty.POST, - uri=f"{SERVER_ADDRESS}{constants.URL_STATEMENT_PATH}", - body=post_statement_callback) + server.register("POST", constants.URL_STATEMENT_PATH, post_statement_callback) # bind get statement for result retrieval - httpretty.register_uri( - method=httpretty.GET, - uri=f"{SERVER_ADDRESS}{constants.URL_STATEMENT_PATH}/20210817_140827_00000_arvdv/1", - body=get_statement_callback) + server.register("GET", f"{constants.URL_STATEMENT_PATH}/{_QUERY_ID}/1", get_statement_callback) # bind get token get_token_callback = GetTokenCallback(token_server, token) - httpretty.register_uri( - method=httpretty.GET, - uri=token_server, - body=get_token_callback) + server.register("GET", f"/{TOKEN_PATH}/{challenge_id}", get_token_callback) redirect_handler = RedirectHandler() @@ -99,7 +93,8 @@ def test_token_retrieved_once_per_auth_instance(sample_post_response_data, sampl "coordinator", user="test", auth=OAuth2Authentication(redirect_auth_url_handler=redirect_handler), - http_scheme=constants.HTTPS + http_scheme=constants.HTTPS, + http_session=server.client(), ) as conn: conn.cursor().execute("SELECT 1") conn.cursor().execute("SELECT 2") @@ -107,10 +102,7 @@ def test_token_retrieved_once_per_auth_instance(sample_post_response_data, sampl # bind get token get_token_callback = GetTokenCallback(token_server, token) - httpretty.register_uri( - method=httpretty.GET, - uri=token_server, - body=get_token_callback) + server.register("GET", f"/{TOKEN_PATH}/{challenge_id}", get_token_callback) redirect_handler = RedirectHandler() @@ -118,16 +110,16 @@ def test_token_retrieved_once_per_auth_instance(sample_post_response_data, sampl "coordinator", user="test", auth=OAuth2Authentication(redirect_auth_url_handler=redirect_handler), - http_scheme=constants.HTTPS + http_scheme=constants.HTTPS, + http_session=server.client(), ) as conn2: conn2.cursor().execute("SELECT 1") conn2.cursor().execute("SELECT 2") conn2.cursor().execute("SELECT 3") - assert len(_get_token_requests(challenge_id)) == 1 + assert len(_get_token_requests(server, challenge_id)) == 1 -@httprettified def test_token_retrieved_once_when_authentication_instance_is_shared(sample_post_response_data, sample_get_response_data): token = str(uuid.uuid4()) @@ -136,27 +128,20 @@ def test_token_retrieved_once_when_authentication_instance_is_shared(sample_post redirect_server = f"{REDIRECT_RESOURCE}/{challenge_id}" token_server = f"{TOKEN_RESOURCE}/{challenge_id}" + server = MockTrinoServer() + post_statement_callback = PostStatementCallback(redirect_server, token_server, [token], sample_post_response_data) get_statement_callback = PostStatementCallback(redirect_server, token_server, [token], sample_get_response_data) # bind post statement to submit query - httpretty.register_uri( - method=httpretty.POST, - uri=f"{SERVER_ADDRESS}{constants.URL_STATEMENT_PATH}", - body=post_statement_callback) + server.register("POST", constants.URL_STATEMENT_PATH, post_statement_callback) # bind get statement for result retrieval - httpretty.register_uri( - method=httpretty.GET, - uri=f"{SERVER_ADDRESS}{constants.URL_STATEMENT_PATH}/20210817_140827_00000_arvdv/1", - body=get_statement_callback) + server.register("GET", f"{constants.URL_STATEMENT_PATH}/{_QUERY_ID}/1", get_statement_callback) # bind get token get_token_callback = GetTokenCallback(token_server, token) - httpretty.register_uri( - method=httpretty.GET, - uri=token_server, - body=get_token_callback) + server.register("GET", f"/{TOKEN_PATH}/{challenge_id}", get_token_callback) redirect_handler = RedirectHandler() @@ -166,7 +151,8 @@ def test_token_retrieved_once_when_authentication_instance_is_shared(sample_post "coordinator", user="test", auth=authentication, - http_scheme=constants.HTTPS + http_scheme=constants.HTTPS, + http_session=server.client(), ) as conn: conn.cursor().execute("SELECT 1") conn.cursor().execute("SELECT 2") @@ -174,26 +160,23 @@ def test_token_retrieved_once_when_authentication_instance_is_shared(sample_post # bind get token get_token_callback = GetTokenCallback(token_server, token) - httpretty.register_uri( - method=httpretty.GET, - uri=token_server, - body=get_token_callback) + server.register("GET", f"/{TOKEN_PATH}/{challenge_id}", get_token_callback) with connect( "coordinator", user="test", auth=authentication, - http_scheme=constants.HTTPS + http_scheme=constants.HTTPS, + http_session=server.client(), ) as conn2: conn2.cursor().execute("SELECT 1") conn2.cursor().execute("SELECT 2") conn2.cursor().execute("SELECT 3") - assert len(_post_statement_requests()) == 7 - assert len(_get_token_requests(challenge_id)) == 1 + assert len(_post_statement_requests(server)) == 7 + assert len(_get_token_requests(server, challenge_id)) == 1 -@httprettified def test_token_retrieved_once_when_multithreaded(sample_post_response_data, sample_get_response_data): token = str(uuid.uuid4()) challenge_id = str(uuid.uuid4()) @@ -201,27 +184,20 @@ def test_token_retrieved_once_when_multithreaded(sample_post_response_data, samp redirect_server = f"{REDIRECT_RESOURCE}/{challenge_id}" token_server = f"{TOKEN_RESOURCE}/{challenge_id}" + server = MockTrinoServer() + post_statement_callback = PostStatementCallback(redirect_server, token_server, [token], sample_post_response_data) get_statement_callback = PostStatementCallback(redirect_server, token_server, [token], sample_get_response_data) # bind post statement to submit query - httpretty.register_uri( - method=httpretty.POST, - uri=f"{SERVER_ADDRESS}{constants.URL_STATEMENT_PATH}", - body=post_statement_callback) + server.register("POST", constants.URL_STATEMENT_PATH, post_statement_callback) # bind get statement for result retrieval - httpretty.register_uri( - method=httpretty.GET, - uri=f"{SERVER_ADDRESS}{constants.URL_STATEMENT_PATH}/20210817_140827_00000_arvdv/1", - body=get_statement_callback) + server.register("GET", f"{constants.URL_STATEMENT_PATH}/{_QUERY_ID}/1", get_statement_callback) # bind get token get_token_callback = GetTokenCallback(token_server, token) - httpretty.register_uri( - method=httpretty.GET, - uri=token_server, - body=get_token_callback) + server.register("GET", f"/{TOKEN_PATH}/{challenge_id}", get_token_callback) redirect_handler = RedirectHandler() @@ -231,7 +207,8 @@ def test_token_retrieved_once_when_multithreaded(sample_post_response_data, samp "coordinator", user="test", auth=authentication, - http_scheme=constants.HTTPS + http_scheme=constants.HTTPS, + http_session=server.client(), ) class RunningThread(threading.Thread): @@ -256,27 +233,19 @@ def run(self) -> None: for thread in threads: thread.join() - assert len(_get_token_requests(challenge_id)) == 1 + assert len(_get_token_requests(server, challenge_id)) == 1 -@patch("trino.dbapi.trino.client") -def test_tags_are_set_when_specified(mock_client): +def test_tags_are_set_when_specified(): client_tags = ["TAG1", "TAG2"] with connect("sample_trino_cluster:443", client_tags=client_tags) as conn: - conn.cursor().execute("SOME FAKE QUERY") - - _, passed_client_tags = mock_client.ClientSession.call_args - assert passed_client_tags["client_tags"] == client_tags + assert conn._client_session.client_tags == client_tags -@patch("trino.dbapi.trino.client") -def test_role_is_set_when_specified(mock_client): +def test_role_is_set_when_specified(): roles = {"system": "finance"} with connect("sample_trino_cluster:443", roles=roles) as conn: - conn.cursor().execute("SOME FAKE QUERY") - - _, passed_role = mock_client.ClientSession.call_args - assert passed_role["roles"] == roles + assert conn._client_session.roles == {"system": "ROLE{finance}"} def test_hostname_parsing(): @@ -398,11 +367,14 @@ def test_no_error_when_auth_over_http_with_allow_insecure_auth(): assert request._http_scheme == constants.HTTP +def _statement_path(query_id, token): + return f"{constants.URL_STATEMENT_PATH}/{query_id}/{token}" + + def _statement_uri(query_id, token): - return f"{SERVER_ADDRESS}{constants.URL_STATEMENT_PATH}/{query_id}/{token}" + return f"{SERVER_ADDRESS}{_statement_path(query_id, token)}" -@httprettified def test_cursor_close_does_not_cancel_finished_update_query(): """Regression test for https://github.com/trinodb/trino-python-client/issues/601 @@ -411,8 +383,7 @@ def test_cursor_close_does_not_cancel_finished_update_query(): cursor without fetching must drain that nextUri instead of issuing a DELETE, otherwise the already-completed statement is reported as USER_CANCELED. """ - query_id = "20210817_140827_00000_arvdv" - statement_path = f"{SERVER_ADDRESS}{constants.URL_STATEMENT_PATH}" + query_id = _QUERY_ID post_response = { "id": query_id, @@ -446,12 +417,13 @@ def test_cursor_close_does_not_cancel_finished_update_query(): "stats": {"state": "FINISHED"}, } - httpretty.register_uri(method=httpretty.POST, uri=statement_path, body=json.dumps(post_response)) - httpretty.register_uri(method=httpretty.GET, uri=_statement_uri(query_id, 1), body=json.dumps(update_response)) - httpretty.register_uri(method=httpretty.GET, uri=_statement_uri(query_id, 2), body=json.dumps(final_response)) - httpretty.register_uri(method=httpretty.DELETE, uri=_statement_uri(query_id, 2), status=204) + server = MockTrinoServer() + server.register("POST", constants.URL_STATEMENT_PATH, json=post_response) + server.register("GET", _statement_path(query_id, 1), json=update_response) + server.register("GET", _statement_path(query_id, 2), json=final_response) + server.register("DELETE", _statement_path(query_id, 2), status=204, text="") - with connect("coordinator", user="test", http_scheme=constants.HTTPS) as conn: + with connect("coordinator", user="test", http_scheme=constants.HTTPS, http_session=server.client()) as conn: cur = conn.cursor() cur.execute("INSERT INTO some_table VALUES (1), (2), (3)") # execute() must have drained the query to a terminal state. @@ -459,17 +431,14 @@ def test_cursor_close_does_not_cancel_finished_update_query(): assert cur.rowcount == 1000 cur.close() - delete_requests = [r for r in httpretty.latest_requests() if r.method == "DELETE"] - assert delete_requests == [], "closing a finished update query must not issue a cancel" + assert server.requests(method="DELETE") == [], "closing a finished update query must not issue a cancel" -@httprettified def test_cursor_close_cancels_unfinished_query(): """Closing a cursor whose result set has not been fully consumed must still cancel the running query so the server can free its resources. """ - query_id = "20210817_140827_00000_arvdv" - statement_path = f"{SERVER_ADDRESS}{constants.URL_STATEMENT_PATH}" + query_id = _QUERY_ID post_response = { "id": query_id, @@ -491,18 +460,18 @@ def test_cursor_close_cancels_unfinished_query(): "stats": {"state": "RUNNING"}, } - httpretty.register_uri(method=httpretty.POST, uri=statement_path, body=json.dumps(post_response)) - httpretty.register_uri(method=httpretty.GET, uri=_statement_uri(query_id, 1), body=json.dumps(data_response)) - httpretty.register_uri(method=httpretty.DELETE, uri=_statement_uri(query_id, 2), status=204) + server = MockTrinoServer() + server.register("POST", constants.URL_STATEMENT_PATH, json=post_response) + server.register("GET", _statement_path(query_id, 1), json=data_response) + server.register("DELETE", _statement_path(query_id, 2), status=204, text="") - with connect("coordinator", user="test", http_scheme=constants.HTTPS) as conn: + with connect("coordinator", user="test", http_scheme=constants.HTTPS, http_session=server.client()) as conn: cur = conn.cursor() cur.execute("SELECT x FROM some_table") assert cur._query.finished is False cur.close() - delete_requests = [r for r in httpretty.latest_requests() if r.method == "DELETE"] - assert len(delete_requests) == 1, "closing an unfinished query must cancel it" + assert len(server.requests(method="DELETE")) == 1, "closing an unfinished query must cancel it" @pytest.mark.parametrize( diff --git a/trino/_protocol.py b/trino/_protocol.py new file mode 100644 index 00000000..d14e68f3 --- /dev/null +++ b/trino/_protocol.py @@ -0,0 +1,1013 @@ +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +I/O-free pieces of the Trino HTTP protocol shared by the synchronous client +(:mod:`trino.client`) and the asynchronous client (:mod:`trino.aio`). + +Everything in this module is transport-agnostic: header construction and +parsing, response processing, retry decisions, session state, and segment +decoding. The actual HTTP calls live in the transport-specific modules. +""" +from __future__ import annotations + +import abc +import base64 +import copy +import random +import re +import threading +import urllib.parse +import warnings +from abc import abstractmethod +from collections.abc import Iterator +from collections.abc import Mapping +from collections.abc import MutableMapping +from dataclasses import dataclass +from datetime import datetime +from email.utils import parsedate_to_datetime +from enum import Enum +from typing import Any +from typing import Callable +from typing import cast +from typing import Dict +from typing import List +from typing import Literal +from typing import Optional +from typing import Tuple +from typing import TypedDict +from typing import Union +from zoneinfo import ZoneInfo + +try: + import lz4.block +except ImportError as err: + _LZ4_ERROR: Optional[str] = str(err) +else: + _LZ4_ERROR = None + +try: + import orjson as json +except ImportError: + import json + +try: + import zstandard +except ImportError as err: + _ZSTD_ERROR: Optional[str] = str(err) +else: + _ZSTD_ERROR = None + +import trino.logging +from trino import constants +from trino import exceptions +from trino._version import __version__ +from trino.exceptions import TrinoExternalError +from trino.exceptions import TrinoQueryError +from trino.exceptions import TrinoUserError +from trino.mapper import RowMapper +from trino.mapper import RowMapperFactory + +logger = trino.logging.get_logger(__name__) + +_HEADER_EXTRA_CREDENTIAL_KEY_REGEX = re.compile(r'^\S[^\s=]*$') + +ENCODINGS = ["json+zstd", "json+lz4", "json"] +CODECS_UNAVAILABLE = {} +if _LZ4_ERROR: + CODECS_UNAVAILABLE["lz4"] = _LZ4_ERROR +if _ZSTD_ERROR: + CODECS_UNAVAILABLE["zstd"] = _ZSTD_ERROR + +ROLE_PATTERN = re.compile(r"^ROLE\{(.*)\}$") + +# HTTP status codes that are retried because they are transient coordinator or +# gateway conditions: Too Many Requests, Bad Gateway, Service Unavailable and +# Gateway Timeout. +RETRIABLE_STATUS_CODES = (429, 502, 503, 504) + + +class CaseInsensitiveDict(MutableMapping[str, Any]): + """ + A case-insensitive ``dict``-like object for HTTP headers. + + Keys are matched case-insensitively but the case of the last key set is + preserved for iteration. Unlike ``httpx2.Headers`` it tolerates ``None`` + values, which the client uses to mean "header not sent"; ``None``-valued + entries are filtered out before a request hits the wire. + """ + + def __init__(self, data: Optional[Mapping[str, Any]] = None, **kwargs: Any) -> None: + # Maps lowercased key -> (actual key, value) + self._store: Dict[str, Tuple[str, Any]] = {} + if data is not None: + self.update(data) + self.update(kwargs) + + def __setitem__(self, key: str, value: Any) -> None: + self._store[key.lower()] = (key, value) + + def __getitem__(self, key: str) -> Any: + return self._store[key.lower()][1] + + def __delitem__(self, key: str) -> None: + del self._store[key.lower()] + + def __iter__(self) -> Iterator[str]: + return (original_key for original_key, _ in self._store.values()) + + def __len__(self) -> int: + return len(self._store) + + def __eq__(self, other: object) -> bool: + if isinstance(other, Mapping): + other_dict = {k.lower(): v for k, v in other.items()} + return {k: v for k, (_, v) in self._store.items()} == other_dict + return NotImplemented + + def __repr__(self) -> str: + return repr(dict(self.items())) + + def copy(self) -> CaseInsensitiveDict: + return CaseInsensitiveDict(dict(self.items())) + + +def wire_headers(headers: Mapping[str, Any]) -> Dict[str, str]: + """Drop ``None``-valued headers; only the remaining ones are sent on the wire.""" + return {key: value for key, value in headers.items() if value is not None} + + +def get_header_values(headers: Mapping[str, str], header: str) -> List[str]: + return [val.strip() for val in headers[header].split(",")] + + +def get_session_property_values(headers: Mapping[str, str], header: str) -> List[Tuple[str, str]]: + kvs = get_header_values(headers, header) + return [ + (k.strip(), urllib.parse.unquote_plus(v.strip())) + for k, v in (kv.split("=", 1) for kv in kvs if kv) + ] + + +def get_prepared_statement_values(headers: Mapping[str, str], header: str) -> List[Tuple[str, str]]: + kvs = get_header_values(headers, header) + return [ + (k.strip(), urllib.parse.unquote_plus(v.strip())) + for k, v in (kv.split("=", 1) for kv in kvs if kv) + ] + + +def get_roles_values(headers: Mapping[str, str], header: str) -> List[Tuple[str, str]]: + kvs = get_header_values(headers, header) + return [ + (k.strip(), urllib.parse.unquote_plus(v.strip())) + for k, v in (kv.split("=", 1) for kv in kvs if kv) + ] + + +class ClientSession: + """ + Manage the current Client Session properties of a specific connection. This class is thread-safe. + + :param user: associated with the query. It is useful for access control + and query scheduling. + :param authorization_user: associated with the query. It is useful for access control + and query scheduling. + :param source: associated with the query. It is useful for access + control and query scheduling. + :param catalog: to query. The *catalog* is associated with a Trino + connector. This variable sets the default catalog used + by SQL statements. For example, if *catalog* is set + to ``some_catalog``, the SQL statement + ``SELECT * FROM some_schema.some_table`` will actually + query the table + ``some_catalog.some_schema.some_table``. + :param schema: to query. The *schema* is a logical abstraction to group + table. This variable sets the default schema used by + SQL statements. For example, if *schema* is set to + ``some_schema``, the SQL statement + ``SELECT * FROM some_table`` will actually query the + table ``some_catalog.some_schema.some_table``. + :param properties: set specific Trino behavior for the current + session. Please refer to the output of + ``SHOW SESSION`` to check the available + properties. + :param headers: HTTP headers to POST/GET in the HTTP requests + :param extra_credential: extra credentials. as list of ``(key, value)`` + tuples. + :param client_tags: Client tags as list of strings. + :param roles: roles for the current session. Some connectors do not + support role management. See connector documentation for more details. + :param timezone: The timezone for query processing. Defaults to the system's local timezone. + :param encoding: The encoding for the spooling protocol. Defaults to None. + """ + + def __init__( + self, + user: str, + authorization_user: Optional[str] = None, + catalog: Optional[str] = None, + schema: Optional[str] = None, + source: Optional[str] = None, + properties: Optional[Dict[str, str]] = None, + headers: Optional[Dict[str, str]] = None, + transaction_id: Optional[str] = None, + extra_credential: Optional[List[Tuple[str, str]]] = None, + client_tags: Optional[List[str]] = None, + roles: Optional[Union[Dict[str, str], str]] = None, + timezone: Optional[str] = None, + encoding: Optional[Union[str, List[str]]] = None, + heartbeat_interval: Optional[float] = constants.DEFAULT_HEARTBEAT_INTERVAL, + ): + self._object_lock = threading.Lock() + self._prepared_statements: Dict[str, str] = {} + + self._user = user + self._authorization_user = authorization_user + self._catalog = catalog + self._schema = schema + self._source = source + self._properties = properties.copy() if properties is not None else {} + self._headers = headers.copy() if headers is not None else {} + self._transaction_id = transaction_id + self._extra_credential = extra_credential + self._client_tags = client_tags.copy() if client_tags is not None else list() + self._roles = self._format_roles(roles) if roles is not None else {} + if timezone: # Check timezone validity + ZoneInfo(timezone) + self._timezone = timezone + else: + from tzlocal import get_localzone_name + self._timezone = get_localzone_name() + self._encoding = encoding + self._heartbeat_interval = heartbeat_interval + + @property + def user(self) -> str: + return self._user + + @property + def authorization_user(self) -> Optional[str]: + with self._object_lock: + return self._authorization_user + + @authorization_user.setter + def authorization_user(self, authorization_user: Optional[str]) -> None: + with self._object_lock: + self._authorization_user = authorization_user + + @property + def catalog(self) -> Optional[str]: + with self._object_lock: + return self._catalog + + @catalog.setter + def catalog(self, catalog: Optional[str]) -> None: + with self._object_lock: + self._catalog = catalog + + @property + def schema(self) -> Optional[str]: + with self._object_lock: + return self._schema + + @schema.setter + def schema(self, schema: Optional[str]) -> None: + with self._object_lock: + self._schema = schema + + @property + def source(self) -> Optional[str]: + return self._source + + @property + def properties(self) -> Dict[str, str]: + with self._object_lock: + return self._properties + + @properties.setter + def properties(self, properties: Dict[str, str]) -> None: + with self._object_lock: + self._properties = properties + + @property + def headers(self) -> Dict[str, str]: + return self._headers + + @property + def transaction_id(self) -> Optional[str]: + with self._object_lock: + return self._transaction_id + + @transaction_id.setter + def transaction_id(self, transaction_id: Optional[str]) -> None: + with self._object_lock: + self._transaction_id = transaction_id + + @property + def extra_credential(self) -> Optional[List[Tuple[str, str]]]: + return self._extra_credential + + @property + def client_tags(self) -> List[str]: + return self._client_tags + + @property + def roles(self) -> Dict[str, str]: + with self._object_lock: + return self._roles + + @roles.setter + def roles(self, roles: Dict[str, str]) -> None: + with self._object_lock: + self._roles = roles + + @property + def prepared_statements(self) -> Dict[str, str]: + return self._prepared_statements + + @prepared_statements.setter + def prepared_statements(self, prepared_statements: Dict[str, str]) -> None: + with self._object_lock: + self._prepared_statements = prepared_statements + + @property + def timezone(self) -> str: + with self._object_lock: + return self._timezone + + @property + def encoding(self) -> Optional[Union[str, List[str]]]: + with self._object_lock: + return self._encoding + + @property + def heartbeat_interval(self) -> Optional[float]: + return self._heartbeat_interval + + @staticmethod + def _format_roles(roles: Union[Dict[str, str], str]) -> Dict[str, str]: + if isinstance(roles, str): + roles = {"system": roles} + formatted_roles = {} + for catalog, role in roles.items(): + is_legacy_role_pattern = ROLE_PATTERN.match(role) is not None + if role in ("NONE", "ALL") or is_legacy_role_pattern: + if is_legacy_role_pattern: + warnings.warn(f"A role '{role}' is provided using a legacy format. " + "Please remove the ROLE{} wrapping. Support for the legacy format might be " + "removed in a future release.", + DeprecationWarning) + formatted_roles[catalog] = role + else: + formatted_roles[catalog] = f"ROLE{{{role}}}" + return formatted_roles + + def __getstate__(self) -> Dict[str, Any]: + state = self.__dict__.copy() + del state["_object_lock"] + return state + + def __setstate__(self, state: Dict[str, Any]) -> None: + self.__dict__.update(state) + self._object_lock = threading.Lock() + + +@dataclass +class TrinoStatus: + id: str + stats: Dict[str, str] + warnings: List[Any] + info_uri: str + next_uri: Optional[str] + update_type: Optional[str] + update_count: Optional[int] + rows: Union[List[Any], Dict[str, Any]] + columns: List[Any] + + def __repr__(self) -> str: + return ( + "TrinoStatus(" + "id={}, stats={{...}}, warnings={}, info_uri={}, next_uri={}, rows=" + ")".format( + self.id, + len(self.warnings), + self.info_uri, + self.next_uri, + len(self.rows), + ) + ) + + +class _DelayExponential: + def __init__( + self, + base: float = 0.1, # 100ms + exponent: float = 2, + jitter: bool = True, + max_delay: float = 1800, # 30 min + ) -> None: + self._base = base + self._exponent = exponent + self._jitter = jitter + self._max_delay = max_delay + + def __call__(self, attempt: int) -> float: + delay = float(self._base) * (self._exponent ** attempt) + if self._jitter: + delay *= random.random() + delay = min(float(self._max_delay), delay) + return delay + + +def _parse_retry_after_header(retry_after: Union[int, str]) -> float: + if isinstance(retry_after, int): + return retry_after + elif isinstance(retry_after, str) and retry_after.isdigit(): + return int(retry_after) + else: + retry_date = parsedate_to_datetime(retry_after) + now = datetime.utcnow() + return (retry_date - now).total_seconds() + + +def needs_retry(response: Any) -> bool: + """Retry decision for a response that did not raise: transient gateway + statuses, or a 200 with an empty body (transient under load).""" + status_code = getattr(response, "status_code", None) + if status_code in RETRIABLE_STATUS_CODES: + return True + return status_code == 200 and not getattr(response, "text", "").strip() + + +def retry_after_seconds(response: Any) -> Optional[float]: + """Server-mandated delay before retrying, from a 429's Retry-After header.""" + if getattr(response, "status_code", None) == 429 and "Retry-After" in response.headers: + return _parse_retry_after_header(response.headers.get("Retry-After")) + return None + + +class _TrinoRequestBase(abc.ABC): + """ + Transport-agnostic half of a Trino request: URL construction, protocol + header construction, and response processing. Subclasses provide the + actual HTTP verbs on top of an ``httpx2.Client`` or ``httpx2.AsyncClient``. + """ + + def __init__( + self, + host: str, + port: int, + client_session: ClientSession, + http_scheme: Optional[str] = None, + ) -> None: + self._client_session = client_session + self._host = 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 + else: + self._http_scheme = http_scheme + + @property + def transaction_id(self) -> Optional[str]: + return self._client_session.transaction_id + + @transaction_id.setter + def transaction_id(self, value: Optional[str]) -> None: + self._client_session.transaction_id = value + + @property + def http_headers(self) -> CaseInsensitiveDict: + headers = CaseInsensitiveDict() + + headers[constants.HEADER_CATALOG] = self._client_session.catalog + headers[constants.HEADER_SCHEMA] = self._client_session.schema + headers[constants.HEADER_SOURCE] = self._client_session.source + if self._client_session.authorization_user is not None: + headers[constants.HEADER_ORIGINAL_USER] = self._client_session.user + headers[constants.HEADER_USER] = self._client_session.authorization_user + else: + headers[constants.HEADER_USER] = self._client_session.user + headers[constants.HEADER_TIMEZONE] = self._client_session.timezone + if self._client_session.encoding is None: + if not CODECS_UNAVAILABLE: + pass + else: + encoding = [ + enc + for enc in ENCODINGS + if (enc.split("+")[1] if "+" in enc else None) not in CODECS_UNAVAILABLE + ] + headers[constants.HEADER_ENCODING] = ",".join(encoding) + elif isinstance(self._client_session.encoding, list): + headers[constants.HEADER_ENCODING] = ",".join(self._client_session.encoding) + elif isinstance(self._client_session.encoding, str): + headers[constants.HEADER_ENCODING] = self._client_session.encoding + else: + raise ValueError("Invalid type for encoding: expected str or list") + headers[constants.HEADER_CLIENT_CAPABILITIES] = constants.CLIENT_CAPABILITIES + + headers["user-agent"] = f"{constants.CLIENT_NAME}/{__version__}" + if len(self._client_session.roles.values()): + headers[constants.HEADER_ROLE] = ",".join( + # ``name`` must not contain ``=`` + "{}={}".format(catalog, urllib.parse.quote(str(role))) + for catalog, role in self._client_session.roles.items() + ) + if self._client_session.client_tags is not None and len(self._client_session.client_tags) > 0: + headers[constants.HEADER_CLIENT_TAGS] = ",".join(self._client_session.client_tags) + + headers[constants.HEADER_SESSION] = ",".join( + # ``name`` must not contain ``=`` + "{}={}".format(name, urllib.parse.quote(str(value))) + for name, value in self._client_session.properties.items() + ) + + if len(self._client_session.prepared_statements) != 0: + # ``name`` must not contain ``=`` + headers[constants.HEADER_PREPARED_STATEMENT] = ",".join( + "{}={}".format(name, urllib.parse.quote_plus(statement)) + for name, statement in self._client_session.prepared_statements.items() + ) + + # merge custom http headers + for key in self._client_session.headers: + if key in headers.keys(): + raise ValueError("cannot override reserved HTTP header {}".format(key)) + headers.update(self._client_session.headers) + + transaction_id = self._client_session.transaction_id + headers[constants.HEADER_TRANSACTION] = transaction_id + + if self._client_session.extra_credential is not None and \ + len(self._client_session.extra_credential) > 0: + + for tup in self._client_session.extra_credential: + self._verify_extra_credential(tup) + + # HTTP 1.1 section 4.2 combine multiple extra credentials into a + # comma-separated value + # extra credential value is encoded per spec (application/x-www-form-urlencoded MIME format) + headers[constants.HEADER_EXTRA_CREDENTIAL] = \ + ", ".join( + [f"{tup[0]}={urllib.parse.quote_plus(str(tup[1]))}" + for tup in self._client_session.extra_credential]) + + return headers + + 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 + ) + + @property + def statement_url(self) -> str: + return self.get_url(constants.URL_STATEMENT_PATH) + + @property + def next_uri(self) -> Optional[str]: + return self._next_uri + + @staticmethod + def _process_error( + error: Dict[str, Any], query_id: Optional[str] + ) -> Union[TrinoExternalError, TrinoQueryError, TrinoUserError]: + error_type = error["errorType"] + if error_type == "EXTERNAL": + raise exceptions.TrinoExternalError(error, query_id) + elif error_type == "USER_ERROR": + return exceptions.TrinoUserError(error, query_id) + + return exceptions.TrinoQueryError(error, query_id) + + @staticmethod + def raise_response_error(http_response: Any) -> None: + if http_response.status_code == 502: + raise exceptions.Http502Error("error 502: bad gateway") + + if http_response.status_code == 503: + raise exceptions.Http503Error("error 503: service unavailable") + + if http_response.status_code == 504: + raise exceptions.Http504Error("error 504: gateway timeout") + + raise exceptions.HttpError( + "error {}{}".format( + http_response.status_code, + ": {}".format(http_response.content) if http_response.content else "", + ) + ) + + def process(self, http_response: Any) -> TrinoStatus: + if http_response.is_error: + self.raise_response_error(http_response) + + try: + http_response.encoding = "utf-8" + except ValueError: + # .text was already accessed (e.g. by the retry logic), so the body + # has been decoded with the detected encoding already. + pass + if not http_response.text.strip(): + raise exceptions.TrinoConnectionError( + "received empty response from server (status 200)" + ) + response = json.loads(http_response.text) + if "error" in response and response["error"]: + raise self._process_error(response["error"], response.get("id")) + + if constants.HEADER_CLEAR_SESSION in http_response.headers: + for prop in get_header_values( + http_response.headers, constants.HEADER_CLEAR_SESSION + ): + self._client_session.properties.pop(prop, None) + + if constants.HEADER_SET_SESSION in http_response.headers: + for key, value in get_session_property_values( + http_response.headers, constants.HEADER_SET_SESSION + ): + self._client_session.properties[key] = value + + if constants.HEADER_SET_CATALOG in http_response.headers: + self._client_session.catalog = http_response.headers[constants.HEADER_SET_CATALOG] + + if constants.HEADER_SET_SCHEMA in http_response.headers: + self._client_session.schema = http_response.headers[constants.HEADER_SET_SCHEMA] + + if constants.HEADER_SET_ROLE in http_response.headers: + for key, value in get_roles_values( + http_response.headers, constants.HEADER_SET_ROLE + ): + self._client_session.roles[key] = value + + if constants.HEADER_ADDED_PREPARE in http_response.headers: + for name, statement in get_prepared_statement_values( + http_response.headers, constants.HEADER_ADDED_PREPARE + ): + self._client_session.prepared_statements[name] = statement + + if constants.HEADER_DEALLOCATED_PREPARE in http_response.headers: + for name in get_header_values( + http_response.headers, constants.HEADER_DEALLOCATED_PREPARE + ): + self._client_session.prepared_statements.pop(name, None) + + if constants.HEADER_SET_AUTHORIZATION_USER in http_response.headers: + self._client_session.authorization_user = http_response.headers[constants.HEADER_SET_AUTHORIZATION_USER] + + if constants.HEADER_RESET_AUTHORIZATION_USER in http_response.headers: + self._client_session.authorization_user = None + + self._next_uri = response.get("nextUri") + + data = response.get("data") if response.get("data") else [] + + return TrinoStatus( + id=response["id"], + stats=response["stats"], + warnings=response.get("warnings", []), + info_uri=response["infoUri"], + next_uri=self._next_uri, + update_type=response.get("updateType"), + update_count=response.get("updateCount"), + rows=data, + columns=response.get("columns"), + ) + + @staticmethod + def _verify_extra_credential(header: Tuple[str, str]) -> None: + """ + Verifies that key has ASCII only and non-whitespace characters. + """ + key = header[0] + + if not _HEADER_EXTRA_CREDENTIAL_KEY_REGEX.match(key): + raise ValueError(f"whitespace or '=' are disallowed in extra credential '{key}'") + + try: + key.encode().decode('ascii') + except UnicodeDecodeError: + raise ValueError(f"only ASCII characters are allowed in extra credential '{key}'") + + +class _TrinoQueryBase(abc.ABC): + """Transport-agnostic state of the execution of a SQL statement by Trino.""" + + def __init__( + self, + query: str, + legacy_primitive_types: bool = False, + fetch_mode: Literal["mapped", "segments"] = "mapped", + stats_callback: Optional[Callable[[Dict[str, Any]], None]] = None + ) -> None: + self._query_id: Optional[str] = None + self._stats: Dict[Any, Any] = {} + self._info_uri: Optional[str] = None + self._warnings: List[Dict[Any, Any]] = [] + self._columns: Optional[List[str]] = None + self._finished = False + self._cancelled = False + self._update_type: Optional[str] = None + self._update_count: Optional[int] = None + self._next_uri: Optional[str] = None + self._query = query + self._legacy_primitive_types = legacy_primitive_types + self._row_mapper: Optional[Any] = None + self._fetch_mode = fetch_mode + self._stats_callback = stats_callback + + @property + def query_id(self) -> Optional[str]: + return self._query_id + + @property + def query(self) -> Optional[str]: + return self._query + + @property + def stats(self) -> Dict[Any, Any]: + return self._stats + + @property + def update_type(self) -> Optional[str]: + return self._update_type + + @property + def update_count(self) -> Optional[int]: + return self._update_count + + @property + def warnings(self) -> List[Dict[Any, Any]]: + return self._warnings + + @property + def info_uri(self) -> Optional[str]: + return self._info_uri + + @property + def finished(self) -> bool: + return self._finished + + def is_finished(self) -> bool: + warnings.warn("is_finished is deprecated, use finished instead", DeprecationWarning) + return self.finished + + @property + def cancelled(self) -> bool: + return self._cancelled + + def _update_state(self, status: TrinoStatus) -> None: + self._stats.update(status.stats) + self._update_type = status.update_type + self._update_count = status.update_count + self._next_uri = status.next_uri + if not self._row_mapper and status.columns: + self._row_mapper = RowMapperFactory().create(columns=status.columns, + legacy_primitive_types=self._legacy_primitive_types) + if status.columns: + self._columns = status.columns + self._report_stats() + + def _report_stats(self) -> None: + if self._stats_callback is not None: + # Pass a deep copy so the callback cannot mutate internal query state. + self._stats_callback(copy.deepcopy(self._stats)) + + def _to_segments(self, rows: _SpooledProtocolResponseTO) -> List[DecodableSegment]: + encoding = rows["encoding"] + metadata = rows["metadata"] if "metadata" in rows else None + segments: List[Segment] = [] + for segment in rows["segments"]: + segment_type = segment["type"] + if segment_type == SegmentType.INLINE: + inline_segment = cast("_InlineSegmentTO", segment) + segments.append(InlineSegment(inline_segment)) + elif segment_type == SegmentType.SPOOLED: + segments.append(self._create_spooled_segment(cast("_SpooledSegmentTO", segment))) + else: + raise ValueError(f"Unsupported segment type: {segment_type}") + + return list(map(lambda segment: DecodableSegment(encoding, metadata, segment), segments)) + + @abstractmethod + def _create_spooled_segment(self, segment: _SpooledSegmentTO) -> Segment: + """Build the transport-specific spooled segment (sync or async download).""" + + +# Trino Spooled protocol transfer objects +class _SpooledProtocolResponseTO(TypedDict): + encoding: Literal["json", "json+std", "json+lz4"] + metadata: _SegmentMetadataTO + segments: List[_SegmentTO] + + +class _SegmentMetadataTO(TypedDict): + uncompressedSize: str + segmentSize: str + + +class _SegmentTO(_SegmentMetadataTO): + type: Literal["spooled", "inline"] + metadata: _SegmentMetadataTO + + +class _SpooledSegmentTO(_SegmentTO): + uri: str + ackUri: str + headers: Dict[str, List[str]] + + +class _InlineSegmentTO(_SegmentTO): + data: str + + +class SegmentType(str, Enum): + """Enum with string values that can be compared to strings.""" + INLINE = "inline" + SPOOLED = "spooled" + + +class Segment(abc.ABC): + """ + Abstract base class representing a segment of data produced by the spooling protocol. + + Attributes: + metadata (property): Metadata associated with the segment. + rows (property): Returns the decoded and mapped data. + """ + def __init__(self, segment: _SegmentTO) -> None: + self._segment = segment + + @property + def metadata(self) -> _SegmentMetadataTO: + return self._segment["metadata"] + + +class InlineSegment(Segment): + """ + A subclass of Segment that handles inline data segments. The data is base64 encoded and + requires mapping to rows using the provided row_mapper. + + Attributes: + rows (property): The data in the segment, decoded and mapped from the base64 encoded data. + """ + def __init__(self, segment: _InlineSegmentTO) -> None: + super().__init__(segment) + self._segment = cast(_InlineSegmentTO, segment) + + @property + def data(self) -> bytes: + return base64.b64decode(cast(_InlineSegmentTO, self._segment)["data"]) + + def __repr__(self) -> str: + return f"InlineSegment(metadata={self.metadata})" + + +class DecodableSegment: + """ + Represents a collection of spooled segments of data, with an encoding format. + + Attributes: + encoding (str): The encoding format of the spooled data. + metadata (_SegmentMetadataTO): Metadata for all segments in the query + segment (Segment): The spooled segment data + """ + def __init__(self, encoding: str, metadata: _SegmentMetadataTO, segment: Segment) -> None: + self._encoding = encoding + self._metadata = metadata + self._segment = segment + + @property + def encoding(self) -> str: + return self._encoding + + @property + def segment(self) -> Segment: + return self._segment + + @property + def metadata(self) -> _SegmentMetadataTO: + return self._metadata + + def __repr__(self) -> str: + return (f"DecodableSegment(encoding={self._encoding}, metadata={self._metadata}, segment={self._segment})") + + +def spooling_request_headers( + uri: str, + segment_headers: Dict[str, List[str]], + custom_headers: Dict[str, str], + coordinator_host: Optional[str], +) -> Dict[str, str]: + """ + Headers for a spooled segment data/ack request. + + Forward user-supplied custom headers (e.g. auth gateway headers) only when the + request targets the Trino coordinator, never to external storage (e.g. S3 presigned + URLs) where such headers can break the request. The per-segment protocol headers + returned by the coordinator always take precedence. + """ + headers: Dict[str, str] = {} + if coordinator_host is not None and urllib.parse.urlsplit(uri).hostname == coordinator_host: + headers.update(custom_headers) + for key, values in segment_headers.items(): + if len(values) > 1: + raise ValueError(f"Header '{key}' contains multiple values: {values}") + headers[key] = values[0] + return headers + + +class QueryDataDecoder(abc.ABC): + @abstractmethod + def decode(self, data: bytes, metadata: _SegmentMetadataTO) -> List[List[Any]]: + pass + + +class JsonQueryDataDecoder(QueryDataDecoder): + def __init__(self, mapper: RowMapper) -> None: + self._mapper = mapper + + def decode(self, data: bytes, metadata: _SegmentMetadataTO) -> List[List[Any]]: + return self._mapper.map(json.loads(data.decode("utf8"))) + + +class CompressedQueryDataDecoder(QueryDataDecoder): + def __init__(self, delegate: QueryDataDecoder) -> None: + self._delegate = delegate + + @abstractmethod + def decompress(self, data: bytes, metadata: _SegmentMetadataTO) -> bytes: + pass + + def decode(self, data: bytes, metadata: _SegmentMetadataTO) -> List[List[Any]]: + if "uncompressedSize" not in metadata: + # Data not compressed - below threshold + return self._delegate.decode(data, metadata) + + # Data is compressed + expected_compressed_size = metadata["segmentSize"] + if not len(data) == expected_compressed_size: + raise RuntimeError(f"Expected to read {expected_compressed_size} bytes but got {len(data)}") + decompressed_data = self.decompress(data, metadata) + expected_uncompressed_size = metadata["uncompressedSize"] + if not len(decompressed_data) == expected_uncompressed_size: + raise RuntimeError( + "Decompressed size does not match expected segment size, " + f"expected {expected_uncompressed_size}, got {len(decompressed_data)}" + ) + return self._delegate.decode(decompressed_data, metadata) + + +class ZStdQueryDataDecoder(CompressedQueryDataDecoder): + def __init__(self, delegate: QueryDataDecoder) -> None: + super().__init__(delegate) + self._decompressor: Optional[Any] = None + + def decompress(self, data: bytes, metadata: _SegmentMetadataTO) -> bytes: + if self._decompressor is None: + self._decompressor = zstandard.ZstdDecompressor() + return self._decompressor.decompress(data) + + +class Lz4QueryDataDecoder(CompressedQueryDataDecoder): + def decompress(self, data: bytes, metadata: _SegmentMetadataTO) -> bytes: + expected_uncompressed_size = metadata["uncompressedSize"] + decoded_bytes = lz4.block.decompress(data, uncompressed_size=int(expected_uncompressed_size)) + return decoded_bytes + + +class CompressedQueryDataDecoderFactory(): + def __init__(self, mapper: RowMapper) -> None: + self._mapper = mapper + + def create(self, encoding: str) -> QueryDataDecoder: + if encoding == "json+zstd": + if "zstd" in CODECS_UNAVAILABLE: + raise ValueError( + f"zstd is not installed so json+zstd encoding is not supported: {CODECS_UNAVAILABLE['zstd']}" + ) + return ZStdQueryDataDecoder(JsonQueryDataDecoder(self._mapper)) + elif encoding == "json+lz4": + if "lz4" in CODECS_UNAVAILABLE: + raise ValueError( + f"lz4 is not installed so json+lz4 encoding is not supported: {CODECS_UNAVAILABLE['lz4']}" + ) + return Lz4QueryDataDecoder(JsonQueryDataDecoder(self._mapper)) + elif encoding == "json": + return JsonQueryDataDecoder(self._mapper) + else: + raise ValueError(f"Unsupported encoding: {encoding}") diff --git a/trino/_spnego.py b/trino/_spnego.py new file mode 100644 index 00000000..e1429c4e --- /dev/null +++ b/trino/_spnego.py @@ -0,0 +1,199 @@ +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +SPNEGO (Kerberos/GSSAPI "Negotiate") authentication flow for httpx2. + +This is an in-repo port of the flow implemented by requests_kerberos / +requests_gssapi / httpx-gssapi, built directly on python-gssapi and the +``httpx2.Auth`` generator protocol so it works with both the synchronous and +asynchronous clients. It backs both ``trino.auth.KerberosAuthentication`` and +``trino.auth.GSSAPIAuthentication``. + +Requires the ``gssapi`` package (the ``trino[gssapi]``/``trino[kerberos]`` +extras) and a GSSAPI implementation such as MIT Kerberos. Unlike +requests_kerberos, Windows SSPI is not supported. +""" +from __future__ import annotations + +import base64 +import re +import typing + +import httpx2 + +from trino.exceptions import SPNEGOExchangeError + +# Mutual-authentication modes; the values match requests_kerberos/requests_gssapi. +MUTUAL_REQUIRED = 1 +MUTUAL_OPTIONAL = 2 +MUTUAL_DISABLED = 3 + +_NEGOTIATE_TOKEN_PATTERN = re.compile(r"Negotiate\s*([^,\s]*)", flags=re.IGNORECASE) + +# The maximum number of 401 challenges answered before giving up, matching +# the requests_gssapi/httpx-gssapi behavior. +_MAX_CHALLENGES = 2 + + +def _import_gssapi() -> typing.Any: + try: + import gssapi + except ImportError: + raise RuntimeError("unable to import gssapi") + return gssapi + + +class SPNEGOAuth(httpx2.Auth): + """ + httpx2 authentication flow answering HTTP Negotiate (SPNEGO) challenges. + + :param target_name: the GSSAPI target. Either a ``gssapi.Name``, a host + name string (authenticates the default ``HTTP`` service at that host) + or ``None`` (authenticates ``service_name`` at the request's host). + :param service_name: Kerberos service, defaults to ``HTTP``. + :param hostname_override: authenticate against this host instead of the + host of the request URL. + :param mutual_authentication: MUTUAL_REQUIRED, MUTUAL_OPTIONAL or + MUTUAL_DISABLED. When enabled, the server's final Negotiate token is + verified; on MUTUAL_REQUIRED a missing or invalid token raises + :class:`~trino.exceptions.SPNEGOExchangeError`. + :param opportunistic_auth: send an initial token preemptively instead of + waiting for a 401 challenge. + :param delegate: request credential delegation. + :param creds: explicit ``gssapi.Credentials`` to use. + :param sanitize_mutual_error_response: accepted for signature compatibility + with requests_kerberos/requests_gssapi; response bodies of failed + exchanges are never exposed as authenticated content by this client. + """ + + def __init__( + self, + *, + target_name: typing.Any = None, + service_name: typing.Optional[str] = None, + hostname_override: typing.Optional[str] = None, + mutual_authentication: int = MUTUAL_DISABLED, + opportunistic_auth: bool = False, + delegate: bool = False, + creds: typing.Any = None, + sanitize_mutual_error_response: bool = True, + ) -> None: + # Fail early when gssapi is missing, like the requests-era classes did. + _import_gssapi() + self.target_name = target_name + self.service_name = service_name + self.hostname_override = hostname_override + self.mutual_authentication = mutual_authentication + self.opportunistic_auth = opportunistic_auth + self.delegate = delegate + self.creds = creds + self.sanitize_mutual_error_response = sanitize_mutual_error_response + + def _resolve_target_name(self, request_host: str) -> typing.Any: + gssapi = _import_gssapi() + if isinstance(self.target_name, gssapi.Name): + return self.target_name + # A plain-string target is a host name carrying the default service. + host = self.target_name or self.hostname_override or request_host + service = self.service_name or "HTTP" + return gssapi.Name(f"{service}@{host}", gssapi.NameType.hostbased_service) + + def _make_context(self, request_host: str) -> typing.Any: + gssapi = _import_gssapi() + flags = [gssapi.RequirementFlag.out_of_sequence_detection] + if self.delegate: + flags.append(gssapi.RequirementFlag.delegate_to_peer) + if self.mutual_authentication != MUTUAL_DISABLED: + flags.append(gssapi.RequirementFlag.mutual_authentication) + try: + return gssapi.SecurityContext( + name=self._resolve_target_name(request_host), + creds=self.creds, + mech=gssapi.OID.from_int_seq("1.3.6.1.5.5.2"), # SPNEGO + flags=flags, + usage="initiate", + ) + except gssapi.exceptions.GSSError as exc: + raise SPNEGOExchangeError(f"failed to initialize GSSAPI context: {exc}") from exc + + @staticmethod + def _step(context: typing.Any, in_token: typing.Optional[bytes]) -> typing.Optional[bytes]: + gssapi = _import_gssapi() + try: + return context.step(in_token) + except gssapi.exceptions.GSSError as exc: + raise SPNEGOExchangeError(f"GSSAPI token exchange failed: {exc}") from exc + + @staticmethod + def _extract_negotiate_token(response: httpx2.Response) -> typing.Optional[bytes]: + auth_header = response.headers.get("WWW-Authenticate", "") + match = _NEGOTIATE_TOKEN_PATTERN.search(auth_header) + if match is None: + return None + token = match.group(1) + return base64.b64decode(token) if token else b"" + + @staticmethod + def _has_negotiate_challenge(response: httpx2.Response) -> bool: + return "negotiate" in response.headers.get("WWW-Authenticate", "").lower() + + def _set_header(self, request: httpx2.Request, out_token: bytes) -> None: + request.headers["Authorization"] = "Negotiate " + base64.b64encode(out_token).decode() + + def auth_flow( + self, request: httpx2.Request + ) -> typing.Generator[httpx2.Request, httpx2.Response, None]: + host = request.url.host + context = None + + if self.opportunistic_auth: + context = self._make_context(host) + out_token = self._step(context, None) + if out_token: + self._set_header(request, out_token) + + response = yield request + + challenges_answered = 0 + while ( + response.status_code == 401 + and self._has_negotiate_challenge(response) + and challenges_answered < _MAX_CHALLENGES + ): + in_token = self._extract_negotiate_token(response) + if context is None: + context = self._make_context(host) + out_token = self._step(context, in_token or None) + if out_token is None: + break + self._set_header(request, out_token) + response = yield request + challenges_answered += 1 + + if self.mutual_authentication != MUTUAL_DISABLED and not response.is_error: + self._verify_mutual_auth(context, response) + + def _verify_mutual_auth(self, context: typing.Any, response: httpx2.Response) -> None: + in_token = self._extract_negotiate_token(response) + if not in_token: + if self.mutual_authentication == MUTUAL_REQUIRED: + raise SPNEGOExchangeError( + "mutual authentication required but the server sent no Negotiate token" + ) + return + if context is None: + raise SPNEGOExchangeError( + "mutual authentication required but no GSSAPI exchange took place" + ) + self._step(context, in_token) + if not context.complete: + raise SPNEGOExchangeError("mutual authentication failed: incomplete GSSAPI context") diff --git a/trino/aio/__init__.py b/trino/aio/__init__.py new file mode 100644 index 00000000..f26b312e --- /dev/null +++ b/trino/aio/__init__.py @@ -0,0 +1,33 @@ +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Asynchronous Trino client; see :mod:`trino.aio.dbapi` for usage.""" +from trino.aio.client import AsyncSegmentIterator +from trino.aio.client import AsyncSpooledSegment +from trino.aio.client import AsyncTrinoQuery +from trino.aio.client import AsyncTrinoRequest +from trino.aio.client import AsyncTrinoResult +from trino.aio.dbapi import connect +from trino.aio.dbapi import Connection +from trino.aio.dbapi import Cursor +from trino.aio.dbapi import SegmentCursor + +__all__ = [ + "connect", + "Connection", + "Cursor", + "SegmentCursor", + "AsyncTrinoQuery", + "AsyncTrinoRequest", + "AsyncTrinoResult", + "AsyncSegmentIterator", + "AsyncSpooledSegment", +] diff --git a/trino/aio/client.py b/trino/aio/client.py new file mode 100644 index 00000000..943ad6a9 --- /dev/null +++ b/trino/aio/client.py @@ -0,0 +1,709 @@ +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Asynchronous variant of :mod:`trino.client` built on ``httpx2.AsyncClient``. + +The protocol logic (headers, response processing, retry decisions, segment +decoding) is shared with the synchronous client through +:mod:`trino._protocol`; this module only contains the ``await``-shaped +control flow. +""" +from __future__ import annotations + +import asyncio +import copy +import functools +from collections.abc import AsyncIterator +from typing import Any +from typing import Callable +from typing import cast +from typing import Dict +from typing import List +from typing import Literal +from typing import Optional +from typing import Set +from typing import Tuple +from typing import Union + +import httpx2 +from httpx2 import AsyncClient +from httpx2 import Response + +import trino.logging +from trino import constants +from trino import exceptions +from trino._protocol import _DelayExponential +from trino._protocol import _SpooledProtocolResponseTO +from trino._protocol import _SpooledSegmentTO +from trino._protocol import _TrinoQueryBase +from trino._protocol import _TrinoRequestBase +from trino._protocol import ClientSession +from trino._protocol import CompressedQueryDataDecoderFactory +from trino._protocol import DecodableSegment +from trino._protocol import InlineSegment +from trino._protocol import needs_retry +from trino._protocol import retry_after_seconds +from trino._protocol import Segment +from trino._protocol import spooling_request_headers +from trino._protocol import wire_headers +from trino.auth import Authentication +from trino.client import MAX_ATTEMPTS +from trino.client import PROXIES +from trino.mapper import RowMapper + +__all__ = [ + "AsyncTrinoQuery", + "AsyncTrinoRequest", + "AsyncTrinoResult", + "AsyncSegmentIterator", + "AsyncSpooledSegment", +] + +logger = trino.logging.get_logger(__name__) + + +class _AsyncRetryWithExponentialBackoff: + def __init__( + self, base=0.1, exponent=2, jitter=True, max_delay=1800 # 100ms # 30 min + ): + self._get_delay = _DelayExponential(base, exponent, jitter, max_delay) + + async def retry(self, func, args, kwargs, err, attempt): + await asyncio.sleep(self._get_delay(attempt)) + + +def _retry_with_async(handle_retry, handled_exceptions, conditions, max_attempts): + def wrapper(func): + @functools.wraps(func) + async def decorated(*args, **kwargs): + error = None + result = None + for attempt in range(1, max_attempts + 1): + try: + result = await func(*args, **kwargs) + if any(guard(result) for guard in conditions): + retry_after = retry_after_seconds(result) + if retry_after is not None: + await asyncio.sleep(retry_after) + else: + await handle_retry.retry(func, args, kwargs, None, attempt) + continue + return result + except Exception as err: + error = err + if any(isinstance(err, exc) for exc in handled_exceptions): + await handle_retry.retry(func, args, kwargs, err, attempt) + continue + break + logger.info("failed after %s attempts", attempt) + if error is not None: + raise error + return result + + return decorated + + return wrapper + + +class AsyncTrinoRequest(_TrinoRequestBase): + """ + Asynchronous twin of :class:`trino.client.TrinoRequest` on an + ``httpx2.AsyncClient``. The constructor parameters are identical; the HTTP + verbs are coroutines. + """ + + http = httpx2 + + HTTP_EXCEPTIONS: Tuple[Any, ...] = ( + http.TransportError, + ) + + def __init__( + self, + host: str, + port: int, + client_session: ClientSession, + http_session: Optional[AsyncClient] = None, + http_scheme: Optional[str] = None, + auth: Optional[Authentication] = constants.DEFAULT_AUTH, + max_attempts: int = MAX_ATTEMPTS, + request_timeout: Union[float, Tuple[float, float]] = constants.DEFAULT_REQUEST_TIMEOUT, + handle_retry=_AsyncRetryWithExponentialBackoff(), + verify: Union[bool, str] = True, + task_registry: Optional[Set[asyncio.Task]] = None, + ) -> None: + super().__init__(host, port, client_session, http_scheme) + + self._owns_http_session = http_session is None + if http_session is not None: + self._http_session = http_session + if auth is not None: + self._apply_auth_to_existing_client(http_session, auth) + else: + self._http_session = self.create_http_client( + verify=verify, timeout=request_timeout, auth=auth + ) + self._verify = getattr(self._http_session, "_trino_verify", verify) + self._http_session.headers.update(wire_headers(self.http_headers)) + self._exceptions = self.HTTP_EXCEPTIONS + self._auth = auth + if self._auth: + self._exceptions += self._auth.get_exceptions() + + self._request_timeout = self.http.Timeout(request_timeout) + self._handle_retry = handle_retry + # Fire-and-forget acknowledgment tasks; kept strongly referenced here + # and awaited by drain()/aclose() so none is garbage collected mid-run. + self._task_registry: Set[asyncio.Task] = task_registry if task_registry is not None else set() + # Requests spawned by unauthenticated(); their clients are closed with this one. + self._children: List[AsyncTrinoRequest] = [] + self.max_attempts = max_attempts + + @classmethod + def create_http_client( + cls, + verify: Union[bool, str] = True, + timeout: Union[float, Tuple[float, float], None] = constants.DEFAULT_REQUEST_TIMEOUT, + auth: Optional[Authentication] = None, + **kwargs: Any, + ) -> AsyncClient: + """Async twin of :meth:`trino.client.TrinoRequest.create_http_client`.""" + client_kwargs: Dict[str, Any] = { + "verify": verify, + "http2": True, + "follow_redirects": True, + "timeout": cls.http.Timeout(timeout), + } + if PROXIES: + client_kwargs["proxy"] = PROXIES.get("all://") + auth_arguments: Dict[str, Any] = {} + http_auth = None + if auth is not None: + auth_arguments = auth.get_client_arguments() + client_kwargs.update(auth_arguments) + http_auth = auth.get_http_auth() + client_kwargs.update(kwargs) + client = cls.http.AsyncClient(auth=http_auth, **client_kwargs) + client._trino_verify = client_kwargs["verify"] + client._trino_client_arguments = frozenset(auth_arguments) + return client + + @staticmethod + def _apply_auth_to_existing_client(http_session: AsyncClient, auth: Authentication) -> None: + required = auth.get_client_arguments() + provided = getattr(http_session, "_trino_client_arguments", frozenset()) + missing = set(required) - set(provided) + if missing: + raise exceptions.TrinoConnectionError( + f"{type(auth).__name__} requires HTTP client construction options {sorted(missing)}; " + "configure them on your own httpx2 client, or omit http_session to let the " + "client be created for you." + ) + http_session.auth = auth.get_http_auth() + + def unauthenticated(self) -> AsyncTrinoRequest: + request = AsyncTrinoRequest( + host=self._host, + port=self._port, + max_attempts=self.max_attempts, + request_timeout=self._request_timeout, + handle_retry=self._handle_retry, + client_session=ClientSession(user=self._client_session.user), + verify=self._verify, + task_registry=self._task_registry) + self._children.append(request) + return request + + def register_task(self, coro) -> asyncio.Task: + """Run a fire-and-forget coroutine, keeping a strong reference to the + task until it completes so it survives garbage collection.""" + task = asyncio.ensure_future(coro) + self._task_registry.add(task) + task.add_done_callback(self._task_registry.discard) + return task + + async def drain(self) -> None: + """Wait for pending background tasks and close spawned child clients.""" + pending = list(self._task_registry) + if pending: + await asyncio.gather(*pending, return_exceptions=True) + children, self._children = self._children, [] + for child in children: + await child.aclose() + + async def aclose(self) -> None: + await self.drain() + if self._owns_http_session: + await self._http_session.aclose() + + @property + def max_attempts(self) -> int: + return self._max_attempts + + @max_attempts.setter + def max_attempts(self, value: int) -> None: + self._max_attempts = value + if value == 1: # No retry + self._get = self._http_session.get + self._post = self._http_session.post + self._delete = self._http_session.delete + self._head = self._http_session.head + return + + with_retry = _retry_with_async( + self._handle_retry, + handled_exceptions=self._exceptions, + conditions=(needs_retry,), + max_attempts=self._max_attempts, + ) + self._get = with_retry(self._http_session.get) + self._post = with_retry(self._http_session.post) + self._delete = with_retry(self._http_session.delete) + self._head = with_retry(self._http_session.head) + + async def post(self, sql: str, additional_http_headers: Optional[Dict[str, Any]] = None) -> Response: + data = sql.encode("utf-8") + http_headers = copy.deepcopy(self.http_headers) + http_headers.update(additional_http_headers or {}) + http_headers.setdefault(constants.HEADER_CONTENT_TYPE, constants.CONTENT_TYPE_TEXT_UTF8) + + return await self._post( + self.statement_url, + content=data, + headers=wire_headers(http_headers), + timeout=self._request_timeout, + ) + + async def get(self, url: str) -> Response: + return await self._get( + url, + headers=wire_headers(self.http_headers), + timeout=self._request_timeout, + ) + + async def delete(self, url: str) -> Response: + return await self._delete(url, timeout=self._request_timeout) + + async def head(self, url: str) -> Response: + return await self._head( + url, + headers=wire_headers(self.http_headers), + timeout=self._request_timeout, + follow_redirects=False, + ) + + +async def _aiterate(rows) -> AsyncIterator[List[Any]]: + """Iterate rows that may be a list or an asynchronous iterable.""" + if hasattr(rows, "__aiter__"): + async for row in rows: + yield row + else: + for row in rows: + yield row + + +async def _achain(*iterables) -> AsyncIterator[List[Any]]: + """Chain multiple (possibly asynchronous) iterables of rows.""" + for rows in iterables: + async for row in _aiterate(rows): + yield row + + +async def _aprepend(first_row, rest) -> AsyncIterator[List[Any]]: + yield first_row + async for row in _aiterate(rest): + yield row + + +class AsyncTrinoResult: + """ + Asynchronous twin of :class:`trino.client.TrinoResult`: an async iterator + over the rows of a query. Iteration state lives on the instance so a + transient error (e.g. a failed spooled segment download) propagates to the + caller while the iterator stays usable. + """ + + def __init__(self, query: AsyncTrinoQuery, rows) -> None: + self._query = query + self._rows = rows + self._rownumber = 0 + self._current_batch: Optional[AsyncIterator[Any]] = None + self._next_rows = None + + @property + def rows(self): + return self._rows + + @rows.setter + def rows(self, rows): + self._rows = rows + + @property + def rownumber(self) -> int: + return self._rownumber + + def __aiter__(self) -> AsyncTrinoResult: + return self + + async def __anext__(self): + while True: + if self._current_batch is None: + if self._query.finished and self._rows is None: + raise StopAsyncIteration + self._next_rows = await self._query.fetch() if not self._query.finished else None + self._current_batch = _aiterate(self._rows) + + try: + row = await self._current_batch.__anext__() + except StopAsyncIteration: + self._rows = self._next_rows + self._next_rows = None + self._current_batch = None + continue + self._rownumber += 1 + return row + + +class AsyncTrinoQuery(_TrinoQueryBase): + """Represent the asynchronous execution of a SQL statement by Trino.""" + + def __init__( + self, + request: AsyncTrinoRequest, + query: str, + legacy_primitive_types: bool = False, + fetch_mode: Literal["mapped", "segments"] = "mapped", + stats_callback: Optional[Callable[[Dict[str, Any]], None]] = None + ) -> None: + super().__init__(query, legacy_primitive_types, fetch_mode, stats_callback) + self._request = request + self._result: Optional[AsyncTrinoResult] = None + + @property + def columns(self): + """The columns of the result set, as known after execute() finished. + + Unlike the synchronous client this does not lazily fetch further pages; + execute() already polls until rows or a terminal state arrived. + """ + return self._columns + + @property + def result(self): + return self._result + + async def execute(self, additional_http_headers=None) -> AsyncTrinoResult: + """Initiate a Trino query by sending the SQL statement + + This is the first HTTP request sent to the coordinator. + It sets the query_id and returns a Result object used to + track the rows returned by the query. + """ + if self.cancelled: + raise exceptions.TrinoUserError("Query has been cancelled", self.query_id) + + try: + response = await self._request.post(self._query, additional_http_headers) + except httpx2.HTTPError as e: + raise trino.exceptions.TrinoConnectionError("failed to execute: {}".format(e)) + status = self._request.process(response) + self._info_uri = status.info_uri + self._query_id = status.id + self._stats.update({"queryId": self.query_id}) + self._update_state(status) + self._warnings = getattr(status, "warnings", []) + if status.next_uri is None: + self._finished = True + + rows = self._row_mapper.map(status.rows) if self._row_mapper else status.rows + self._result = AsyncTrinoResult(self, rows) + + # Block until rows are available, the query finishes, or it is + # canceled; see TrinoQuery.execute for the direct-vs-spooling shape + # of these loops. + while not self.finished and not self.cancelled and self._result.rows == []: + new_rows = await self.fetch() + if isinstance(new_rows, list): + self._result.rows += new_rows + else: + try: + first_row = await new_rows.__anext__() + self._result.rows = _aprepend(first_row, new_rows) + break + except StopAsyncIteration: + self._result.rows = [] + + # Drain the trailing pages of update statements so the query reaches a + # terminal state; see TrinoQuery.execute for the full rationale. + while self._update_type is not None and not self.finished and not self.cancelled: + new_rows = await self.fetch() + if isinstance(self._result.rows, list) and isinstance(new_rows, list): + self._result.rows += new_rows + else: + self._result.rows = _achain(self._result.rows, new_rows) + + return self._result + + async def fetch(self) -> Union[List[Union[List[Any], Any]], AsyncIterator[List[Any]]]: + """Continue fetching data for the current query_id""" + try: + response = await self._request.get(self._request.next_uri) + except httpx2.HTTPError as e: + raise trino.exceptions.TrinoConnectionError("failed to fetch: {}".format(e)) + status = self._request.process(response) + self._update_state(status) + if status.next_uri is None: + self._finished = True + + if not self._row_mapper: + return [] + + rows = status.rows + if isinstance(status.rows, dict): + # spooling protocol + rows = cast(_SpooledProtocolResponseTO, rows) + spooled = self._to_segments(rows) + if self._fetch_mode == "segments": + return spooled + return AsyncSegmentIterator( + spooled, + self._row_mapper, + request=self._request, + heartbeat_interval=self._request._client_session.heartbeat_interval, + ) + elif isinstance(status.rows, list): + return self._row_mapper.map(rows) + else: + raise ValueError(f"Unexpected type: {type(status.rows)}") + + def _create_spooled_segment(self, segment: _SpooledSegmentTO) -> Segment: + return AsyncSpooledSegment( + segment, + self._request.unauthenticated(), + coordinator_host=self._request._host, + custom_headers=dict(self._request._client_session.headers), + ) + + async def cancel(self) -> None: + """Cancel the current query""" + if self._next_uri is None: + return + + logger.debug("cancelling query: %s", self.query_id) + try: + response = await self._request.delete(self._next_uri) + except httpx2.HTTPError as e: + raise trino.exceptions.TrinoConnectionError("failed to cancel query: {}".format(e)) + if response.status_code == httpx2.codes.NO_CONTENT: + self._cancelled = True + logger.debug("query cancelled: %s", self.query_id) + return + + self._request.raise_response_error(response) + + +class AsyncSpooledSegment(Segment): + """ + Asynchronous twin of :class:`trino.client.SpooledSegment`: the segment data + is downloaded with ``await segment.data()`` and acknowledgments run as + asyncio tasks registered on the owning request. + """ + + def __init__( + self, + segment: _SpooledSegmentTO, + request: AsyncTrinoRequest, + coordinator_host: Optional[str] = None, + custom_headers: Optional[Dict[str, str]] = None, + ) -> None: + super().__init__(segment) + self._segment = cast(_SpooledSegmentTO, segment) + self._request = request + self._coordinator_host = coordinator_host + self._custom_headers = custom_headers or {} + + async def data(self) -> bytes: + # No timeout on the data path: downloading a large segment from + # external storage may legitimately take a long time. + http_response = await self._send_spooling_request(self.uri, timeout=None) + if http_response.is_error: + self._request.raise_response_error(http_response) + return http_response.content + + @property + def uri(self) -> str: + return self._segment["uri"] + + @property + def ack_uri(self) -> str: + return self._segment["ackUri"] + + @property + def headers(self) -> Dict[str, List[str]]: + return self._segment.get("headers", {}) + + def acknowledge(self) -> None: + async def acknowledge_request(): + try: + http_response = await self._send_spooling_request(self.ack_uri, timeout=2) + if http_response.is_error: + self._request.raise_response_error(http_response) + except Exception as e: + logger.error(f"Failed to acknowledge spooling request for segment {self}: {e}") + + # Fire and forget; the task is awaited when the connection is closed. + self._request.register_task(acknowledge_request()) + + async def _send_spooling_request(self, uri: str, **kwargs) -> Response: + headers = spooling_request_headers(uri, self.headers, self._custom_headers, self._coordinator_host) + return await self._request._get(uri, headers=headers, **kwargs) + + def __repr__(self): + return ( + f"AsyncSpooledSegment(metadata={self.metadata})" + ) + + +class _AsyncRequestHeartbeat: + """ + Asynchronous twin of :class:`trino.client._RequestHeartbeat`: an asyncio + task periodically sending HEAD requests to the request's next URI while a + spooled segment downloads. + """ + MAX_FAILURES = 3 + + def __init__(self, request: AsyncTrinoRequest, interval: float) -> None: + self._request = request + self._interval = interval + self._stop_event = asyncio.Event() + self._task: Optional[asyncio.Task] = None + + async def __aenter__(self) -> _AsyncRequestHeartbeat: + self._task = asyncio.ensure_future(self._run()) + return self + + async def __aexit__(self, *_) -> None: + self._stop_event.set() + if self._task is not None: + await self._task + + async def _run(self) -> None: + """ + Run the heartbeat loop. + + Exit when the stop event is set, the query completed or if the error + count exceeds MAX_FAILURES. + """ + failures = 0 + + while True: + try: + await asyncio.wait_for(self._stop_event.wait(), timeout=self._interval) + return + except asyncio.TimeoutError: + pass + + uri = self._request.next_uri + if uri is None: + return + + try: + response = await self._request.head(uri) + if response.status_code in (404, 405): + logger.warning("The server does not support heartbeat calls") + return + if response.is_error: + failures += 1 + else: + failures = 0 + except Exception: + failures += 1 + + if failures >= self.MAX_FAILURES: + logger.warning(f"Stopping the heartbeat after {self.MAX_FAILURES} consecutive errors") + return + + +class AsyncSegmentIterator: + def __init__( + self, + segments: Union[DecodableSegment, List[DecodableSegment]], + mapper: RowMapper, + *, + request: Optional[AsyncTrinoRequest] = None, + heartbeat_interval: Optional[float] = None, + ) -> None: + self._segments = iter(segments if isinstance(segments, List) else [segments]) + self._mapper = mapper + self._decoder = None + self._rows = iter([]) + self._finished = False + self._current_segment: Optional[DecodableSegment] = None + # Segment whose decoding failed. Retried on the next call instead of being acknowledged and skipped. + self._pending_segment: Optional[DecodableSegment] = None + if (request is not None) != bool(heartbeat_interval): + raise ValueError("request and heartbeat_interval must be both provided or both omitted") + self._request = request + self._heartbeat_interval = heartbeat_interval + + def __aiter__(self) -> AsyncSegmentIterator: + return self + + async def __anext__(self) -> List[Any]: + # If rows are exhausted, fetch the next segment + while True: + try: + return next(self._rows) + except StopIteration: + if self._finished: + raise StopAsyncIteration + await self._load_next_segment() + + async def _load_next_segment(self): + # A segment is acknowledged only after its rows were decoded successfully. If the previous attempt failed + # mid-decode (e.g. the spooled segment download failed) the same segment is retried instead of being skipped. + if self._pending_segment is None: + if self._current_segment: + segment = self._current_segment.segment + if isinstance(segment, AsyncSpooledSegment): + segment.acknowledge() + self._current_segment = None + + try: + self._pending_segment = next(self._segments) + except StopIteration: + self._finished = True + return + + if self._decoder is None: + self._decoder = CompressedQueryDataDecoderFactory(self._mapper) \ + .create(self._pending_segment.encoding) + + segment = self._pending_segment.segment + if isinstance(segment, AsyncSpooledSegment): + if self._request and self._heartbeat_interval: + # Downloading a spooled segment may take some time. In the meantime, send heartbeat + # requests so the coordinator doesn't think we lost interest and close the query. + async with _AsyncRequestHeartbeat(self._request, self._heartbeat_interval): + data = await segment.data() + else: + data = await segment.data() + elif isinstance(segment, InlineSegment): + data = segment.data + else: + raise ValueError(f"Unsupported segment type: {type(segment)}") + + rows = self._decoder.decode(data, segment.metadata) + self._rows = iter(rows) + self._current_segment = self._pending_segment + self._pending_segment = None diff --git a/trino/aio/dbapi.py b/trino/aio/dbapi.py new file mode 100644 index 00000000..8b48f22a --- /dev/null +++ b/trino/aio/dbapi.py @@ -0,0 +1,463 @@ +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" + +Asynchronous, DBAPI-like interface to Trino built on ``httpx2.AsyncClient``. + +The API mirrors :mod:`trino.dbapi` with coroutine methods, but it is not a +PEP 249 implementation (PEP 249 is a synchronous specification): + +:: + + import trino.aio + + conn = trino.aio.connect(host="coordinator", port=8080, user="user") + cur = conn.cursor() + await cur.execute("SELECT * FROM system.runtime.nodes") + rows = await cur.fetchall() + await conn.close() + +Transactions are not supported by the asynchronous client yet. +""" +import asyncio +from typing import Any +from typing import Callable +from typing import Dict +from typing import List +from typing import Optional +from typing import Set +from typing import Union +from urllib.parse import urlparse + +import httpx2 + +import trino.client +import trino.exceptions +import trino.logging +from trino import constants +from trino.aio.client import AsyncTrinoQuery +from trino.aio.client import AsyncTrinoRequest +from trino.dbapi import _default_spooling_encoding +from trino.dbapi import _require_tls_for_auth +from trino.dbapi import _resolve_scheme_and_port +from trino.dbapi import _USE_DEFAULT_ENCODING +from trino.dbapi import Cursor as _SyncCursor +from trino.dbapi import DescribeOutput +from trino.dbapi import must_use_legacy_prepared_statements +from trino.exceptions import NotSupportedError +from trino.transaction import IsolationLevel +from trino.transaction import NO_TRANSACTION + +__all__ = ["connect", "Connection", "Cursor", "SegmentCursor"] + +logger = trino.logging.get_logger(__name__) + + +def connect(*args, **kwargs): + """Constructor for creating a connection to the database. + + See class :py:class:`Connection` for arguments. No I/O happens until the + first statement is executed. + + :returns: a :py:class:`Connection` object. + """ + return Connection(*args, **kwargs) + + +class Connection: + """Asynchronous Trino connection. Accepts the same arguments as + :class:`trino.dbapi.Connection` except that transactions (a non-default + ``isolation_level``) are not supported yet.""" + + def __init__( + self, + host: str, + port=None, + user=None, + source=constants.DEFAULT_SOURCE, + catalog=constants.DEFAULT_CATALOG, + schema=constants.DEFAULT_SCHEMA, + session_properties=None, + http_headers=None, + http_scheme=None, + auth=constants.DEFAULT_AUTH, + extra_credential=None, + max_attempts=constants.DEFAULT_MAX_ATTEMPTS, + request_timeout=constants.DEFAULT_REQUEST_TIMEOUT, + isolation_level=IsolationLevel.AUTOCOMMIT, + verify=True, + http_session: Optional[httpx2.AsyncClient] = None, + client_tags=None, + legacy_primitive_types=False, + legacy_prepared_statements=None, + roles=None, + timezone=None, + encoding: Union[str, List[str]] = _USE_DEFAULT_ENCODING, + heartbeat_interval: Optional[float] = constants.DEFAULT_HEARTBEAT_INTERVAL, + allow_insecure_auth: bool = False, + ): + if isolation_level != IsolationLevel.AUTOCOMMIT: + raise NotSupportedError("transactions are not supported by the asynchronous client yet") + + # Automatically assign http_schema, port based on hostname + parsed_host = urlparse(host, allow_fragments=False) + + if encoding is _USE_DEFAULT_ENCODING: + encoding = _default_spooling_encoding() + + 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 + self.schema = schema + self.session_properties = session_properties + self._client_session = trino.client.ClientSession( + user=user, + catalog=catalog, + schema=schema, + source=source, + properties=session_properties, + headers=http_headers, + transaction_id=NO_TRANSACTION, + extra_credential=extra_credential, + client_tags=client_tags, + roles=roles, + timezone=timezone, + encoding=encoding, + heartbeat_interval=heartbeat_interval, + ) + if http_session is None: + self._http_session = AsyncTrinoRequest.create_http_client( + verify=verify, timeout=request_timeout, auth=auth + ) + else: + self._http_session = http_session + self.http_headers = http_headers + + self.http_scheme, self.port = _resolve_scheme_and_port(parsed_host, port, http_scheme) + _require_tls_for_auth(auth, self.http_scheme, allow_insecure_auth) + + self.auth = auth + self.extra_credential = extra_credential + self.max_attempts = max_attempts + self.request_timeout = request_timeout + self.client_tags = client_tags + + self._isolation_level = isolation_level + self.legacy_primitive_types = legacy_primitive_types + self.legacy_prepared_statements = legacy_prepared_statements + + # Pending fire-and-forget tasks (spooled segment acknowledgments) and + # the requests spawned by this connection, both drained on close(). + self._task_registry: Set[asyncio.Task] = set() + self._requests: List[AsyncTrinoRequest] = [] + + @property + def isolation_level(self): + return self._isolation_level + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc_value, traceback): + await self.close() + + async def close(self): + for request in self._requests: + await request.drain() + await self._http_session.aclose() + + def start_transaction(self): + raise NotSupportedError("transactions are not supported by the asynchronous client yet") + + async def commit(self): + # Autocommit is the only supported mode; there is nothing to commit. + pass + + async def rollback(self): + raise trino.exceptions.NotSupportedError( + "transactions are not supported by the asynchronous client yet") + + def _create_request(self) -> AsyncTrinoRequest: + request = AsyncTrinoRequest( + self.host, + self.port, + self._client_session, + self._http_session, + self.http_scheme, + self.auth, + self.max_attempts, + self.request_timeout, + task_registry=self._task_registry, + ) + self._requests.append(request) + return request + + def cursor( + self, + cursor_style: str = "row", + legacy_primitive_types: bool = None, + stats_callback: Optional[Callable[[Dict[str, Any]], None]] = None): + """Return a new :py:class:`Cursor` object using the connection.""" + cursor_class = { + # Add any custom Cursor classes here + "segment": SegmentCursor, + "row": Cursor + }.get(cursor_style.lower(), Cursor) + + return cursor_class( + self, + self._create_request(), + legacy_primitive_types=( + legacy_primitive_types + if legacy_primitive_types is not None + else self.legacy_primitive_types + ), + stats_callback=stats_callback + ) + + async def _use_legacy_prepared_statements(self): + if self.legacy_prepared_statements is not None: + return self.legacy_prepared_statements + + value = must_use_legacy_prepared_statements.get((self.host, self.port)) + if value is None: + try: + query = AsyncTrinoQuery( + self._create_request(), + query="EXECUTE IMMEDIATE 'SELECT 1'") + await query.execute() + value = False + except Exception as e: + logger.warning( + "EXECUTE IMMEDIATE not available for %s:%s; defaulting to legacy prepared statements (%s)", + self.host, self.port, e) + value = True + must_use_legacy_prepared_statements.put((self.host, self.port), value) + return value + + +class Cursor(_SyncCursor): + """Asynchronous database cursor. Shares the parameter formatting and + description logic with :class:`trino.dbapi.Cursor`; the execute and fetch + methods are coroutines and rows are iterated with ``async for``.""" + + def __init__( + self, + connection, + request: AsyncTrinoRequest, + legacy_primitive_types: bool = False, + stats_callback: Optional[Callable[[Dict[str, Any]], None]] = None): + if not isinstance(connection, Connection): + raise ValueError( + "connection must be a trino.aio Connection object: {}".format(type(connection)) + ) + self._connection = connection + self._request = request + + self.arraysize = 1 + self._iterator = None + self._query = None + self._legacy_primitive_types = legacy_primitive_types + self._stats_callback = stats_callback + + def __iter__(self): + raise TypeError("use 'async for' to iterate an asynchronous cursor") + + def __aiter__(self): + return self._iterator + + def __enter__(self): + raise TypeError("use 'async with' with an asynchronous cursor") + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc_value, traceback): + await self.close() + + async def _prepare_statement(self, statement: str, name: str) -> None: + sql = f"PREPARE {name} FROM {statement}" + query = AsyncTrinoQuery(self.connection._create_request(), query=sql, + legacy_primitive_types=self._legacy_primitive_types) + await query.execute() + + def _execute_prepared_statement( + self, + statement_name, + params + ): + sql = 'EXECUTE ' + statement_name + ' USING ' + ','.join(map(self._format_prepared_param, params)) + return AsyncTrinoQuery( + self._request, + query=sql, + legacy_primitive_types=self._legacy_primitive_types, + stats_callback=self._stats_callback) + + def _execute_immediate_statement(self, statement: str, params): + sql = "EXECUTE IMMEDIATE '" + statement.replace("'", "''") + \ + "' USING " + ",".join(map(self._format_prepared_param, params)) + return AsyncTrinoQuery( + self.connection._create_request(), + query=sql, + legacy_primitive_types=self._legacy_primitive_types, + stats_callback=self._stats_callback) + + async def _deallocate_prepared_statement(self, statement_name: str) -> None: + sql = 'DEALLOCATE PREPARE ' + statement_name + query = AsyncTrinoQuery(self.connection._create_request(), query=sql, + legacy_primitive_types=self._legacy_primitive_types) + await query.execute() + + async def execute(self, operation, params=None): + if params: + assert isinstance(params, (list, tuple)), ( + 'params must be a list or tuple containing the query ' + 'parameter values' + ) + + if await self.connection._use_legacy_prepared_statements(): + statement_name = self._generate_unique_statement_name() + await self._prepare_statement(operation, statement_name) + + try: + # Send execute statement and assign the return value to `results` + # as it will be returned by the function + self._query = self._execute_prepared_statement( + statement_name, params + ) + self._iterator = (await self._query.execute()).__aiter__() + finally: + # Send deallocate statement + # At this point the query can be deallocated since it has already + # been executed + await self._deallocate_prepared_statement(statement_name) + else: + self._query = self._execute_immediate_statement(operation, params) + self._iterator = (await self._query.execute()).__aiter__() + + else: + self._query = AsyncTrinoQuery(self._request, query=operation, + legacy_primitive_types=self._legacy_primitive_types, + stats_callback=self._stats_callback) + self._iterator = (await self._query.execute()).__aiter__() + return self + + async def executemany(self, operation, seq_of_params): + """ + Prepare a database operation (query or command) and then execute it + against all parameter sequences found in ``seq_of_params``; see + :meth:`trino.dbapi.Cursor.executemany`. + """ + for parameters in seq_of_params[:-1]: + await self.execute(operation, parameters) + await self.fetchall() + if self._query.update_type is None: + raise NotSupportedError("Query must return update type") + if seq_of_params: + await self.execute(operation, seq_of_params[-1]) + else: + await self.execute(operation) + return self + + async def fetchone(self) -> Optional[List[Any]]: + """ + Fetch the next row of a query result set, returning a single + sequence, or None when no more data is available. + """ + try: + assert self._iterator is not None + return await self._iterator.__anext__() + except StopAsyncIteration: + return None + except trino.exceptions.HttpError as err: + raise trino.exceptions.OperationalError(str(err)) + + async def fetchmany(self, size=None) -> List[List[Any]]: + """ + Fetch the next set of rows of a query result, returning a list of + rows. An empty list is returned when no more rows are available. + """ + if size is None: + size = self.arraysize + + rows = [] + for _ in range(size): + row = await self.fetchone() + if row is None: + break + rows.append(row) + return rows + + async def describe(self, sql: str) -> List[DescribeOutput]: + """ + List the output columns of a SQL statement; see + :meth:`trino.dbapi.Cursor.describe`. + """ + statement_name = self._generate_unique_statement_name() + await self._prepare_statement(sql, statement_name) + try: + sql = f"DESCRIBE OUTPUT {statement_name}" + self._query = AsyncTrinoQuery( + self._request, + query=sql, + legacy_primitive_types=self._legacy_primitive_types, + ) + result = await self._query.execute() + rows = [row async for row in result] + finally: + await self._deallocate_prepared_statement(statement_name) + + return list(map(lambda x: DescribeOutput.from_row(x), rows)) + + async def fetchall(self) -> List[List[Any]]: + rows = [] + while True: + row = await self.fetchone() + if row is None: + return rows + rows.append(row) + + async def cancel(self): + if self._query is None: + return + await self._query.cancel() + + async def close(self): + await self.cancel() + # TODO: Cancel not only the last query executed on this cursor + # but also any other outstanding queries executed through this cursor. + + +class SegmentCursor(Cursor): + def __init__( + self, + connection, + request: AsyncTrinoRequest, + legacy_primitive_types: bool = False, + stats_callback: Optional[Callable[[Dict[str, Any]], None]] = None): + super().__init__( + connection, request, legacy_primitive_types=legacy_primitive_types, stats_callback=stats_callback) + if self.connection._client_session.encoding is None: + raise ValueError("SegmentCursor can only be used if encoding is set on the connection") + + async def execute(self, operation, params=None): + if params: + # TODO: refactor code to allow for params to be supported + raise ValueError("params not supported") + + self._query = AsyncTrinoQuery(self._request, query=operation, + legacy_primitive_types=self._legacy_primitive_types, + fetch_mode="segments", + stats_callback=self._stats_callback) + self._iterator = (await self._query.execute()).__aiter__() + return self diff --git a/trino/auth.py b/trino/auth.py index 783d0f5e..7abe32b8 100644 --- a/trino/auth.py +++ b/trino/auth.py @@ -10,12 +10,15 @@ # See the License for the specific language governing permissions and # limitations under the License. import abc +import asyncio import importlib import json import os import re import threading import webbrowser +from collections.abc import AsyncGenerator +from collections.abc import Generator from collections.abc import Mapping from typing import Any from typing import Callable @@ -23,17 +26,12 @@ from typing import List from typing import Optional from typing import Tuple -from urllib.parse import urlparse -from requests import PreparedRequest -from requests import Request -from requests import Response -from requests import Session -from requests.auth import AuthBase -from requests.auth import extract_cookies_to_jar +import httpx2 import trino.logging from trino import exceptions +from trino._spnego import SPNEGOAuth from trino.constants import HEADER_ORIGINAL_USER from trino.constants import HEADER_USER from trino.constants import MAX_NT_PASSWORD_SIZE @@ -42,13 +40,49 @@ class Authentication(metaclass=abc.ABCMeta): + """ + Extension point for Trino authentication mechanisms. + + httpx builds its TLS and environment configuration when a client is + constructed, so an authentication is asked for two things: + + - :meth:`get_client_arguments`: constructor arguments merged into the + ``httpx2.Client``/``httpx2.AsyncClient`` the connection creates + (``verify``, ``cert`` and ``trust_env``). + - :meth:`get_http_auth`: the ``httpx2.Auth`` instance attached to the + client, or ``None`` when the mechanism only needs client arguments. + """ + + def get_client_arguments(self) -> Dict[str, Any]: + return {} + @abc.abstractmethod - def set_http_session(self, http_session: Session) -> Session: + def get_http_auth(self) -> Optional[httpx2.Auth]: pass def get_exceptions(self) -> Tuple[Any, ...]: return tuple() + def set_http_session(self, http_session: Any) -> Any: + raise NotImplementedError( + "set_http_session was removed when the client migrated from requests to httpx2. " + "Implement get_http_auth() (returning an httpx2.Auth) and, for verify/cert/trust_env, " + "get_client_arguments() instead." + ) + + +def _gssapi_credentials(principal: Optional[str]) -> Any: + if principal: + try: + import gssapi + except ImportError: + raise RuntimeError("unable to import gssapi") + + name = gssapi.Name(principal, gssapi.NameType.user) + return gssapi.Credentials(name=name, usage="initiate") + + return None + class KerberosAuthentication(Authentication): MUTUAL_REQUIRED = 1 @@ -77,35 +111,27 @@ def __init__( self._delegate = delegate self._ca_bundle = ca_bundle - def set_http_session(self, http_session: Session) -> Session: - try: - import requests_kerberos - except ImportError: - raise RuntimeError("unable to import requests_kerberos") - + def get_client_arguments(self) -> Dict[str, Any]: if self._config: os.environ["KRB5_CONFIG"] = self._config - http_session.trust_env = False - http_session.auth = requests_kerberos.HTTPKerberosAuth( - mutual_authentication=self._mutual_authentication, - force_preemptive=self._force_preemptive, + arguments: Dict[str, Any] = {"trust_env": False} + if self._ca_bundle: + arguments["verify"] = self._ca_bundle + return arguments + + def get_http_auth(self) -> Optional[httpx2.Auth]: + return SPNEGOAuth( + service_name=self._service_name, hostname_override=self._hostname_override, - sanitize_mutual_error_response=self._sanitize_mutual_error_response, - principal=self._principal, + mutual_authentication=self._mutual_authentication, + opportunistic_auth=self._force_preemptive, delegate=self._delegate, - service=self._service_name, + creds=_gssapi_credentials(self._principal), + sanitize_mutual_error_response=self._sanitize_mutual_error_response, ) - if self._ca_bundle: - http_session.verify = self._ca_bundle - return http_session def get_exceptions(self) -> Tuple[Any, ...]: - try: - from requests_kerberos.exceptions import KerberosExchangeError - - return KerberosExchangeError, - except ImportError: - raise RuntimeError("unable to import requests_kerberos") + return exceptions.SPNEGOExchangeError, def __eq__(self, other: object) -> bool: if not isinstance(other, KerberosAuthentication): @@ -148,38 +174,26 @@ def __init__( self._delegate = delegate self._ca_bundle = ca_bundle - def set_http_session(self, http_session: Session) -> Session: - try: - import requests_gssapi - except ImportError: - raise RuntimeError("unable to import requests_gssapi") - + def get_client_arguments(self) -> Dict[str, Any]: if self._config: os.environ["KRB5_CONFIG"] = self._config - http_session.trust_env = False - http_session.auth = requests_gssapi.HTTPSPNEGOAuth( + arguments: Dict[str, Any] = {"trust_env": False} + if self._ca_bundle: + arguments["verify"] = self._ca_bundle + return arguments + + def get_http_auth(self) -> Optional[httpx2.Auth]: + return SPNEGOAuth( + target_name=self._get_target_name(self._hostname_override, self._service_name), mutual_authentication=self._mutual_authentication, opportunistic_auth=self._force_preemptive, - target_name=self._get_target_name(self._hostname_override, self._service_name), - sanitize_mutual_error_response=self._sanitize_mutual_error_response, - creds=self._get_credentials(self._principal), delegate=self._delegate, + creds=self._get_credentials(self._principal), + sanitize_mutual_error_response=self._sanitize_mutual_error_response, ) - if self._ca_bundle: - http_session.verify = self._ca_bundle - return http_session def _get_credentials(self, principal: Optional[str] = None) -> Any: - if principal: - try: - import gssapi - except ImportError: - raise RuntimeError("unable to import gssapi") - - name = gssapi.Name(principal, gssapi.NameType.user) - return gssapi.Credentials(name=name, usage="initiate") - - return None + return _gssapi_credentials(principal) def _get_target_name( self, @@ -201,12 +215,7 @@ def _get_target_name( return hostname_override def get_exceptions(self) -> Tuple[Any, ...]: - try: - from requests_gssapi.exceptions import SPNEGOExchangeError - - return SPNEGOExchangeError, - except ImportError: - raise RuntimeError("unable to import requests_kerberos") + return exceptions.SPNEGOExchangeError, def __eq__(self, other: object) -> bool: if not isinstance(other, GSSAPIAuthentication): @@ -227,14 +236,8 @@ def __init__(self, username: str, password: str): self._username = username self._password = password - def set_http_session(self, http_session: Session) -> Session: - try: - import requests.auth - except ImportError: - raise RuntimeError("unable to import requests.auth") - - http_session.auth = requests.auth.HTTPBasicAuth(self._username, self._password) - return http_session + def get_http_auth(self) -> Optional[httpx2.Auth]: + return httpx2.BasicAuth(self._username, self._password) def get_exceptions(self) -> Tuple[Any, ...]: return () @@ -245,7 +248,7 @@ def __eq__(self, other: object) -> bool: return self._username == other._username and self._password == other._password -class _BearerAuth(AuthBase): +class _BearerAuth(httpx2.Auth): """ Custom implementation of Authentication class for bearer token """ @@ -253,9 +256,9 @@ class _BearerAuth(AuthBase): def __init__(self, token: str): self.token = token - def __call__(self, r: PreparedRequest) -> PreparedRequest: - r.headers["Authorization"] = "Bearer " + self.token - return r + def auth_flow(self, request: httpx2.Request) -> Generator[httpx2.Request, httpx2.Response, None]: + request.headers["Authorization"] = "Bearer " + self.token + yield request class JWTAuthentication(Authentication): @@ -263,9 +266,8 @@ class JWTAuthentication(Authentication): def __init__(self, token: str): self.token = token - def set_http_session(self, http_session: Session) -> Session: - http_session.auth = _BearerAuth(self.token) - return http_session + def get_http_auth(self) -> Optional[httpx2.Auth]: + return _BearerAuth(self.token) def get_exceptions(self) -> Tuple[Any, ...]: return () @@ -414,7 +416,13 @@ def store_token_to_cache(self, key: Optional[str], token: str) -> None: "information.") from e -class _OAuth2TokenBearer(AuthBase): +# Sentinel yielded by the OAuth2 core flow instead of a request when another +# thread or task is already running the OAuth2 exchange; the drivers translate +# it into an appropriate (blocking or event-loop-friendly) wait. +_WAIT_FOR_OAUTH = object() + + +class _OAuth2TokenBearer(httpx2.Auth): """ Custom implementation of Trino OAuth2 based authentication to get the token """ @@ -424,44 +432,89 @@ class _OAuth2TokenBearer(AuthBase): def __init__(self, redirect_auth_url_handler: Callable[[str], None]): self._redirect_auth_url = redirect_auth_url_handler keyring_cache = _OAuth2KeyRingTokenCache() - self._token_cache = keyring_cache if keyring_cache.is_keyring_available() else _OAuth2TokenInMemoryCache() + self._token_cache: _OAuth2TokenCache = \ + keyring_cache if keyring_cache.is_keyring_available() else _OAuth2TokenInMemoryCache() self._token_lock = threading.Lock() + # Held by the thread/task currently performing the OAuth2 exchange so + # concurrent requests wait for its token instead of also opening a + # browser window. self._inside_oauth_attempt_lock = threading.Lock() - self._inside_oauth_attempt_blocker = threading.Event() - def __call__(self, r: PreparedRequest) -> PreparedRequest: - host = self._determine_host(r.url) - user = self._determine_user(r.headers) + def sync_auth_flow( + self, request: httpx2.Request + ) -> Generator[httpx2.Request, httpx2.Response, None]: + flow = self._flow(request) + item = next(flow) + while True: + if item is _WAIT_FOR_OAUTH: + # Block until the thread running the OAuth2 exchange finishes. + with self._inside_oauth_attempt_lock: + pass + item = flow.send(None) + continue + response = yield item + response.read() + try: + item = flow.send(response) + except StopIteration: + return + + async def async_auth_flow( + self, request: httpx2.Request + ) -> AsyncGenerator[httpx2.Request, httpx2.Response]: + flow = self._flow(request) + item = next(flow) + while True: + if item is _WAIT_FOR_OAUTH: + # Never block the event loop; the lock may be held by another + # thread, so poll it instead of awaiting a loop-bound primitive. + while self._inside_oauth_attempt_lock.locked(): + await asyncio.sleep(0.05) + item = flow.send(None) + continue + response = yield item + await response.aread() + try: + item = flow.send(response) + except StopIteration: + return + + def _flow(self, request: httpx2.Request) -> Generator[Any, Any, None]: + host = request.url.host + user = self._determine_user(request.headers) key = self._construct_cache_key(host, user) token = self._get_token_from_cache(key) if token is not None: - r.headers['Authorization'] = "Bearer " + token - - r.register_hook('response', self._authenticate) - - return r + request.headers["Authorization"] = "Bearer " + token - def _authenticate(self, response: Response, **kwargs: Any) -> Optional[Response]: + response = yield request if not 400 <= response.status_code < 500: - return response + return - acquired = self._inside_oauth_attempt_lock.acquire(blocking=False) - if acquired: + if self._inside_oauth_attempt_lock.acquire(blocking=False): try: # Lock is acquired, attempt the OAuth2 flow - self._attempt_oauth(response, **kwargs) - self._inside_oauth_attempt_blocker.set() + token = yield from self._attempt_oauth(response) + self._store_token_to_cache(key, token) finally: self._inside_oauth_attempt_lock.release() - self._inside_oauth_attempt_blocker.clear() else: - # Lock is not acquired, we are already in the OAuth2 flow, so we block until OAuth2 flow is finished. - self._inside_oauth_attempt_blocker.wait() + # We are already in the OAuth2 flow on another thread or task; + # wait until it finishes and pick up the token it cached. + yield _WAIT_FOR_OAUTH + token = self._get_token_from_cache(key) - return self._retry_request(response, **kwargs) + # Retry the original request with the fresh token, carrying over any + # cookies the failed response may have set. + if token is not None: + request.headers["Authorization"] = "Bearer " + token + cookies = httpx2.Cookies() + cookies.extract_cookies(response) + cookies.set_cookie_header(request) + yield request - def _attempt_oauth(self, response: Response, **kwargs: Any) -> None: + def _attempt_oauth(self, response: httpx2.Response) -> Generator[Any, Any, str]: # we have to handle the authentication, may be token the token expired, or it wasn't there at all auth_info = response.headers.get('WWW-Authenticate') if not auth_info: @@ -485,57 +538,32 @@ def _attempt_oauth(self, response: Response, **kwargs: Any) -> None: # tell app that use this url to proceed with the authentication self._redirect_auth_url(auth_server) - # Consume content and release the original connection - # to allow our new request to reuse the same one. - response.content - response.close() - - token = self._get_token(token_server, response, **kwargs) + # Token polls reuse the timeout the original request was sent with. + extensions = {} + timeout = response.request.extensions.get("timeout") + if timeout is not None: + extensions["timeout"] = timeout - request = response.request - host = self._determine_host(request.url) - user = self._determine_user(request.headers) - key = self._construct_cache_key(host, user) - self._store_token_to_cache(key, token) - - def _retry_request(self, response: Response, **kwargs: Any) -> Optional[Response]: - request = response.request.copy() - extract_cookies_to_jar(request._cookies, response.request, response.raw) - request.prepare_cookies(request._cookies) - - host = self._determine_host(response.request.url) - user = self._determine_user(request.headers) - key = self._construct_cache_key(host, user) - token = self._get_token_from_cache(key) - if token is not None: - request.headers['Authorization'] = "Bearer " + token - retry_response = response.connection.send(request, **kwargs) - retry_response.history.append(response) - retry_response.request = request - return retry_response - - def _get_token(self, token_server: str, response: Response, **kwargs: Any) -> str: attempts = 0 while attempts < self.MAX_OAUTH_ATTEMPTS: attempts += 1 - with response.connection.send(Request( - method='GET', url=token_server).prepare(), **kwargs) as response: - if response.status_code == 200: - token_response = json.loads(response.text) - token = token_response.get('token') - if token: - return token - error = token_response.get('error') - if error: - raise exceptions.TrinoAuthError(f"Error while getting the token: {error}") - else: - token_server = token_response.get('nextUri') - logger.debug(f"nextURi auth token server: {token_server}") + token_response = yield httpx2.Request("GET", token_server, extensions=extensions) + if token_response.status_code == 200: + body = json.loads(token_response.text) + token = body.get('token') + if token: + return token + error = body.get('error') + if error: + raise exceptions.TrinoAuthError(f"Error while getting the token: {error}") else: - raise exceptions.TrinoAuthError( - f"Error while getting the token response " - f"status code: {response.status_code}, " - f"body: {response.text}") + token_server = body.get('nextUri') + logger.debug(f"nextURi auth token server: {token_server}") + else: + raise exceptions.TrinoAuthError( + f"Error while getting the token response " + f"status code: {token_response.status_code}, " + f"body: {token_response.text}") raise exceptions.TrinoAuthError("Exceeded max attempts while getting the token") @@ -547,10 +575,6 @@ def _store_token_to_cache(self, key: Optional[str], token: str) -> None: with self._token_lock: self._token_cache.store_token_to_cache(key, token) - @staticmethod - def _determine_host(url: Optional[str]) -> Any: - return urlparse(url).hostname - @staticmethod def _determine_user(headers: Mapping[Any, Any]) -> Optional[Any]: return headers.get(HEADER_ORIGINAL_USER, headers.get(HEADER_USER)) @@ -586,9 +610,8 @@ def __init__(self, redirect_auth_url_handler: CompositeRedirectHandler = Composi self._redirect_auth_url = redirect_auth_url_handler self._bearer = _OAuth2TokenBearer(self._redirect_auth_url) - def set_http_session(self, http_session: Session) -> Session: - http_session.auth = self._bearer - return http_session + def get_http_auth(self) -> Optional[httpx2.Auth]: + return self._bearer def get_exceptions(self) -> Tuple[Any, ...]: return () @@ -604,9 +627,11 @@ def __init__(self, cert: str, key: str): self._cert = cert self._key = key - def set_http_session(self, http_session: Session) -> Session: - http_session.cert = (self._cert, self._key) - return http_session + def get_client_arguments(self) -> Dict[str, Any]: + return {"cert": (self._cert, self._key)} + + def get_http_auth(self) -> Optional[httpx2.Auth]: + return None def get_exceptions(self) -> Tuple[Any, ...]: return () diff --git a/trino/client.py b/trino/client.py index 9d4956bf..ffede7e3 100644 --- a/trino/client.py +++ b/trino/client.py @@ -34,25 +34,14 @@ """ from __future__ import annotations -import abc import atexit -import base64 import copy import functools import itertools import os -import random -import re import threading -import urllib.parse -import warnings -from abc import abstractmethod from collections.abc import Iterator from concurrent.futures import ThreadPoolExecutor -from dataclasses import dataclass -from datetime import datetime -from email.utils import parsedate_to_datetime -from enum import Enum from time import sleep from typing import Any from typing import Callable @@ -62,46 +51,50 @@ from typing import Literal from typing import Optional from typing import Tuple -from typing import TypedDict from typing import Union -from zoneinfo import ZoneInfo - -try: - import lz4.block -except ImportError as err: - _LZ4_ERROR = str(err) -else: - _LZ4_ERROR = None - -try: - import orjson as json -except ImportError: - import json - -import requests -from requests import Response -from requests import Session -from requests.structures import CaseInsensitiveDict - -try: - import zstandard -except ImportError as err: - _ZSTD_ERROR = str(err) -else: - _ZSTD_ERROR = None +import httpx2 +from httpx2 import Client +from httpx2 import Response import trino.logging from trino import constants from trino import exceptions -from trino._version import __version__ +from trino._protocol import _DelayExponential +from trino._protocol import _InlineSegmentTO # noqa: F401 re-export +from trino._protocol import _parse_retry_after_header # noqa: F401 re-export +from trino._protocol import _SegmentMetadataTO # noqa: F401 re-export +from trino._protocol import _SegmentTO # noqa: F401 re-export +from trino._protocol import _SpooledProtocolResponseTO +from trino._protocol import _SpooledSegmentTO +from trino._protocol import _TrinoQueryBase +from trino._protocol import _TrinoRequestBase +from trino._protocol import CaseInsensitiveDict # noqa: F401 re-export +from trino._protocol import ClientSession +from trino._protocol import CODECS_UNAVAILABLE # noqa: F401 re-export +from trino._protocol import CompressedQueryDataDecoder # noqa: F401 re-export +from trino._protocol import CompressedQueryDataDecoderFactory +from trino._protocol import DecodableSegment +from trino._protocol import ENCODINGS # noqa: F401 re-export +from trino._protocol import get_header_values # noqa: F401 re-export +from trino._protocol import get_prepared_statement_values # noqa: F401 re-export +from trino._protocol import get_roles_values # noqa: F401 re-export +from trino._protocol import get_session_property_values # noqa: F401 re-export +from trino._protocol import InlineSegment +from trino._protocol import JsonQueryDataDecoder # noqa: F401 re-export +from trino._protocol import Lz4QueryDataDecoder # noqa: F401 re-export +from trino._protocol import needs_retry +from trino._protocol import QueryDataDecoder +from trino._protocol import retry_after_seconds +from trino._protocol import ROLE_PATTERN # noqa: F401 re-export +from trino._protocol import Segment +from trino._protocol import SegmentType # noqa: F401 re-export +from trino._protocol import spooling_request_headers +from trino._protocol import TrinoStatus # noqa: F401 re-export +from trino._protocol import wire_headers +from trino._protocol import ZStdQueryDataDecoder # noqa: F401 re-export from trino.auth import Authentication -from trino.exceptions import TrinoExternalError -from trino.exceptions import TrinoQueryError -from trino.exceptions import TrinoUserError from trino.mapper import RowMapper -from trino.mapper import RowMapperFactory - __all__ = [ "ClientSession", @@ -126,301 +119,13 @@ def close_executor(): MAX_ATTEMPTS = constants.DEFAULT_MAX_ATTEMPTS SOCKS_PROXY = os.environ.get("SOCKS_PROXY") +# httpx configures proxies at client construction time, so PROXIES uses the +# httpx mounts-style single "all://" key instead of requests' per-scheme dict. if SOCKS_PROXY: - PROXIES = {"http": "socks5://" + SOCKS_PROXY, "https": "socks5://" + SOCKS_PROXY} + PROXIES = {"all://": "socks5://" + SOCKS_PROXY} else: PROXIES = {} -_HEADER_EXTRA_CREDENTIAL_KEY_REGEX = re.compile(r'^\S[^\s=]*$') - -ENCODINGS = ["json+zstd", "json+lz4", "json"] -CODECS_UNAVAILABLE = {} -if _LZ4_ERROR: - CODECS_UNAVAILABLE["lz4"] = _LZ4_ERROR -if _ZSTD_ERROR: - CODECS_UNAVAILABLE["zstd"] = _ZSTD_ERROR - -ROLE_PATTERN = re.compile(r"^ROLE\{(.*)\}$") - - -class ClientSession: - """ - Manage the current Client Session properties of a specific connection. This class is thread-safe. - - :param user: associated with the query. It is useful for access control - and query scheduling. - :param authorization_user: associated with the query. It is useful for access control - and query scheduling. - :param source: associated with the query. It is useful for access - control and query scheduling. - :param catalog: to query. The *catalog* is associated with a Trino - connector. This variable sets the default catalog used - by SQL statements. For example, if *catalog* is set - to ``some_catalog``, the SQL statement - ``SELECT * FROM some_schema.some_table`` will actually - query the table - ``some_catalog.some_schema.some_table``. - :param schema: to query. The *schema* is a logical abstraction to group - table. This variable sets the default schema used by - SQL statements. For example, if *schema* is set to - ``some_schema``, the SQL statement - ``SELECT * FROM some_table`` will actually query the - table ``some_catalog.some_schema.some_table``. - :param properties: set specific Trino behavior for the current - session. Please refer to the output of - ``SHOW SESSION`` to check the available - properties. - :param headers: HTTP headers to POST/GET in the HTTP requests - :param extra_credential: extra credentials. as list of ``(key, value)`` - tuples. - :param client_tags: Client tags as list of strings. - :param roles: roles for the current session. Some connectors do not - support role management. See connector documentation for more details. - :param timezone: The timezone for query processing. Defaults to the system's local timezone. - :param encoding: The encoding for the spooling protocol. Defaults to None. - """ - - def __init__( - self, - user: str, - authorization_user: Optional[str] = None, - catalog: Optional[str] = None, - schema: Optional[str] = None, - source: Optional[str] = None, - properties: Optional[Dict[str, str]] = None, - headers: Optional[Dict[str, str]] = None, - transaction_id: Optional[str] = None, - extra_credential: Optional[List[Tuple[str, str]]] = None, - client_tags: Optional[List[str]] = None, - roles: Optional[Union[Dict[str, str], str]] = None, - timezone: Optional[str] = None, - encoding: Optional[Union[str, List[str]]] = None, - heartbeat_interval: Optional[float] = constants.DEFAULT_HEARTBEAT_INTERVAL, - ): - self._object_lock = threading.Lock() - self._prepared_statements: Dict[str, str] = {} - - self._user = user - self._authorization_user = authorization_user - self._catalog = catalog - self._schema = schema - self._source = source - self._properties = properties.copy() if properties is not None else {} - self._headers = headers.copy() if headers is not None else {} - self._transaction_id = transaction_id - self._extra_credential = extra_credential - self._client_tags = client_tags.copy() if client_tags is not None else list() - self._roles = self._format_roles(roles) if roles is not None else {} - if timezone: # Check timezone validity - ZoneInfo(timezone) - self._timezone = timezone - else: - from tzlocal import get_localzone_name - self._timezone = get_localzone_name() - self._encoding = encoding - self._heartbeat_interval = heartbeat_interval - - @property - def user(self) -> str: - return self._user - - @property - def authorization_user(self) -> Optional[str]: - with self._object_lock: - return self._authorization_user - - @authorization_user.setter - def authorization_user(self, authorization_user: Optional[str]) -> None: - with self._object_lock: - self._authorization_user = authorization_user - - @property - def catalog(self) -> Optional[str]: - with self._object_lock: - return self._catalog - - @catalog.setter - def catalog(self, catalog: Optional[str]) -> None: - with self._object_lock: - self._catalog = catalog - - @property - def schema(self) -> Optional[str]: - with self._object_lock: - return self._schema - - @schema.setter - def schema(self, schema: Optional[str]) -> None: - with self._object_lock: - self._schema = schema - - @property - def source(self) -> Optional[str]: - return self._source - - @property - def properties(self) -> Dict[str, str]: - with self._object_lock: - return self._properties - - @properties.setter - def properties(self, properties: Dict[str, str]) -> None: - with self._object_lock: - self._properties = properties - - @property - def headers(self) -> Dict[str, str]: - return self._headers - - @property - def transaction_id(self) -> Optional[str]: - with self._object_lock: - return self._transaction_id - - @transaction_id.setter - def transaction_id(self, transaction_id: Optional[str]) -> None: - with self._object_lock: - self._transaction_id = transaction_id - - @property - def extra_credential(self) -> Optional[List[Tuple[str, str]]]: - return self._extra_credential - - @property - def client_tags(self) -> List[str]: - return self._client_tags - - @property - def roles(self) -> Dict[str, str]: - with self._object_lock: - return self._roles - - @roles.setter - def roles(self, roles: Dict[str, str]) -> None: - with self._object_lock: - self._roles = roles - - @property - def prepared_statements(self) -> Dict[str, str]: - return self._prepared_statements - - @prepared_statements.setter - def prepared_statements(self, prepared_statements: Dict[str, str]) -> None: - with self._object_lock: - self._prepared_statements = prepared_statements - - @property - def timezone(self) -> str: - with self._object_lock: - return self._timezone - - @property - def encoding(self) -> Optional[Union[str, List[str]]]: - with self._object_lock: - return self._encoding - - @property - def heartbeat_interval(self) -> Optional[float]: - return self._heartbeat_interval - - @staticmethod - def _format_roles(roles: Union[Dict[str, str], str]) -> Dict[str, str]: - if isinstance(roles, str): - roles = {"system": roles} - formatted_roles = {} - for catalog, role in roles.items(): - is_legacy_role_pattern = ROLE_PATTERN.match(role) is not None - if role in ("NONE", "ALL") or is_legacy_role_pattern: - if is_legacy_role_pattern: - warnings.warn(f"A role '{role}' is provided using a legacy format. " - "Please remove the ROLE{} wrapping. Support for the legacy format might be " - "removed in a future release.", - DeprecationWarning) - formatted_roles[catalog] = role - else: - formatted_roles[catalog] = f"ROLE{{{role}}}" - return formatted_roles - - def __getstate__(self): - state = self.__dict__.copy() - del state["_object_lock"] - return state - - def __setstate__(self, state): - self.__dict__.update(state) - self._object_lock = threading.Lock() - - -def get_header_values(headers: CaseInsensitiveDict[str], header: str) -> List[str]: - return [val.strip() for val in headers[header].split(",")] - - -def get_session_property_values(headers: CaseInsensitiveDict[str], header: str) -> List[Tuple[str, str]]: - kvs = get_header_values(headers, header) - return [ - (k.strip(), urllib.parse.unquote_plus(v.strip())) - for k, v in (kv.split("=", 1) for kv in kvs if kv) - ] - - -def get_prepared_statement_values(headers: CaseInsensitiveDict[str], header: str) -> List[Tuple[str, str]]: - kvs = get_header_values(headers, header) - return [ - (k.strip(), urllib.parse.unquote_plus(v.strip())) - for k, v in (kv.split("=", 1) for kv in kvs if kv) - ] - - -def get_roles_values(headers: CaseInsensitiveDict[str], header: str) -> List[Tuple[str, str]]: - kvs = get_header_values(headers, header) - return [ - (k.strip(), urllib.parse.unquote_plus(v.strip())) - for k, v in (kv.split("=", 1) for kv in kvs if kv) - ] - - -@dataclass -class TrinoStatus: - id: str - stats: Dict[str, str] - warnings: List[Any] - info_uri: str - next_uri: Optional[str] - update_type: Optional[str] - update_count: Optional[int] - rows: Union[List[Any], Dict[str, Any]] - columns: List[Any] - - def __repr__(self): - return ( - "TrinoStatus(" - "id={}, stats={{...}}, warnings={}, info_uri={}, next_uri={}, rows=" - ")".format( - self.id, - len(self.warnings), - self.info_uri, - self.next_uri, - len(self.rows), - ) - ) - - -class _DelayExponential: - def __init__( - self, base=0.1, exponent=2, jitter=True, max_delay=1800 # 100ms # 30 min - ): - self._base = base - self._exponent = exponent - self._jitter = jitter - self._max_delay = max_delay - - def __call__(self, attempt): - delay = float(self._base) * (self._exponent ** attempt) - if self._jitter: - delay *= random.random() - delay = min(float(self._max_delay), delay) - return delay - class _RetryWithExponentialBackoff: def __init__( @@ -441,7 +146,7 @@ def retry(self): sleep(self._retry_after_header) -class TrinoRequest: +class TrinoRequest(_TrinoRequestBase): """ Manage the HTTP requests of a Trino query. @@ -482,11 +187,10 @@ class TrinoRequest: the client. """ - http = requests + http = httpx2 - HTTP_EXCEPTIONS = ( - http.ConnectionError, - http.Timeout, + HTTP_EXCEPTIONS: Tuple[Any, ...] = ( + http.TransportError, ) def __init__( @@ -494,129 +198,93 @@ def __init__( host: str, port: int, client_session: ClientSession, - http_session: Optional[Session] = None, + http_session: Optional[Client] = None, http_scheme: Optional[str] = None, auth: Optional[Authentication] = constants.DEFAULT_AUTH, max_attempts: int = MAX_ATTEMPTS, request_timeout: Union[float, Tuple[float, float]] = constants.DEFAULT_REQUEST_TIMEOUT, handle_retry=_RetryWithExponentialBackoff(), - verify: bool = True, + verify: Union[bool, str] = True, ) -> None: - self._client_session = client_session - self._host = 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 - else: - self._http_scheme = http_scheme + super().__init__(host, port, client_session, http_scheme) if http_session is not None: self._http_session = http_session + if auth is not None: + self._apply_auth_to_existing_client(http_session, auth) else: - self._http_session = self.http.Session() - self._http_session.verify = verify - self._http_session.headers.update(self.http_headers) + self._http_session = self.create_http_client( + verify=verify, timeout=request_timeout, auth=auth + ) + # httpx clients expose no readable ``verify``; clients built by + # create_http_client carry the value used at construction time. + self._verify = getattr(self._http_session, "_trino_verify", verify) + self._http_session.headers.update(wire_headers(self.http_headers)) self._exceptions = self.HTTP_EXCEPTIONS self._auth = auth if self._auth: - self._auth.set_http_session(self._http_session) self._exceptions += self._auth.get_exceptions() - self._request_timeout = request_timeout + self._request_timeout = self.http.Timeout(request_timeout) self._handle_retry = handle_retry self.max_attempts = max_attempts - @property - def transaction_id(self) -> Optional[str]: - return self._client_session.transaction_id - - @transaction_id.setter - def transaction_id(self, value: Optional[str]) -> None: - self._client_session.transaction_id = value - - @property - def http_headers(self) -> CaseInsensitiveDict[str]: - headers: CaseInsensitiveDict[str] = CaseInsensitiveDict() - - headers[constants.HEADER_CATALOG] = self._client_session.catalog - headers[constants.HEADER_SCHEMA] = self._client_session.schema - headers[constants.HEADER_SOURCE] = self._client_session.source - if self._client_session.authorization_user is not None: - headers[constants.HEADER_ORIGINAL_USER] = self._client_session.user - headers[constants.HEADER_USER] = self._client_session.authorization_user - else: - headers[constants.HEADER_USER] = self._client_session.user - headers[constants.HEADER_TIMEZONE] = self._client_session.timezone - if self._client_session.encoding is None: - if not CODECS_UNAVAILABLE: - pass - else: - encoding = [ - enc - for enc in ENCODINGS - if (enc.split("+")[1] if "+" in enc else None) not in CODECS_UNAVAILABLE - ] - headers[constants.HEADER_ENCODING] = ",".join(encoding) - elif isinstance(self._client_session.encoding, list): - headers[constants.HEADER_ENCODING] = ",".join(self._client_session.encoding) - elif isinstance(self._client_session.encoding, str): - headers[constants.HEADER_ENCODING] = self._client_session.encoding - else: - raise ValueError("Invalid type for encoding: expected str or list") - headers[constants.HEADER_CLIENT_CAPABILITIES] = constants.CLIENT_CAPABILITIES - - headers["user-agent"] = f"{constants.CLIENT_NAME}/{__version__}" - if len(self._client_session.roles.values()): - headers[constants.HEADER_ROLE] = ",".join( - # ``name`` must not contain ``=`` - "{}={}".format(catalog, urllib.parse.quote(str(role))) - for catalog, role in self._client_session.roles.items() - ) - if self._client_session.client_tags is not None and len(self._client_session.client_tags) > 0: - headers[constants.HEADER_CLIENT_TAGS] = ",".join(self._client_session.client_tags) + @classmethod + def create_http_client( + cls, + verify: Union[bool, str] = True, + timeout: Union[float, Tuple[float, float], None] = constants.DEFAULT_REQUEST_TIMEOUT, + auth: Optional[Authentication] = None, + **kwargs: Any, + ) -> Client: + """ + Build the ``httpx2.Client`` used to talk to the coordinator. - headers[constants.HEADER_SESSION] = ",".join( - # ``name`` must not contain ``=`` - "{}={}".format(name, urllib.parse.quote(str(value))) - for name, value in self._client_session.properties.items() - ) + ``verify``, ``cert`` and ``trust_env`` can only be set when an httpx + client is constructed, so authentication implementations contribute + constructor arguments here through ``Authentication.get_client_arguments``. + """ + client_kwargs: Dict[str, Any] = { + "verify": verify, + # HTTP/2 is negotiated via ALPN on TLS connections; plain HTTP and + # servers without h2 support silently fall back to HTTP/1.1. + "http2": True, + # requests followed redirects on GET/POST/DELETE by default; httpx does not. + "follow_redirects": True, + "timeout": cls.http.Timeout(timeout), + } + if PROXIES: + client_kwargs["proxy"] = PROXIES.get("all://") + auth_arguments: Dict[str, Any] = {} + http_auth = None + if auth is not None: + auth_arguments = auth.get_client_arguments() + client_kwargs.update(auth_arguments) + http_auth = auth.get_http_auth() + client_kwargs.update(kwargs) + client = cls.http.Client(auth=http_auth, **client_kwargs) + client._trino_verify = client_kwargs["verify"] + client._trino_client_arguments = frozenset(auth_arguments) + return client - if len(self._client_session.prepared_statements) != 0: - # ``name`` must not contain ``=`` - headers[constants.HEADER_PREPARED_STATEMENT] = ",".join( - "{}={}".format(name, urllib.parse.quote_plus(statement)) - for name, statement in self._client_session.prepared_statements.items() + @staticmethod + def _apply_auth_to_existing_client(http_session: Client, auth: Authentication) -> None: + """ + Attach ``auth`` to an already-constructed client. Only ``client.auth`` + is settable after construction; when the authentication needs + constructor-only options (verify/cert/trust_env) that the client was + not built with, fail loudly instead of silently ignoring them. + """ + required = auth.get_client_arguments() + provided = getattr(http_session, "_trino_client_arguments", frozenset()) + missing = set(required) - set(provided) + if missing: + raise exceptions.TrinoConnectionError( + f"{type(auth).__name__} requires HTTP client construction options {sorted(missing)}; " + "configure them on your own httpx2 client, or omit http_session to let the " + "client be created for you." ) - - # merge custom http headers - for key in self._client_session.headers: - if key in headers.keys(): - raise ValueError("cannot override reserved HTTP header {}".format(key)) - headers.update(self._client_session.headers) - - transaction_id = self._client_session.transaction_id - headers[constants.HEADER_TRANSACTION] = transaction_id - - if self._client_session.extra_credential is not None and \ - len(self._client_session.extra_credential) > 0: - - for tup in self._client_session.extra_credential: - self._verify_extra_credential(tup) - - # HTTP 1.1 section 4.2 combine multiple extra credentials into a - # comma-separated value - # extra credential value is encoded per spec (application/x-www-form-urlencoded MIME format) - headers[constants.HEADER_EXTRA_CREDENTIAL] = \ - ", ".join( - [f"{tup[0]}={urllib.parse.quote_plus(str(tup[1]))}" - for tup in self._client_session.extra_credential]) - - return headers + http_session.auth = auth.get_http_auth() def unauthenticated(self): return TrinoRequest( @@ -626,7 +294,7 @@ def unauthenticated(self): request_timeout=self._request_timeout, handle_retry=self._handle_retry, client_session=ClientSession(user=self._client_session.user), - verify=self._http_session.verify) + verify=self._verify) @property def max_attempts(self) -> int: @@ -645,14 +313,9 @@ def max_attempts(self, value: int) -> None: with_retry = _retry_with( self._handle_retry, handled_exceptions=self._exceptions, - conditions=( - # need retry when there is no exception but the status code is 429, 502, 503, or 504 - lambda response: getattr(response, "status_code", None) - in (429, 502, 503, 504), - # need retry when the server returns 200 with an empty body (transient under load) - lambda response: getattr(response, "status_code", None) == 200 - and not getattr(response, "text", "").strip(), - ), + # Retry when there is no exception but the response is a transient + # error status or an empty 200 body; see _protocol.needs_retry. + conditions=(needs_retry,), max_attempts=self._max_attempts, ) self._get = with_retry(self._http_session.get) @@ -660,19 +323,6 @@ def max_attempts(self, value: int) -> None: self._delete = with_retry(self._http_session.delete) 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 - ) - - @property - def statement_url(self) -> str: - return self.get_url(constants.URL_STATEMENT_PATH) - - @property - def next_uri(self) -> Optional[str]: - return self._next_uri - def post(self, sql: str, additional_http_headers: Optional[Dict[str, Any]] = None) -> Response: data = sql.encode("utf-8") # Deep copy of the http_headers dict since they may be modified for this @@ -688,146 +338,31 @@ def post(self, sql: str, additional_http_headers: Optional[Dict[str, Any]] = Non http_response = self._post( self.statement_url, - data=data, - headers=http_headers, + content=data, + headers=wire_headers(http_headers), timeout=self._request_timeout, - proxies=PROXIES, ) return http_response def get(self, url: str) -> Response: return self._get( url, - headers=self.http_headers, + headers=wire_headers(self.http_headers), timeout=self._request_timeout, - proxies=PROXIES, ) def delete(self, url: str) -> Response: - return self._delete(url, timeout=self._request_timeout, proxies=PROXIES) + return self._delete(url, timeout=self._request_timeout) def head(self, url: str) -> Response: return self._head( url, - headers=self.http_headers, + headers=wire_headers(self.http_headers), timeout=self._request_timeout, - proxies=PROXIES, - ) - - @staticmethod - def _process_error(error, query_id: Optional[str]) -> Union[TrinoExternalError, TrinoQueryError, TrinoUserError]: - error_type = error["errorType"] - if error_type == "EXTERNAL": - raise exceptions.TrinoExternalError(error, query_id) - elif error_type == "USER_ERROR": - return exceptions.TrinoUserError(error, query_id) - - return exceptions.TrinoQueryError(error, query_id) - - @staticmethod - def raise_response_error(http_response: Response) -> None: - if http_response.status_code == 502: - raise exceptions.Http502Error("error 502: bad gateway") - - if http_response.status_code == 503: - raise exceptions.Http503Error("error 503: service unavailable") - - if http_response.status_code == 504: - raise exceptions.Http504Error("error 504: gateway timeout") - - raise exceptions.HttpError( - "error {}{}".format( - http_response.status_code, - ": {}".format(http_response.content) if http_response.content else "", - ) + # requests never followed redirects for HEAD; keep that behavior. + follow_redirects=False, ) - def process(self, http_response: Response) -> TrinoStatus: - if not http_response.ok: - self.raise_response_error(http_response) - - http_response.encoding = "utf-8" - if not http_response.text.strip(): - raise exceptions.TrinoConnectionError( - "received empty response from server (status 200)" - ) - response = json.loads(http_response.text) - if "error" in response and response["error"]: - raise self._process_error(response["error"], response.get("id")) - - if constants.HEADER_CLEAR_SESSION in http_response.headers: - for prop in get_header_values( - http_response.headers, constants.HEADER_CLEAR_SESSION - ): - self._client_session.properties.pop(prop, None) - - if constants.HEADER_SET_SESSION in http_response.headers: - for key, value in get_session_property_values( - http_response.headers, constants.HEADER_SET_SESSION - ): - self._client_session.properties[key] = value - - if constants.HEADER_SET_CATALOG in http_response.headers: - self._client_session.catalog = http_response.headers[constants.HEADER_SET_CATALOG] - - if constants.HEADER_SET_SCHEMA in http_response.headers: - self._client_session.schema = http_response.headers[constants.HEADER_SET_SCHEMA] - - if constants.HEADER_SET_ROLE in http_response.headers: - for key, value in get_roles_values( - http_response.headers, constants.HEADER_SET_ROLE - ): - self._client_session.roles[key] = value - - if constants.HEADER_ADDED_PREPARE in http_response.headers: - for name, statement in get_prepared_statement_values( - http_response.headers, constants.HEADER_ADDED_PREPARE - ): - self._client_session.prepared_statements[name] = statement - - if constants.HEADER_DEALLOCATED_PREPARE in http_response.headers: - for name in get_header_values( - http_response.headers, constants.HEADER_DEALLOCATED_PREPARE - ): - self._client_session.prepared_statements.pop(name, None) - - if constants.HEADER_SET_AUTHORIZATION_USER in http_response.headers: - self._client_session.authorization_user = http_response.headers[constants.HEADER_SET_AUTHORIZATION_USER] - - if constants.HEADER_RESET_AUTHORIZATION_USER in http_response.headers: - self._client_session.authorization_user = None - - self._next_uri = response.get("nextUri") - - data = response.get("data") if response.get("data") else [] - - return TrinoStatus( - id=response["id"], - stats=response["stats"], - warnings=response.get("warnings", []), - info_uri=response["infoUri"], - next_uri=self._next_uri, - update_type=response.get("updateType"), - update_count=response.get("updateCount"), - rows=data, - columns=response.get("columns"), - ) - - @staticmethod - def _verify_extra_credential(header: Tuple[str, str]) -> None: - """ - Verifies that key has ASCII only and non-whitespace characters. - """ - key = header[0] - - if not _HEADER_EXTRA_CREDENTIAL_KEY_REGEX.match(key): - raise ValueError(f"whitespace or '=' are disallowed in extra credential '{key}'") - - try: - key.encode().decode('ascii') - except UnicodeDecodeError: - raise ValueError(f"only ASCII characters are allowed in extra credential '{key}'") - class TrinoResult: """ @@ -889,7 +424,7 @@ def __next__(self): return row -class TrinoQuery: +class TrinoQuery(_TrinoQueryBase): """Represent the execution of a SQL statement by Trino.""" def __init__( @@ -900,31 +435,9 @@ def __init__( fetch_mode: Literal["mapped", "segments"] = "mapped", stats_callback: Optional[Callable[[Dict[str, Any]], None]] = None ) -> None: - self._query_id: Optional[str] = None - self._stats: Dict[Any, Any] = {} - self._info_uri: Optional[str] = None - self._warnings: List[Dict[Any, Any]] = [] - self._columns: Optional[List[str]] = None - self._finished = False - self._cancelled = False + super().__init__(query, legacy_primitive_types, fetch_mode, stats_callback) self._request = request - self._update_type = None - self._update_count = None - self._next_uri = None - self._query = query self._result: Optional[TrinoResult] = None - self._legacy_primitive_types = legacy_primitive_types - self._row_mapper: Optional[RowMapper] = None - self._fetch_mode = fetch_mode - self._stats_callback = stats_callback - - @property - def query_id(self) -> Optional[str]: - return self._query_id - - @property - def query(self) -> Optional[str]: - return self._query @property def columns(self): @@ -949,30 +462,10 @@ def columns(self): self._result.rows = [] return self._columns - @property - def stats(self): - return self._stats - - @property - def update_type(self): - return self._update_type - - @property - def update_count(self): - return self._update_count - - @property - def warnings(self): - return self._warnings - @property def result(self): return self._result - @property - def info_uri(self): - return self._info_uri - def execute(self, additional_http_headers=None) -> TrinoResult: """Initiate a Trino query by sending the SQL statement @@ -986,7 +479,7 @@ def execute(self, additional_http_headers=None) -> TrinoResult: try: response = self._request.post(self._query, additional_http_headers) - except requests.exceptions.RequestException as e: + except httpx2.HTTPError as e: raise trino.exceptions.TrinoConnectionError("failed to execute: {}".format(e)) status = self._request.process(response) self._info_uri = status.info_uri @@ -1040,28 +533,11 @@ def execute(self, additional_http_headers=None) -> TrinoResult: return self._result - def _update_state(self, status): - self._stats.update(status.stats) - self._update_type = status.update_type - self._update_count = status.update_count - self._next_uri = status.next_uri - if not self._row_mapper and status.columns: - self._row_mapper = RowMapperFactory().create(columns=status.columns, - legacy_primitive_types=self._legacy_primitive_types) - if status.columns: - self._columns = status.columns - self._report_stats() - - def _report_stats(self) -> None: - if self._stats_callback is not None: - # Pass a deep copy so the callback cannot mutate internal query state. - self._stats_callback(copy.deepcopy(self._stats)) - def fetch(self) -> Union[List[Union[List[Any], Any]], Iterator[List[Any]]]: """Continue fetching data for the current query_id""" try: response = self._request.get(self._request.next_uri) - except requests.exceptions.RequestException as e: + except httpx2.HTTPError as e: raise trino.exceptions.TrinoConnectionError("failed to fetch: {}".format(e)) status = self._request.process(response) self._update_state(status) @@ -1090,27 +566,13 @@ def fetch(self) -> Union[List[Union[List[Any], Any]], Iterator[List[Any]]]: else: raise ValueError(f"Unexpected type: {type(status.rows)}") - def _to_segments(self, rows: _SpooledProtocolResponseTO) -> List[DecodableSegment]: - encoding = rows["encoding"] - metadata = rows["metadata"] if "metadata" in rows else None - segments = [] - for segment in rows["segments"]: - segment_type = segment["type"] - if segment_type == SegmentType.INLINE: - inline_segment = cast(_InlineSegmentTO, segment) - segments.append(InlineSegment(inline_segment)) - elif segment_type == SegmentType.SPOOLED: - spooled_segment = cast(_SpooledSegmentTO, segment) - segments.append(SpooledSegment( - spooled_segment, - self._request.unauthenticated(), - coordinator_host=self._request._host, - custom_headers=dict(self._request._client_session.headers), - )) - else: - raise ValueError(f"Unsupported segment type: {segment_type}") - - return list(map(lambda segment: DecodableSegment(encoding, metadata, segment), segments)) + def _create_spooled_segment(self, segment: _SpooledSegmentTO) -> Segment: + return SpooledSegment( + segment, + self._request.unauthenticated(), + coordinator_host=self._request._host, + custom_headers=dict(self._request._client_session.headers), + ) def cancel(self) -> None: """Cancel the current query""" @@ -1120,28 +582,15 @@ def cancel(self) -> None: logger.debug("cancelling query: %s", self.query_id) try: response = self._request.delete(self._next_uri) - except requests.exceptions.RequestException as e: + except httpx2.HTTPError as e: raise trino.exceptions.TrinoConnectionError("failed to cancel query: {}".format(e)) - if response.status_code == requests.codes.no_content: + if response.status_code == httpx2.codes.NO_CONTENT: self._cancelled = True logger.debug("query cancelled: %s", self.query_id) return self._request.raise_response_error(response) - def is_finished(self) -> bool: - import warnings - warnings.warn("is_finished is deprecated, use finished instead", DeprecationWarning) - return self.finished - - @property - def finished(self) -> bool: - return self._finished - - @property - def cancelled(self) -> bool: - return self._cancelled - def _retry_with(handle_retry, handled_exceptions, conditions, max_attempts): def wrapper(func): @@ -1153,8 +602,8 @@ def decorated(*args, **kwargs): try: result = func(*args, **kwargs) if any(guard(result) for guard in conditions): - if result.status_code == 429 and "Retry-After" in result.headers: - retry_after = _parse_retry_after_header(result.headers.get("Retry-After")) + retry_after = retry_after_seconds(result) + if retry_after is not None: handle_retry_sleep = _RetryAfterSleep(retry_after) handle_retry_sleep.retry() else: @@ -1177,91 +626,6 @@ def decorated(*args, **kwargs): return wrapper -def _parse_retry_after_header(retry_after): - if isinstance(retry_after, int): - return retry_after - elif isinstance(retry_after, str) and retry_after.isdigit(): - return int(retry_after) - else: - retry_date = parsedate_to_datetime(retry_after) - now = datetime.utcnow() - return (retry_date - now).total_seconds() - - -# Trino Spooled protocol transfer objects -class _SpooledProtocolResponseTO(TypedDict): - encoding: Literal["json", "json+std", "json+lz4"] - metadata: _SegmentMetadataTO - segments: List[_SegmentTO] - - -class _SegmentMetadataTO(TypedDict): - uncompressedSize: str - segmentSize: str - - -class _SegmentTO(_SegmentMetadataTO): - type: Literal["spooled", "inline"] - metadata: _SegmentMetadataTO - - -class _SpooledSegmentTO(_SegmentTO): - uri: str - ackUri: str - headers: Dict[str, List[str]] - - -class _InlineSegmentTO(_SegmentTO): - data: str - - -class SegmentType(str, Enum): - """Enum with string values that can be compared to strings.""" - INLINE = "inline" - SPOOLED = "spooled" - - -class Segment(abc.ABC): - """ - Abstract base class representing a segment of data produced by the spooling protocol. - - Attributes: - metadata (property): Metadata associated with the segment. - rows (property): Returns the decoded and mapped data. - """ - def __init__(self, segment: _SegmentTO) -> None: - self._segment = segment - - @property - @abstractmethod - def data(self): - pass - - @property - def metadata(self) -> _SegmentMetadataTO: - return self._segment["metadata"] - - -class InlineSegment(Segment): - """ - A subclass of Segment that handles inline data segments. The data is base64 encoded and - requires mapping to rows using the provided row_mapper. - - Attributes: - rows (property): The data in the segment, decoded and mapped from the base64 encoded data. - """ - def __init__(self, segment: _InlineSegmentTO) -> None: - super().__init__(segment) - self._segment = cast(_InlineSegmentTO, segment) - - @property - def data(self) -> bytes: - return base64.b64decode(self._segment["data"]) - - def __repr__(self): - return f"InlineSegment(metadata={self.metadata})" - - class SpooledSegment(Segment): """ A subclass of Segment that handles spooled data segments, where data may be compressed and needs to be @@ -1292,8 +656,10 @@ def __init__( @property def data(self) -> bytes: - http_response = self._send_spooling_request(self.uri) - if not http_response.ok: + # No timeout on the data path: downloading a large segment from + # external storage may legitimately take a long time. + http_response = self._send_spooling_request(self.uri, timeout=None) + if http_response.is_error: self._request.raise_response_error(http_response) return http_response.content @@ -1313,25 +679,15 @@ def acknowledge(self) -> None: def acknowledge_request(): try: http_response = self._send_spooling_request(self.ack_uri, timeout=2) - if not http_response.ok: + if http_response.is_error: self._request.raise_response_error(http_response) except Exception as e: logger.error(f"Failed to acknowledge spooling request for segment {self}: {e}") # Start the acknowledgment in the executor thread executor.submit(acknowledge_request) - def _send_spooling_request(self, uri: str, **kwargs) -> requests.Response: - headers: Dict[str, str] = {} - # Forward user-supplied custom headers (e.g. auth gateway headers) only when the - # request targets the Trino coordinator, never to external storage (e.g. S3 presigned - # URLs) where such headers can break the request. The per-segment protocol headers - # returned by the coordinator always take precedence. - if self._coordinator_host is not None and urllib.parse.urlsplit(uri).hostname == self._coordinator_host: - headers.update(self._custom_headers) - for key, values in self.headers.items(): - if len(values) > 1: - raise ValueError(f"Header '{key}' contains multiple values: {values}") - headers[key] = values[0] + def _send_spooling_request(self, uri: str, **kwargs) -> Response: + headers = spooling_request_headers(uri, self.headers, self._custom_headers, self._coordinator_host) return self._request._get(uri, headers=headers, **kwargs) def __repr__(self): @@ -1340,36 +696,6 @@ def __repr__(self): ) -class DecodableSegment: - """ - Represents a collection of spooled segments of data, with an encoding format. - - Attributes: - encoding (str): The encoding format of the spooled data. - metadata (_SegmentMetadataTO): Metadata for all segments in the query - segment (Segment): The spooled segment data - """ - def __init__(self, encoding: str, metadata: _SegmentMetadataTO, segment: Segment) -> None: - self._encoding = encoding - self._metadata = metadata - self._segment = segment - - @property - def encoding(self): - return self._encoding - - @property - def segment(self): - return self._segment - - @property - def metadata(self): - return self._metadata - - def __repr__(self): - return (f"DecodableSegment(encoding={self._encoding}, metadata={self._metadata}, segment={self._segment})") - - class _RequestHeartbeat: """ Heartbeat loop for a trino request. Periodically sends HEAD requests to the request's next URI. @@ -1410,7 +736,7 @@ def _run(self) -> None: if response.status_code in (404, 405): logger.warning("The server does not support heartbeat calls") return - if not response.ok: + if response.is_error: failures += 1 else: failures = 0 @@ -1503,85 +829,3 @@ def decode(self, segment: Segment) -> List[List[Any]]: return self._decoder.decode(spooled_data.data, spooled_data.metadata) else: raise ValueError(f"Unsupported segment type: {type(segment)}") - - -class CompressedQueryDataDecoderFactory(): - def __init__(self, mapper: RowMapper) -> None: - self._mapper = mapper - - def create(self, encoding: str) -> QueryDataDecoder: - if encoding == "json+zstd": - if "zstd" in CODECS_UNAVAILABLE: - raise ValueError( - f"zstd is not installed so json+zstd encoding is not supported: {CODECS_UNAVAILABLE['zstd']}" - ) - return ZStdQueryDataDecoder(JsonQueryDataDecoder(self._mapper)) - elif encoding == "json+lz4": - if "lz4" in CODECS_UNAVAILABLE: - raise ValueError( - f"lz4 is not installed so json+lz4 encoding is not supported: {CODECS_UNAVAILABLE['lz4']}" - ) - return Lz4QueryDataDecoder(JsonQueryDataDecoder(self._mapper)) - elif encoding == "json": - return JsonQueryDataDecoder(self._mapper) - else: - raise ValueError(f"Unsupported encoding: {encoding}") - - -class QueryDataDecoder(abc.ABC): - @abstractmethod - def decode(self, data: bytes, metadata: _SegmentMetadataTO) -> List[List[Any]]: - pass - - -class JsonQueryDataDecoder(QueryDataDecoder): - def __init__(self, mapper: RowMapper) -> None: - self._mapper = mapper - - def decode(self, data: bytes, metadata: Dict[str, Any]) -> List[List[Any]]: - return self._mapper.map(json.loads(data.decode("utf8"))) - - -class CompressedQueryDataDecoder(QueryDataDecoder): - def __init__(self, delegate: QueryDataDecoder) -> None: - self._delegate = delegate - - @abstractmethod - def decompress(self, data: bytes, metadata: _SegmentMetadataTO) -> bytes: - pass - - def decode(self, data: bytes, metadata: _SegmentMetadataTO) -> List[List[Any]]: - if "uncompressedSize" not in metadata: - # Data not compressed - below threshold - return self._delegate.decode(data, metadata) - - # Data is compressed - expected_compressed_size = metadata["segmentSize"] - if not len(data) == expected_compressed_size: - raise RuntimeError(f"Expected to read {expected_compressed_size} bytes but got {len(data)}") - decompressed_data = self.decompress(data, metadata) - expected_uncompressed_size = metadata["uncompressedSize"] - if not len(decompressed_data) == expected_uncompressed_size: - raise RuntimeError( - "Decompressed size does not match expected segment size, " - f"expected {expected_uncompressed_size}, got {len(decompressed_data)}" - ) - return self._delegate.decode(decompressed_data, metadata) - - -class ZStdQueryDataDecoder(CompressedQueryDataDecoder): - def __init__(self, delegate: QueryDataDecoder) -> None: - super().__init__(delegate) - self._decompressor = None - - def decompress(self, data: bytes, metadata: _SegmentMetadataTO) -> bytes: - if self._decompressor is None: - self._decompressor = zstandard.ZstdDecompressor() - return self._decompressor.decompress(data) - - -class Lz4QueryDataDecoder(CompressedQueryDataDecoder): - def decompress(self, data: bytes, metadata: _SegmentMetadataTO) -> bytes: - expected_uncompressed_size = metadata["uncompressedSize"] - decoded_bytes = lz4.block.decompress(data, uncompressed_size=int(expected_uncompressed_size)) - return decoded_bytes diff --git a/trino/dbapi.py b/trino/dbapi.py index 35d576e3..d2118da4 100644 --- a/trino/dbapi.py +++ b/trino/dbapi.py @@ -133,6 +133,48 @@ def connect(*args, **kwargs): _USE_DEFAULT_ENCODING = object() +def _default_spooling_encoding(): + return [ + enc + for enc in trino.client.ENCODINGS + if (enc.split("+")[1] if "+" in enc else None) not in trino.client.CODECS_UNAVAILABLE + ] + + +def _resolve_scheme_and_port(parsed_host, port, http_scheme): + """Resolve the effective HTTP scheme and port; the scheme and port in the + host URL take precedence over the explicit arguments.""" + if parsed_host.scheme: + scheme = parsed_host.scheme + elif http_scheme: + scheme = http_scheme + elif port == constants.DEFAULT_TLS_PORT: + scheme = constants.HTTPS + else: + scheme = constants.HTTP + + default_port = constants.DEFAULT_TLS_PORT if scheme == constants.HTTPS else constants.DEFAULT_PORT + resolved_port = ( + parsed_host.port if parsed_host.port is not None + else port if port is not None + else default_port + ) + return scheme, resolved_port + + +def _require_tls_for_auth(auth, http_scheme, allow_insecure_auth): + if auth is not None and 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'. " + "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." + ) + + 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 @@ -172,11 +214,7 @@ def __init__( parsed_host = urlparse(host, allow_fragments=False) if encoding is _USE_DEFAULT_ENCODING: - encoding = [ - enc - for enc in trino.client.ENCODINGS - if (enc.split("+")[1] if "+" in enc else None) not in trino.client.CODECS_UNAVAILABLE - ] + encoding = _default_spooling_encoding() self.host = host if parsed_host.hostname is None else parsed_host.hostname + parsed_host.path self.user = user @@ -201,43 +239,15 @@ def __init__( ) # mypy cannot follow module import if http_session is None: - self._http_session = trino.client.TrinoRequest.http.Session() - self._http_session.verify = verify + self._http_session = trino.client.TrinoRequest.create_http_client( + verify=verify, timeout=request_timeout, auth=auth + ) else: 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'. " - "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.http_scheme, self.port = _resolve_scheme_and_port(parsed_host, port, http_scheme) + _require_tls_for_auth(auth, self.http_scheme, allow_insecure_auth) self.auth = auth self.extra_credential = extra_credential diff --git a/trino/exceptions.py b/trino/exceptions.py index 6e2e6503..ae00fbca 100644 --- a/trino/exceptions.py +++ b/trino/exceptions.py @@ -74,6 +74,15 @@ class TrinoConnectionError(OperationalError): pass +class SPNEGOExchangeError(TrinoAuthError): + """Raised when a Kerberos/GSSAPI (SPNEGO) token exchange with the server fails.""" + pass + + +# Backwards-compatible alias for the requests_kerberos-era exception name. +KerberosExchangeError = SPNEGOExchangeError + + class TrinoDataError(NotSupportedError): pass diff --git a/trino/transaction.py b/trino/transaction.py index b6b506d6..b11fdc81 100644 --- a/trino/transaction.py +++ b/trino/transaction.py @@ -65,7 +65,7 @@ def request(self) -> trino.client.TrinoRequest: def begin(self) -> None: response = self._request.post(START_TRANSACTION) - if not response.ok: + if response.is_error: raise trino.exceptions.DatabaseError( "failed to start transaction: {}".format(response.status_code) )