Skip to content

Commit 05695e8

Browse files
Add defer_connect config to allow eagerly verifying connection
This commit adds a new connection parameter `defer_connect` which can be set to False to force creating a connection when `trino.dbapi.connect` is called. Any connection errors as a result of that get rewrapped into `trino.exceptions.TrinoConnectionError`. By default `defer_connect` is set to `True` so users can explicitly call `trino.dbapi.Connection.connect` to do the connection check. This doesn't end up actually executing a query on the server because after the initial POST request the nextUri in the response is not followed which leaves the query in QUEUED state. This is not documented in the Trino REST API but the server does behave like this today. The benefit is that we can very cheaply verify if the connection is valid without polluting the server's query history or adding queries to queue.
1 parent 9d417d6 commit 05695e8

File tree

1 file changed

+33
-0
lines changed

1 file changed

+33
-0
lines changed

trino/dbapi.py

+33
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,8 @@
2828
from typing import Any, Dict, List, NamedTuple, Optional # NOQA for mypy types
2929
from urllib.parse import urlparse
3030

31+
from requests.exceptions import RequestException
32+
3133
try:
3234
from zoneinfo import ZoneInfo
3335
except ModuleNotFoundError:
@@ -158,6 +160,7 @@ def __init__(
158160
legacy_prepared_statements=None,
159161
roles=None,
160162
timezone=None,
163+
defer_connect=True,
161164
):
162165
# Automatically assign http_schema, port based on hostname
163166
parsed_host = urlparse(host, allow_fragments=False)
@@ -203,6 +206,36 @@ def __init__(
203206
self.legacy_primitive_types = legacy_primitive_types
204207
self.legacy_prepared_statements = legacy_prepared_statements
205208

209+
if not defer_connect:
210+
self.connect()
211+
212+
def connect(self) -> None:
213+
connection_test_request = trino.client.TrinoRequest(
214+
self.host,
215+
self.port,
216+
self._client_session,
217+
self._http_session,
218+
self.http_scheme,
219+
self.auth,
220+
self.redirect_handler,
221+
self.max_attempts,
222+
self.request_timeout,
223+
verify=self._http_session.verify,
224+
)
225+
try:
226+
test_response = connection_test_request.post("<not-going-to-be-executed>")
227+
if not test_response.ok:
228+
raise trino.exceptions.TrinoConnectionError(
229+
"error {}{}".format(
230+
test_response.status_code,
231+
": {}".format(test_response.content)
232+
if test_response.content
233+
else "",
234+
)
235+
)
236+
except RequestException as e:
237+
raise trino.exceptions.TrinoConnectionError("connection failed: {}".format(e))
238+
206239
@property
207240
def isolation_level(self):
208241
return self._isolation_level

0 commit comments

Comments
 (0)