feat: BI-6553 Add Trino timeouts (#1221)

Author: khamitovdr

lib/dl_connector_trino/dl_connector_trino/core/adapters.py

@@ -69,13 +69,19 @@
 
 
 class CustomHTTPAdapter(HTTPAdapter):
-    def __init__(self, ssl_ca: str, *args: Any, **kwargs: Any) -> None:
+    """
+    This custom adapter is here to create an SSL context with a custom CA certificate provided as a string instead of a file path.
+    """
+
+    def __init__(self, ssl_ca: str | None = None, *args: Any, **kwargs: Any) -> None:
         self.ssl_ca = ssl_ca
         super().__init__(*args, **kwargs)
 
     def init_poolmanager(self, connections: int, maxsize: int, block: bool = False, **pool_kwargs: Any) -> None:
-        # Use a secure context with the provided SSL CA
-        context = ssl.create_default_context(cadata=self.ssl_ca)
+        if self.ssl_ca is None:
+            context = ssl.create_default_context(capath=get_root_certificates_path())
+        else:
+            context = ssl.create_default_context(cadata=self.ssl_ca)
         super().init_poolmanager(connections, maxsize, block, ssl_context=context, **pool_kwargs)
 
 
@@ -122,10 +128,6 @@ class TrinoDefaultAdapter(BaseClassicAdapter[TrinoConnTargetDTO]):
     EXTRA_EXC_CLS = (sa_exc.DBAPIError,)
 
     def get_conn_line(self, db_name: str | None = None, params: dict[str, Any] | None = None) -> str:
-        # We do not expect to transfer any additional parameters when creating the engine.
-        # This check is needed to track if it still passed.
-        assert params is None
-
         params = params or {}
         return trino_url(
             host=self._target_dto.host,
@@ -136,11 +138,24 @@ def get_conn_line(self, db_name: str | None = None, params: dict[str, Any] | Non
             **params,
         )
 
+    def _get_http_session(self) -> requests.Session:
+        session = requests.Session()
+        adapter = CustomHTTPAdapter(ssl_ca=self._target_dto.ssl_ca)
+        session.mount("http://", adapter)
+        session.mount("https://", adapter)
+        return session
+
     def get_connect_args(self) -> dict[str, Any]:
-        args: dict[str, Any] = {
-            **super().get_connect_args(),
-            "legacy_primitive_types": True,
-        }
+        timeout = (
+            self._target_dto.connect_timeout,
+            self._target_dto.total_timeout,
+        )
+        args: dict[str, Any] = super().get_connect_args() | dict(
+            http_scheme="https" if self._target_dto.ssl_enable else "http",
+            http_session=self._get_http_session(),
+            legacy_primitive_types=True,
+            request_timeout=timeout,
+        )
         if self._target_dto.auth_type is TrinoAuthType.none:
             pass
         elif self._target_dto.auth_type is TrinoAuthType.password:
@@ -150,18 +165,6 @@ def get_connect_args(self) -> dict[str, Any]:
         else:
             raise NotImplementedError(f"{self._target_dto.auth_type.name} authentication is not supported yet")
 
-        if not self._target_dto.ssl_enable:
-            args["http_scheme"] = "http"
-            return args
-
-        args["http_scheme"] = "https"
-        if self._target_dto.ssl_ca:
-            session = requests.Session()
-            session.mount("https://", CustomHTTPAdapter(self._target_dto.ssl_ca))
-            args["http_session"] = session
-        else:
-            args["verify"] = get_root_certificates_path()
-
         return args
 
     def execute_by_steps(self, db_adapter_query: DBAdapterQuery) -> Generator[ExecutionStep, None, None]:

lib/dl_connector_trino/dl_connector_trino/core/conn_options.py

@@ -0,0 +1,9 @@
+import attr
+
+from dl_core.connection_models.conn_options import ConnectOptions
+
+
+@attr.s(frozen=True, hash=True)
+class TrinoConnectOptions(ConnectOptions):
+    connect_timeout: int | None = attr.ib(default=None)
+    total_timeout: int | None = attr.ib(default=None)

lib/dl_connector_trino/dl_connector_trino/core/connection_executors.py

@@ -3,6 +3,7 @@
 from dl_core.connection_executors.async_sa_executors import DefaultSqlAlchemyConnExecutor
 
 from dl_connector_trino.core.adapters import TrinoDefaultAdapter
+from dl_connector_trino.core.conn_options import TrinoConnectOptions
 from dl_connector_trino.core.dto import TrinoConnDTO
 from dl_connector_trino.core.target_dto import TrinoConnTargetDTO
 
@@ -11,6 +12,7 @@
 class TrinoConnExecutor(DefaultSqlAlchemyConnExecutor[TrinoDefaultAdapter]):
     TARGET_ADAPTER_CLS = TrinoDefaultAdapter
     _conn_dto: TrinoConnDTO = attr.ib()
+    _conn_options: TrinoConnectOptions = attr.ib()
 
     async def _make_target_conn_dto_pool(self) -> list[TrinoConnTargetDTO]:
         return [
@@ -26,5 +28,7 @@ async def _make_target_conn_dto_pool(self) -> list[TrinoConnTargetDTO]:
                 jwt=self._conn_dto.jwt,
                 ssl_enable=self._conn_dto.ssl_enable,
                 ssl_ca=self._conn_dto.ssl_ca,
+                connect_timeout=self._conn_options.connect_timeout,
+                total_timeout=self._conn_options.total_timeout,
             )
         ]

lib/dl_connector_trino/dl_connector_trino/core/target_dto.py

@@ -1,5 +1,3 @@
-from typing import Optional
-
 import attr
 from typing_extensions import Self
 
@@ -9,18 +7,20 @@
 from dl_connector_trino.core.constants import TrinoAuthType
 
 
-@attr.s(frozen=True)
+@attr.s(frozen=True, kw_only=True)
 class TrinoConnTargetDTO(ConnTargetDTO):
     host: str = attr.ib()
     port: int = attr.ib()
     username: str = attr.ib()
-    auth_type: TrinoAuthType = attr.ib(kw_only=True)
-    password: Optional[str] = attr.ib(repr=False, kw_only=True, default=None)
-    jwt: Optional[str] = attr.ib(repr=False, kw_only=True, default=None)
-    ssl_enable: bool = attr.ib(kw_only=True, default=False)
-    ssl_ca: Optional[str] = attr.ib(kw_only=True, default=None)
-
-    def get_effective_host(self) -> Optional[str]:
+    auth_type: TrinoAuthType = attr.ib()
+    password: str | None = attr.ib(repr=False, default=None)
+    jwt: str | None = attr.ib(repr=False, default=None)
+    ssl_enable: bool = attr.ib(default=False)
+    ssl_ca: str | None = attr.ib(default=None)
+    connect_timeout: int | None = attr.ib(default=None)
+    total_timeout: int | None = attr.ib(default=None)
+
+    def get_effective_host(self) -> str | None:
         return self.host
 
     @classmethod

lib/dl_connector_trino/dl_connector_trino/core/us_connection.py

@@ -21,6 +21,7 @@
 from dl_i18n.localizer_base import Localizer
 
 from dl_connector_trino.api.i18n.localizer import Translatable
+from dl_connector_trino.core.conn_options import TrinoConnectOptions
 from dl_connector_trino.core.constants import (
     CONNECTION_TYPE_TRINO,
     SOURCE_TYPE_TRINO_SUBSELECT,
@@ -81,6 +82,14 @@ class DataModel(ConnectionSQL.DataModel):
         jwt: str | None = attr.ib(repr=secrepr, default=None)
         listing_sources: ListingSources = attr.ib()
 
+    def get_conn_options(self) -> TrinoConnectOptions:
+        base = super().get_conn_options()
+        return base.to_subclass(
+            TrinoConnectOptions,
+            connect_timeout=1,
+            total_timeout=80,
+        )
+
     def get_data_source_template_templates(self, localizer: Localizer) -> list[DataSourceTemplate]:
         result: list[DataSourceTemplate] = []