From d0823f6fee651f08045c73ff6498f4454a78d58f Mon Sep 17 00:00:00 2001 From: Paolo Quadri Date: Wed, 22 Nov 2023 12:13:22 +0100 Subject: [PATCH] feat: configurable metabase api timeotus --- dbtmetabase/__init__.py | 43 +++++++++++++++++++++++++++--- dbtmetabase/metabase.py | 10 ++++--- dbtmetabase/models/interface.py | 15 ++++++----- tests/test_click_custom_objects.py | 11 ++++++++ 4 files changed, 66 insertions(+), 13 deletions(-) create mode 100644 tests/test_click_custom_objects.py diff --git a/dbtmetabase/__init__.py b/dbtmetabase/__init__.py index c9fc959..33a72ba 100644 --- a/dbtmetabase/__init__.py +++ b/dbtmetabase/__init__.py @@ -1,14 +1,14 @@ -import logging import functools -from pathlib import Path -from typing import Iterable, Optional, Callable, Any, Dict +import logging import os +from pathlib import Path +from typing import Any, Callable, Dict, Iterable, Optional import click import yaml from .logger import logging as package_logger -from .models.interface import MetabaseInterface, DbtInterface +from .models.interface import DbtInterface, MetabaseInterface from .utils import get_version, load_config __all__ = ["MetabaseInterface", "DbtInterface"] @@ -24,6 +24,7 @@ "MB_HOST", "MB_DATABASE", "MB_SESSION_TOKEN", + "MB_TIMEOUT", ] @@ -114,6 +115,27 @@ def process_value(self, ctx: click.Context, value: Any) -> Any: return value +class IntOrNoneParamType(click.ParamType): + name = "integer_or_none" + + def convert( + self, value: Any, param: Optional[click.Parameter], ctx: Optional[click.Context] + ) -> Optional[int]: + if isinstance(value, int): + return value + + if value.lower() == "none": + return None + + try: + return int(value) + except ValueError: + self.fail(f"{value!r} is not a valid integer", param, ctx) + + +INT_OR_NONE = IntOrNoneParamType() + + 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. @@ -289,6 +311,13 @@ def shared_opts(func: Callable) -> Callable: multiple=True, help="Additional HTTP request header to be sent to Metabase.", ) + @click.option( + "--metabase_requests_timeout", + cls=OptionAcceptableFromConfig, + type=INT_OR_NONE, + default=15, + help="Set the value for single requests timeout", + ) @functools.wraps(func) def wrapper(*args, **kwargs): return func(*args, **kwargs) @@ -559,6 +588,7 @@ def models( metabase_sync: bool = True, metabase_sync_timeout: Optional[int] = None, metabase_exclude_sources: bool = False, + metabase_requests_timeout: Optional[int] = 15, dbt_include_tags: bool = True, dbt_docs_url: Optional[str] = None, verbose: bool = False, @@ -585,6 +615,7 @@ def models( metabase_sync (bool, optional): Attempt to synchronize Metabase schema with local models. Defaults to True. metabase_sync_timeout (Optional[int], optional): 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. Defaults to None. metabase_exclude_sources (bool, optional): Flag to skip exporting sources to Metabase. Defaults to False. + metabase_requests_timeout (int, optional): Set the timeout for the single Metabase requests. 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. @@ -625,6 +656,7 @@ def models( sync_timeout=metabase_sync_timeout, exclude_sources=metabase_exclude_sources, http_extra_headers=http_extra_headers, + requests_timeout=metabase_requests_timeout, ) # Load client @@ -685,6 +717,7 @@ def exposures( metabase_verify: Optional[str] = None, metabase_sync: bool = True, metabase_sync_timeout: Optional[int] = None, + metabase_requests_timeout: Optional[int] = 15, output_path: str = ".", output_name: str = "metabase_exposures.yml", include_personal_collections: bool = False, @@ -711,6 +744,7 @@ def exposures( metabase_verify (Optional[str], optional): Path to custom certificate bundle to be used by Metabase client. Defaults to None. metabase_sync (bool, optional): Attempt to synchronize Metabase schema with local models. Defaults to True. metabase_sync_timeout (Optional[int], optional): 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. Defaults to None. + metabase_requests_timeout (int, optional): Set the timeout for the single Metabase requests. output_path (str): Output path for generated exposure yaml. Defaults to "." local dir. 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. @@ -749,6 +783,7 @@ def exposures( sync=metabase_sync, sync_timeout=metabase_sync_timeout, http_extra_headers=http_extra_headers, + requests_timeout=metabase_requests_timeout, ) # Load client diff --git a/dbtmetabase/metabase.py b/dbtmetabase/metabase.py index e410075..7aa643f 100644 --- a/dbtmetabase/metabase.py +++ b/dbtmetabase/metabase.py @@ -128,6 +128,7 @@ def __init__( sync_timeout: Optional[int] = None, exclude_sources: bool = False, http_extra_headers: Optional[dict] = None, + requests_timeout: Optional[int] = 15, ): """Constructor. @@ -165,7 +166,7 @@ def __init__( self.table_map: MutableMapping = {} self.models_exposed: List = [] self.native_query: str = "" - + self.requests_timeout = requests_timeout # This regex is looking for from and join clauses, and extracting the table part. # It won't recognize some valid sql table references, such as `from "table with spaces"`. self.exposure_parser = re.compile(r"[FfJj][RrOo][OoIi][MmNn]\s+([\w.\"]+)") @@ -723,7 +724,10 @@ def increase_indent(self, flow=False, indentless=False): creator = self.api("get", f"/api/user/{exposure['creator_id']}") except requests.exceptions.HTTPError as error: creator = {} - if error.response.status_code != 404: + if ( + error.response is not None + and error.response.status_code != 404 + ): raise creator_email = creator.get("email") @@ -982,7 +986,7 @@ def api( response = self.session.request( method, f"{self.base_url}{path}", - timeout=15, + timeout=self.requests_timeout, **kwargs, ) diff --git a/dbtmetabase/models/interface.py b/dbtmetabase/models/interface.py index 436ac73..b885f3a 100644 --- a/dbtmetabase/models/interface.py +++ b/dbtmetabase/models/interface.py @@ -1,17 +1,17 @@ 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 class MetabaseInterface: @@ -33,6 +33,7 @@ def __init__( sync_timeout: Optional[int] = None, exclude_sources: bool = False, http_extra_headers: Optional[dict] = None, + requests_timeout: Optional[int] = 15, ): """Constructor. @@ -62,6 +63,7 @@ def __init__( self.verify = verify self.cert = cert self.http_extra_headers = dict(http_extra_headers) if http_extra_headers else {} + self.requests_timeout = requests_timeout # Metabase Sync self.sync = sync self.sync_timeout = sync_timeout @@ -109,6 +111,7 @@ def prepare_metabase_client(self, dbt_models: Optional[List[MetabaseModel]] = No sync=self.sync, sync_timeout=self.sync_timeout, exclude_sources=self.exclude_sources, + requests_timeout=self.requests_timeout, ) # Sync and attempt schema alignment prior to execution; if timeout is not explicitly set, proceed regardless of success diff --git a/tests/test_click_custom_objects.py b/tests/test_click_custom_objects.py new file mode 100644 index 0000000..fd7e898 --- /dev/null +++ b/tests/test_click_custom_objects.py @@ -0,0 +1,11 @@ +import unittest + +from dbtmetabase import INT_OR_NONE + + +class TestDbtFolderReader(unittest.TestCase): + def test_none_value(self): + assert INT_OR_NONE("None", None, None) is None + + def test_integer_value(self): + assert INT_OR_NONE("15", None, None) == 15