Skip to content

Commit

Permalink
Add defer_connect config to allow eagerly verifying connection
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
Shaheer-rossoneri14 committed Aug 6, 2023
1 parent 6eaef4b commit 676c42a
Showing 1 changed file with 33 additions and 0 deletions.
33 changes: 33 additions & 0 deletions trino/dbapi.py
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,8 @@
from typing import Any, Dict, List, NamedTuple, Optional # NOQA for mypy types
from urllib.parse import urlparse

from requests.exceptions import RequestException

try:
from zoneinfo import ZoneInfo
except ModuleNotFoundError:
Expand Down Expand Up @@ -157,6 +159,7 @@ def __init__(
legacy_prepared_statements=None,
roles=None,
timezone=None,
defer_connect=True,
):
# Automatically assign http_schema, port based on hostname
parsed_host = urlparse(host, allow_fragments=False)
Expand Down Expand Up @@ -201,6 +204,36 @@ def __init__(
self.legacy_primitive_types = legacy_primitive_types
self.legacy_prepared_statements = legacy_prepared_statements

if not defer_connect:
self.connect()

def connect(self) -> None:
connection_test_request = trino.client.TrinoRequest(
self.host,
self.port,
self._client_session,
self._http_session,
self.http_scheme,
self.auth,
self.redirect_handler,
self.max_attempts,
self.request_timeout,
verify=self._http_session.verify,
)
try:
test_response = connection_test_request.post("<not-going-to-be-executed>")
if not test_response.ok:
raise trino.exceptions.TrinoConnectionError(
"error {}{}".format(
test_response.status_code,
": {}".format(test_response.content)
if test_response.content
else "",
)
)
except RequestException as e:
raise trino.exceptions.TrinoConnectionError("connection failed: {}".format(e))

@property
def isolation_level(self):
return self._isolation_level
Expand Down

0 comments on commit 676c42a

Please sign in to comment.