diff --git a/dbtmetabase/metabase.py b/dbtmetabase/metabase.py index b688b98..8ca7ee6 100644 --- a/dbtmetabase/metabase.py +++ b/dbtmetabase/metabase.py @@ -129,6 +129,7 @@ def __init__( exclude_sources: bool = False, http_extra_headers: Optional[dict] = None, http_timeout: int = 15, + http_adapter: Optional[HTTPAdapter] = None, ): """Constructor. @@ -146,16 +147,22 @@ def __init__( sync_timeout (Optional[int], optional): Synchronization timeout (in secs). Defaults to None. http_extra_headers {dict} -- HTTP headers to be used by the Metabase client. (default: {None}) exclude_sources {bool} -- Exclude exporting sources. (default: {False}) + http_adapter: (Optional[HTTPAdapter], optional) Provide custom HTTP adapter implementation for requests to use. Defaults to None. """ + self.base_url = f"{'http' if use_http else 'https'}://{host}" self.http_timeout = http_timeout self.session = requests.Session() self.session.verify = verify self.session.cert = cert + if http_extra_headers is not None: self.session.headers.update(http_extra_headers) - adaptor = HTTPAdapter(max_retries=Retry(total=3, backoff_factor=0.5)) - self.session.mount(self.base_url, adaptor) + + if not http_adapter: + http_adapter = HTTPAdapter(max_retries=Retry(total=3, backoff_factor=0.5)) + + self.session.mount(self.base_url, http_adapter) session_header = session_id or self.get_session_id(user, password) self.session.headers["X-Metabase-Session"] = session_header diff --git a/dbtmetabase/models/interface.py b/dbtmetabase/models/interface.py index cdeb37e..658b83d 100644 --- a/dbtmetabase/models/interface.py +++ b/dbtmetabase/models/interface.py @@ -1,17 +1,18 @@ import logging from os.path import expandvars -from typing import Optional, Union, List, Tuple, MutableMapping, Iterable +from typing import Iterable, List, MutableMapping, Optional, Tuple, Union -from .metabase import MetabaseModel +from ..metabase import MetabaseClient +from ..parsers.dbt import DbtReader +from ..parsers.dbt_folder import DbtFolderReader +from ..parsers.dbt_manifest import DbtManifestReader from .exceptions import ( NoDbtPathSupplied, NoDbtSchemaSupplied, NoMetabaseCredentialsSupplied, ) -from ..parsers.dbt import DbtReader -from ..parsers.dbt_folder import DbtFolderReader -from ..parsers.dbt_manifest import DbtManifestReader -from ..metabase import MetabaseClient +from .metabase import MetabaseModel +from requests.adapters import HTTPAdapter class MetabaseInterface: @@ -34,6 +35,7 @@ def __init__( exclude_sources: bool = False, http_extra_headers: Optional[dict] = None, http_timeout: int = 15, + http_adapter: Optional[HTTPAdapter] = None, ): """Constructor. @@ -50,6 +52,7 @@ def __init__( sync_timeout (Optional[int], optional): Synchronization timeout (in secs). Defaults to None. exclude_sources (bool, optional): Exclude exporting sources. Defaults to False. http_extra_headers (Optional[dict], optional): HTTP headers to be used by the Metabase client. Defaults to None. + http_adapter: (Optional[HTTPAdapter], optional) Provide custom HTTP adapter implementation for requests to use. Defaults to None. """ # Metabase Client @@ -64,6 +67,7 @@ def __init__( self.cert = cert self.http_extra_headers = dict(http_extra_headers) if http_extra_headers else {} self.http_timeout = http_timeout + self.http_adapter = http_adapter # Metabase Sync self.sync = sync self.sync_timeout = sync_timeout @@ -112,8 +116,8 @@ def prepare_metabase_client(self, dbt_models: Optional[List[MetabaseModel]] = No sync_timeout=self.sync_timeout, exclude_sources=self.exclude_sources, http_timeout=self.http_timeout, + http_adapter=self.http_adapter, ) - # Sync and attempt schema alignment prior to execution; if timeout is not explicitly set, proceed regardless of success if self.sync: self._client.sync_and_wait(self.database, dbt_models)