Allow custom http adapter for certificate handling (#188)

gouline · Dec 5, 2023 · 848c81d · 848c81d
1 parent 88ddbde
commit 848c81d
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 9 deletions.
diff --git 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
@@ -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)