From dc98d01ad54af8003ef56241ccf56fb42195ab01 Mon Sep 17 00:00:00 2001 From: Dat Nguyen Date: Sat, 24 Feb 2024 17:23:15 +0700 Subject: [PATCH] Feat/83 python api (cont) (#94) * feat: python api adding (cont) * chore: test api * chore: add tests * chore: cq * docs: update readme [skip ci] --- README.md | 54 ++++++++++++++++++++++++++++++- dbterd/adapters/base.py | 11 +++++-- dbterd/api/__init__.py | 55 +++++++++++++++++++------------- dbterd/cli/params.py | 2 +- dbterd/default.py | 2 +- samples/dbtresto/erd.py | 20 ++++++++++++ tests/unit/adapters/test_base.py | 15 ++++++++- tests/unit/api/test_api.py | 36 +++++++++++++++++++++ tests/unit/test_api.py | 12 ------- tests/unit/test_default.py | 4 +-- 10 files changed, 169 insertions(+), 42 deletions(-) create mode 100644 samples/dbtresto/erd.py create mode 100644 tests/unit/api/test_api.py delete mode 100644 tests/unit/test_api.py diff --git a/README.md b/README.md index a4a7837..666fe15 100644 --- a/README.md +++ b/README.md @@ -21,7 +21,7 @@ dbterd --version ## Quick examine with existing samples
- Click me + Play with CLI ```bash # select all models in dbt_resto @@ -53,6 +53,58 @@ dbterd --version
+
+ Play with Python API (whole ERD) + + ```python + from dbterd.api import DbtErd + + erd = DbtErd().get_erd() + print("erd (dbml):", erd) + + erd = DbtErd(target="mermaid").get_erd() + print("erd (mermaid):", erd) + ``` + +
+ + +
+ Play with Python API (1 model's ERD) + + ```python + from dbterd.api import DbtErd + + dim_prize_erd = DbtErd(target="mermaid").get_model_erd( + node_unique_id="model.dbt_resto.dim_prize" + ) + print("erd of dim_date (mermaid):", dim_prize_erd) + ``` + + Here is the output: + + ```mermaid + erDiagram + "MODEL.DBT_RESTO.DIM_PRIZE" { + varchar prize_key + nvarchar prize_name + int prize_order + } + "MODEL.DBT_RESTO.FACT_RESULT" { + varchar fact_result_key + varchar box_key + varchar prize_key + date date_key + int no_of_won + float prize_value + float prize_paid + int is_prize_taken + } + "MODEL.DBT_RESTO.FACT_RESULT" }|--|| "MODEL.DBT_RESTO.DIM_PRIZE": prize_key + ``` + +
+ ## Quick DEMO Check [Quick Demo](https://dbterd.datnguyen.de/latest/nav/guide/targets/generate-dbml.html) out! And, following is the sample result using `dbdocs`: diff --git a/dbterd/adapters/base.py b/dbterd/adapters/base.py index 727b206..955fc89 100644 --- a/dbterd/adapters/base.py +++ b/dbterd/adapters/base.py @@ -1,5 +1,6 @@ import os from pathlib import Path +from types import NotImplementedType from typing import List, Tuple import click @@ -238,7 +239,10 @@ def __run_by_strategy( if not kwargs.get("api"): self.__save_result(path=kwargs.get("output"), data=result) - return result + if type(result) == NotImplementedType: + return result + + return result[1] def __run_metadata_by_strategy(self, **kwargs) -> Tuple[List[Table], List[Ref]]: """Metadata - Read artifacts and export the diagram file following the target""" @@ -250,4 +254,7 @@ def __run_metadata_by_strategy(self, **kwargs) -> Tuple[List[Table], List[Ref]]: if not kwargs.get("api"): self.__save_result(path=kwargs.get("output"), data=result) - return result + if type(result) == NotImplementedType: + return result + + return result[1] diff --git a/dbterd/api/__init__.py b/dbterd/api/__init__.py index c3307a6..922a21c 100644 --- a/dbterd/api/__init__.py +++ b/dbterd/api/__init__.py @@ -1,3 +1,5 @@ +import logging +from pathlib import Path from typing import List, Tuple from click import Command, Context @@ -5,6 +7,9 @@ from dbterd import default from dbterd.adapters.base import Executor from dbterd.adapters.meta import Ref, Table +from dbterd.helpers.log import logger + +logger.setLevel(logging.WARNING) # hide log class DbtErd: @@ -12,34 +17,38 @@ class DbtErd: dbt ERD official API functions. - Usage: + **Usage**: - - Get a whole ERD: - ```python - from dbterd.api import DbtErd - erd = DbtErd().get_erd() - ``` + ## Get a whole ERD - - Get a whole ERD given all models attached to `my_exposure_name`: - ```python - from dbterd.api import DbtErd - erd = DbtErd(select="exposure:my_exposure_name").get_erd() - ``` - See the - [Selection](https://dbterd.datnguyen.de/latest/nav/guide/cli-references.html#dbterd-run-select-s) - page for more details. + ```python + from dbterd.api import DbtErd + erd = DbtErd().get_erd() + ``` - - Get a model (named `model.jaffle_shop.my_model`)'s ERD: - ```python - from dbterd.api import DbtErd - erd = DbtErd().get_model_erd(s - node_fqn="model.jaffle_shop.my_model" - ) - ``` + ## Get a whole ERD given all models attached to `my_exposure_name` + + ```python + from dbterd.api import DbtErd + erd = DbtErd(select="exposure:my_exposure_name").get_erd() + ``` + See the + [Selection](https://dbterd.datnguyen.de/latest/nav/guide/cli-references.html#dbterd-run-select-s) + page for more details. + + ## Get a model (named `model.jaffle_shop.my_model`)'s ERD + + ```python + from dbterd.api import DbtErd + erd = DbtErd().get_model_erd(s + node_fqn="model.jaffle_shop.my_model" + ) + ``` """ def __init__(self, **kwargs) -> None: """Initialize the main Executor given similar input CLI parameters""" + self.params: dict = kwargs """ Mimic CLI params with overriding `api = True`.\n @@ -67,11 +76,13 @@ def __set_params_default_if_not_specified(self) -> None: self.params["resource_type"] = self.params.get( "resource_type", default.default_resource_types() ) - self.params["algo"] = self.params.get("algo", default.deafult_algo()) + self.params["algo"] = self.params.get("algo", default.default_algo()) self.params["entity_name_format"] = self.params.get( "entity_name_format", default.default_entity_name_format() ) self.params["omit_columns"] = self.params.get("omit_columns", False) + self.params["artifacts_dir"] = self.params.get("artifacts_dir", Path.cwd()) + self.params["target"] = self.params.get("target", default.default_target()) def get_erd(self) -> Tuple[List[Table], List[Ref]]: """Generate ERD code for a whole project diff --git a/dbterd/cli/params.py b/dbterd/cli/params.py index 18d13ee..15f58c4 100644 --- a/dbterd/cli/params.py +++ b/dbterd/cli/params.py @@ -43,7 +43,7 @@ def common_params(func): "--algo", "-a", help="Specified algorithm in the way to detect diagram connectors", - default=default.deafult_algo(), + default=default.default_algo(), show_default=True, type=click.STRING, ) diff --git a/dbterd/default.py b/dbterd/default.py index adacc0c..a20658d 100644 --- a/dbterd/default.py +++ b/dbterd/default.py @@ -14,7 +14,7 @@ def default_target() -> str: return "dbml" -def deafult_algo() -> str: +def default_algo() -> str: return "test_relationship" diff --git a/samples/dbtresto/erd.py b/samples/dbtresto/erd.py new file mode 100644 index 0000000..2394f6d --- /dev/null +++ b/samples/dbtresto/erd.py @@ -0,0 +1,20 @@ +from dbterd.api import DbtErd + +erd = DbtErd().get_erd() +print("erd (dbml):", erd) +erd = DbtErd(target="mermaid").get_erd() +print("erd (mermaid):", erd) + +print("===============") +print("===============") +fact_number_erd = DbtErd(target="mermaid").get_model_erd( + node_unique_id="model.dbt_resto.fact_number" +) +print("erd of fact_number (mermaid):", fact_number_erd) + +print("===============") +print("===============") +dim_prize_erd = DbtErd(target="mermaid").get_model_erd( + node_unique_id="model.dbt_resto.dim_prize" +) +print("erd of dim_date (mermaid):", dim_prize_erd) diff --git a/tests/unit/adapters/test_base.py b/tests/unit/adapters/test_base.py index 2e5bec3..e915312 100644 --- a/tests/unit/adapters/test_base.py +++ b/tests/unit/adapters/test_base.py @@ -1,4 +1,5 @@ from pathlib import Path +from types import NotImplementedType from unittest import mock import click @@ -27,6 +28,18 @@ def test___run_metadata_by_strategy(self, mock_query_erd_data, mock_save_result) mock_query_erd_data.assert_called_once() mock_save_result.assert_called_once() + @mock.patch("dbterd.adapters.base.DbtCloudMetadata.query_erd_data") + @mock.patch("dbterd.adapters.base.Executor._Executor__save_result") + def test___run_metadata_by_strategy_with_not_implemented_algo( + self, mock_query_erd_data, mock_save_result + ): + result = Executor( + ctx=click.Context(command=click.BaseCommand("dummy")) + )._Executor__run_metadata_by_strategy(target="dbml", algo="notfound") + assert type(result) == NotImplementedType + mock_query_erd_data.assert_called_once() + mock_save_result.assert_called_once() + @mock.patch("builtins.open") def test___save_result(self, mock_open): Executor( @@ -260,7 +273,7 @@ def test__run_by_strategy__for_api_simple( mock_parent.attach_mock( mock_set_single_node_selection, "mock_set_single_node_selection" ) - mock_dbml_run.return_value = dict(i="irr") + mock_dbml_run.return_value = ("irr", dict(i="irr")) mock_parent.attach_mock(mock_dbml_run, "mock_dbml_run") assert worker._Executor__run_by_strategy(node_unique_id="irr") == dict(i="irr") diff --git a/tests/unit/api/test_api.py b/tests/unit/api/test_api.py new file mode 100644 index 0000000..d6689fc --- /dev/null +++ b/tests/unit/api/test_api.py @@ -0,0 +1,36 @@ +from pathlib import Path +from unittest import mock + +from dbterd import default +from dbterd.api import DbtErd + + +class TestDbtErd: + @mock.patch("dbterd.adapters.base.Executor.run") + def test_get_erd(self, mock_executor_run): + mock_executor_run.return_value = "expected-result" + assert DbtErd().get_erd() == "expected-result" + + @mock.patch("dbterd.adapters.base.Executor.run") + def test_get_model_erd(self, mock_executor_run): + mock_executor_run.return_value = "expected-result" + assert DbtErd().get_model_erd(node_unique_id="any") == "expected-result" + + def test_init_default(self): + actual = DbtErd() + actual_dict = dict(vars(actual)) + del actual_dict["executor"] + assert actual_dict == dict( + params=dict( + api=True, + select=[], + exclude=[], + resource_type=default.default_resource_types(), + algo=default.default_algo(), + entity_name_format=default.default_entity_name_format(), + omit_columns=False, + artifacts_dir=Path.cwd(), + target=default.default_target(), + ) + ) + assert actual.executor.ctx.command.name == "run" diff --git a/tests/unit/test_api.py b/tests/unit/test_api.py deleted file mode 100644 index 94b54f2..0000000 --- a/tests/unit/test_api.py +++ /dev/null @@ -1,12 +0,0 @@ -class TestDbtErd: - def test__init__(self): - pass # TODO - - def test__set_params_default_if_not_specified(self): - pass # TODO - - def test_get_erd(self): - pass # TODO - - def test_get_model_erd(self): - pass # TODO diff --git a/tests/unit/test_default.py b/tests/unit/test_default.py index c1f333e..7f2887b 100644 --- a/tests/unit/test_default.py +++ b/tests/unit/test_default.py @@ -17,5 +17,5 @@ def test_default_target(self, target): assert default.default_target() == target @pytest.mark.parametrize("algo", [("test_relationship")]) - def test_deafult_algo(self, algo): - assert default.deafult_algo() == algo + def test_default_algo(self, algo): + assert default.default_algo() == algo