From 713dbb9c157d7a63891f5cd9d2f3c11ba5b0da2b Mon Sep 17 00:00:00 2001 From: "TANGUY Antoine (SIB)" Date: Wed, 4 Oct 2023 10:30:29 +0200 Subject: [PATCH] feat: allow custom headers to be sent through http requests --- dbtmetabase/__init__.py | 15 ++++++++++++++- dbtmetabase/metabase.py | 4 ++++ dbtmetabase/models/interface.py | 4 ++++ 3 files changed, 22 insertions(+), 1 deletion(-) diff --git a/dbtmetabase/__init__.py b/dbtmetabase/__init__.py index 254bd21..1386e64 100644 --- a/dbtmetabase/__init__.py +++ b/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 @@ -281,6 +281,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 +561,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 +586,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 +623,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 +688,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 +714,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 +747,7 @@ def exposures( database=metabase_database, sync=metabase_sync, sync_timeout=metabase_sync_timeout, + http_extra_headers=http_extra_headers ) # Load client diff --git a/dbtmetabase/metabase.py b/dbtmetabase/metabase.py index 290766f..790e324 100644 --- a/dbtmetabase/metabase.py +++ b/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) diff --git a/dbtmetabase/models/interface.py b/dbtmetabase/models/interface.py index 2342683..436ac73 100644 --- a/dbtmetabase/models/interface.py +++ b/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,