From 848c81dc8c5b0a66a8a364705fe4ec3fb376a4c0 Mon Sep 17 00:00:00 2001
From: Alexis Flipo <63051878+alexisflipo@users.noreply.github.com>
Date: Tue, 5 Dec 2023 12:05:42 +0100
Subject: [PATCH] Allow custom http adapter for certificate handling (#188)

---
 dbtmetabase/metabase.py         | 11 +++++++++--
 dbtmetabase/models/interface.py | 18 +++++++++++-------
 2 files changed, 20 insertions(+), 9 deletions(-)

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)