feat: allow custom headers to be sent through http requests

dbtmetabase/__init__.py

@@ -1,7 +1,7 @@
 import logging
 import functools
 from pathlib import Path
-from typing import Iterable, Optional, Callable, Any
+from typing import Iterable, Optional, Callable, Any, Tuple
 import os
 
 import click
@@ -86,7 +86,8 @@ class OptionAcceptableFromConfig(click.Option):
     more resilence to raising an error when the option exists in the users config.
 
     This also overrides default values for boolean CLI flags (e.g. --use_metabase_http/--use_metabase_https) in options when
-    no CLI flag is passed, but a value is provided in the config file (e.g. metabase_use_http: True)."""
+    no CLI flag is passed, but a value is provided in the config file (e.g. metabase_use_http: True).
+    """
 
     def process_value(self, ctx: click.Context, value: Any) -> Any:
         if value is not None:
@@ -115,10 +116,10 @@ def process_value(self, ctx: click.Context, value: Any) -> Any:
 
 class CommandController(click.Command):
     """This class inherets from click.Command and supplies custom help text renderer to
-    render our docstrings a little prettier as well as a hook in the invoke to load from a config file if it exists."""
+    render our docstrings a little prettier as well as a hook in the invoke to load from a config file if it exists.
+    """
 
     def invoke(self, ctx: click.Context):
-
         if CONFIG:
             for param, value in ctx.params.items():
                 if value is None and param in CONFIG:
@@ -281,6 +282,13 @@ def shared_opts(func: Callable) -> Callable:
         type=int,
         help="Synchronization timeout (in secs). If set, we will fail hard on synchronization failure; if not set, we will proceed after attempting sync regardless of success. Only valid if sync is enabled",
     )
+    @click.option(
+        "--http_extra_headers",
+        cls=OptionAcceptableFromConfig,
+        type=(str, str),
+        multiple=True,
+        help="Additional HTTP request header to be sent to Metabase.",
+    )
     @functools.wraps(func)
     def wrapper(*args, **kwargs):
         return func(*args, **kwargs)
@@ -554,6 +562,7 @@ def models(
     dbt_include_tags: bool = True,
     dbt_docs_url: Optional[str] = None,
     verbose: bool = False,
+    http_extra_headers: Optional[Tuple[Tuple[str, str]]] = None,
 ):
     """Exports model documentation and semantic types from dbt to Metabase.
 
@@ -578,6 +587,7 @@ def models(
         metabase_exclude_sources (bool, optional): Flag to skip exporting sources to Metabase. Defaults to False.
         dbt_include_tags (bool, optional): Flag to append tags to table descriptions in Metabase. Defaults to True.
         dbt_docs_url (Optional[str], optional): Pass in URL to dbt docs site. Appends dbt docs URL for each model to Metabase table description. Defaults to None.
+        http_extra_headers (Optional[str], optional): Additional HTTP request headers to be sent to Metabase. Defaults to None.
         verbose (bool, optional): Flag which signals verbose output. Defaults to False.
     """
 
@@ -614,6 +624,7 @@ def models(
         sync=metabase_sync,
         sync_timeout=metabase_sync_timeout,
         exclude_sources=metabase_exclude_sources,
+        http_extra_headers=http_extra_headers,
     )
 
     # Load client
@@ -678,6 +689,7 @@ def exposures(
     output_name: str = "metabase_exposures.yml",
     include_personal_collections: bool = False,
     collection_excludes: Optional[Iterable] = None,
+    http_extra_headers: Optional[Tuple[Tuple[str, str]]] = None,
     verbose: bool = False,
 ) -> None:
     """Extracts and imports exposures from Metabase to dbt.
@@ -703,6 +715,7 @@ def exposures(
         output_name (str): Output name for generated exposure yaml. Defaults to metabase_exposures.yml.
         include_personal_collections (bool, optional): Flag to include Personal Collections during exposure parsing. Defaults to False.
         collection_excludes (Iterable, optional): Collection names to exclude. Defaults to None.
+        http_extra_headers (Optional[str], optional): Additional HTTP request headers to be sent to Metabase. Defaults to None.
         verbose (bool, optional): Flag which signals verbose output. Defaults to False.
     """
 
@@ -735,6 +748,7 @@ def exposures(
         database=metabase_database,
         sync=metabase_sync,
         sync_timeout=metabase_sync_timeout,
+        http_extra_headers=http_extra_headers,
     )
 
     # Load client

dbtmetabase/metabase.py

@@ -127,6 +127,7 @@ def __init__(
         sync: Optional[bool] = True,
         sync_timeout: Optional[int] = None,
         exclude_sources: bool = False,
+        http_extra_headers: Optional[dict] = None,
     ):
         """Constructor.
 
@@ -142,12 +143,15 @@ def __init__(
             session_id {str} -- Metabase session ID. (default: {None})
             sync (bool, optional): Attempt to synchronize Metabase schema with local models. Defaults to True.
             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})
         """
         self.base_url = f"{'http' if use_http else 'https'}://{host}"
         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)
         session_header = session_id or self.get_session_id(user, password)
@@ -645,7 +649,6 @@ def increase_indent(self, flow=False, indentless=False):
         parsed_exposures = []
 
         for collection in self.collections:
-
             # Exclude collections by name
             if collection["name"] in collection_excludes:
                 continue
@@ -657,7 +660,6 @@ def increase_indent(self, flow=False, indentless=False):
             # Iter through collection
             logger().info(":sparkles: Exploring collection %s", collection["name"])
             for item in self.api("get", f"/api/collection/{collection['id']}/items"):
-
                 # Ensure collection item is of parsable type
                 exposure_type = item["model"]
                 exposure_id = item["id"]
@@ -678,7 +680,6 @@ def increase_indent(self, flow=False, indentless=False):
 
                 # Process exposure
                 if exposure_type == "card":
-
                     # Build header for card and extract models to self.models_exposed
                     header = "### Visualization: {}\n\n".format(
                         exposure.get("display", "Unknown").title()
@@ -689,7 +690,6 @@ def increase_indent(self, flow=False, indentless=False):
                     native_query = self.native_query
 
                 elif exposure_type == "dashboard":
-
                     # We expect this dict key in order to iter through questions
                     if "ordered_cards" not in exposure:
                         continue
@@ -824,7 +824,6 @@ def _extract_card_exposures(
 
             # Find models exposed through joins
             for query_join in query.get("query", {}).get("joins", []):
-
                 # Handle questions based on other question in virtual db
                 if str(query_join.get("source-table", "")).startswith("card__"):
                     self._extract_card_exposures(
@@ -853,7 +852,6 @@ def _extract_card_exposures(
 
             # Parse SQL for exposures through FROM or JOIN clauses
             for sql_ref in re.findall(self.exposure_parser, native_query):
-
                 # Grab just the table / model name
                 clean_exposure = sql_ref.split(".")[-1].strip('"').upper()
 

dbtmetabase/models/interface.py

@@ -32,6 +32,7 @@ def __init__(
         sync: bool = True,
         sync_timeout: Optional[int] = None,
         exclude_sources: bool = False,
+        http_extra_headers: Optional[dict] = None,
     ):
         """Constructor.
 
@@ -47,6 +48,7 @@ def __init__(
             sync (bool, optional): Attempt to synchronize Metabase schema with local models. Defaults to True.
             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.
         """
 
         # Metabase Client
@@ -59,6 +61,7 @@ def __init__(
         self.use_http = use_http
         self.verify = verify
         self.cert = cert
+        self.http_extra_headers = dict(http_extra_headers) if http_extra_headers else {}
         # Metabase Sync
         self.sync = sync
         self.sync_timeout = sync_timeout
@@ -101,6 +104,7 @@ def prepare_metabase_client(self, dbt_models: Optional[List[MetabaseModel]] = No
             use_http=self.use_http,
             verify=self.verify,
             cert=self.cert,
+            http_extra_headers=self.http_extra_headers,
             session_id=self.session_id,
             sync=self.sync,
             sync_timeout=self.sync_timeout,