diff --git a/src/h2o_discovery/__init__.py b/src/h2o_discovery/__init__.py index 4978633..f8bbcd6 100644 --- a/src/h2o_discovery/__init__.py +++ b/src/h2o_discovery/__init__.py @@ -1,5 +1,7 @@ import dataclasses +import datetime import os +import ssl from typing import Optional from typing import Tuple from typing import Union @@ -19,6 +21,9 @@ def discover( environment: Optional[str] = None, discovery_address: Optional[str] = None, config_path: Optional[Union[str, bytes, os.PathLike]] = None, + *, + http_timeout: Optional[datetime.timedelta] = None, + ssl_context: Optional[ssl.SSLContext] = None, ) -> Discovery: """Obtains and returns a Discovery object from the discovery service. @@ -39,6 +44,11 @@ def discover( environment: The H2O Cloud environment URL to use (e.g. https://cloud.h2o.ai). discovery_address: The address of the discovery service. config_path: The path to the H2O CLI configuration file. + http_timeout: The timeout for HTTP requests. Value applies to all of the + timeouts (connect, read, write). + ssl_context: The SSL context to use for HTTPS requests. + If not specified default SSL context is used. + Raises: exceptions.DiscoveryLookupError: If the URI cannot be determined. @@ -47,7 +57,9 @@ def discover( """ uri, cfg = _lookup_and_load(environment, discovery_address, config_path) - discovery = load.load_discovery(client.Client(uri)) + discovery = load.load_discovery( + client.Client(uri=uri, timeout=http_timeout, ssl_context=ssl_context) + ) credentials = load.load_credentials( clients=discovery.clients, config_tokens=cfg.tokens ) @@ -59,6 +71,9 @@ async def discover_async( environment: Optional[str] = None, discovery_address: Optional[str] = None, config_path: Optional[Union[str, bytes, os.PathLike]] = None, + *, + http_timeout: Optional[datetime.timedelta] = None, + ssl_context: Optional[ssl.SSLContext] = None, ) -> Discovery: """Obtains and returns a Discovery object from the discovery service. @@ -79,6 +94,10 @@ async def discover_async( environment: The H2O Cloud environment URL to use (e.g. https://cloud.h2o.ai). discovery_address: The address of the discovery service. config_path: The path to the H2O CLI configuration file. + http_timeout: The timeout for HTTP requests. Value applies to all of the + timeouts (connect, read, write). + ssl_context: The SSL context to use for HTTPS requests. + If not specified default SSL context is used. Raises: exceptions.DiscoveryLookupError: If the URI cannot be determined. @@ -88,7 +107,9 @@ async def discover_async( uri, cfg = _lookup_and_load(environment, discovery_address, config_path) - discovery = await load.load_discovery_async(client.AsyncClient(uri)) + discovery = await load.load_discovery_async( + client.AsyncClient(uri=uri, timeout=http_timeout, ssl_context=ssl_context) + ) credentials = load.load_credentials( clients=discovery.clients, config_tokens=cfg.tokens ) diff --git a/src/h2o_discovery/_internal/client.py b/src/h2o_discovery/_internal/client.py index de65091..3de99be 100644 --- a/src/h2o_discovery/_internal/client.py +++ b/src/h2o_discovery/_internal/client.py @@ -2,7 +2,6 @@ import ssl from typing import List from typing import Optional -from typing import Union import httpx @@ -23,16 +22,12 @@ def __init__( ssl_context: Optional[ssl.SSLContext] = None, ): self._uri = uri - self._ssl_context = ssl_context + self._verify = ssl_context or ssl.create_default_context() self._timeout = 5.0 if timeout is not None: self._timeout = timeout.total_seconds() - self._verify: Union[bool, ssl.SSLContext] = True - if self._ssl_context is not None: - self._verify = self._ssl_context - class Client(_BaseClient): """Synchronous Implementation of the Discovery Service API. @@ -40,9 +35,6 @@ class Client(_BaseClient): Listing methods do pagination and always return all of the available objects. """ - def __init__(self, uri: str): - self._uri = uri - def get_environment(self) -> model.Environment: """Returns the information about the environment.""" with httpx.Client(base_url=self._uri) as client: