Add defer_connect config to allow eagerly verifying connection

tests/unit/sqlalchemy/test_dialect.py

@@ -252,7 +252,8 @@ def test_get_default_isolation_level(self):
         assert isolation_level == "AUTOCOMMIT"
 
     def test_isolation_level(self):
-        dbapi_conn = Connection(host="localhost")
+        # The test only verifies that isolation level is correctly set, no need to attempt actual connection
+        dbapi_conn = Connection(host="localhost", defer_connect=True)
 
         self.dialect.set_isolation_level(dbapi_conn, "SERIALIZABLE")
         assert dbapi_conn._isolation_level == IsolationLevel.SERIALIZABLE

tests/unit/test_dbapi.py

@@ -184,7 +184,8 @@ def test_token_retrieved_once_when_authentication_instance_is_shared(sample_post
         conn2.cursor().execute("SELECT 2")
         conn2.cursor().execute("SELECT 3")
 
-    assert len(_post_statement_requests()) == 7
+    assert len(_post_statement_requests()) == 9
+    # assert only a single token request was sent
     assert len(_get_token_requests(challenge_id)) == 1
 
 
@@ -275,37 +276,38 @@ def test_role_is_set_when_specified(mock_client):
 
 
 def test_hostname_parsing():
-    https_server_with_port = Connection("https://mytrinoserver.domain:9999")
+    # Since this test only verifies URL parsing there is no need to attempt actual connection
+    https_server_with_port = Connection("https://mytrinoserver.domain:9999", defer_connect=True)
     assert https_server_with_port.host == "mytrinoserver.domain"
     assert https_server_with_port.port == 9999
     assert https_server_with_port.http_scheme == constants.HTTPS
 
-    https_server_without_port = Connection("https://mytrinoserver.domain")
+    https_server_without_port = Connection("https://mytrinoserver.domain", defer_connect=True)
     assert https_server_without_port.host == "mytrinoserver.domain"
     assert https_server_without_port.port == 8080
     assert https_server_without_port.http_scheme == constants.HTTPS
 
-    http_server_with_port = Connection("http://mytrinoserver.domain:9999")
+    http_server_with_port = Connection("http://mytrinoserver.domain:9999", defer_connect=True)
     assert http_server_with_port.host == "mytrinoserver.domain"
     assert http_server_with_port.port == 9999
     assert http_server_with_port.http_scheme == constants.HTTP
 
-    http_server_without_port = Connection("http://mytrinoserver.domain")
+    http_server_without_port = Connection("http://mytrinoserver.domain", defer_connect=True)
     assert http_server_without_port.host == "mytrinoserver.domain"
     assert http_server_without_port.port == 8080
     assert http_server_without_port.http_scheme == constants.HTTP
 
-    http_server_with_path = Connection("http://mytrinoserver.domain/some_path")
+    http_server_with_path = Connection("http://mytrinoserver.domain/some_path", defer_connect=True)
     assert http_server_with_path.host == "mytrinoserver.domain/some_path"
     assert http_server_with_path.port == 8080
     assert http_server_with_path.http_scheme == constants.HTTP
 
-    only_hostname = Connection("mytrinoserver.domain")
+    only_hostname = Connection("mytrinoserver.domain", defer_connect=True)
     assert only_hostname.host == "mytrinoserver.domain"
     assert only_hostname.port == 8080
     assert only_hostname.http_scheme == constants.HTTP
 
-    only_hostname_with_path = Connection("mytrinoserver.domain/some_path")
+    only_hostname_with_path = Connection("mytrinoserver.domain/some_path", defer_connect=True)
     assert only_hostname_with_path.host == "mytrinoserver.domain/some_path"
     assert only_hostname_with_path.port == 8080
     assert only_hostname_with_path.http_scheme == constants.HTTP

trino/dbapi.py

@@ -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:
@@ -157,6 +159,7 @@ def __init__(
         legacy_prepared_statements=None,
         roles=None,
         timezone=None,
+        defer_connect=False,
     ):
         # Automatically assign http_schema, port based on hostname
         parsed_host = urlparse(host, allow_fragments=False)
@@ -201,6 +204,31 @@ 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.max_attempts,
+            self.request_timeout,
+            verify=self._http_session.verify,
+        )
+        try:
+            test_response = connection_test_request.post("<not-going-to-be-executed>")
+            response_content = test_response.content if test_response.content else ""
+            if not test_response.ok:
+                raise trino.exceptions.TrinoConnectionError(
+                    "error {}: {}".format(test_response.status_code, response_content))
+
+        except RequestException as e:
+            raise trino.exceptions.TrinoConnectionError("connection failed: {}".format(e))
+
     @property
     def isolation_level(self):
         return self._isolation_level