Skip to content

Commit

Permalink
Feat/83 python api (cont) (#94)
Browse files Browse the repository at this point in the history
* feat: python api adding (cont)

* chore: test api

* chore: add tests

* chore: cq

* docs: update readme [skip ci]
  • Loading branch information
datnguye committed Feb 24, 2024
1 parent 726c061 commit dc98d01
Show file tree
Hide file tree
Showing 10 changed files with 169 additions and 42 deletions.
54 changes: 53 additions & 1 deletion README.md
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@ dbterd --version
## Quick examine with existing samples

<details>
<summary>Click me</summary>
<summary>Play with CLI</summary>

```bash
# select all models in dbt_resto
Expand Down Expand Up @@ -53,6 +53,58 @@ dbterd --version

</details>

<details>
<summary>Play with Python API (whole ERD)</summary>

```python
from dbterd.api import DbtErd

erd = DbtErd().get_erd()
print("erd (dbml):", erd)

erd = DbtErd(target="mermaid").get_erd()
print("erd (mermaid):", erd)
```

</details>


<details>
<summary>Play with Python API (1 model's ERD)</summary>

```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
```

</details>

## 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`:
Expand Down
11 changes: 9 additions & 2 deletions dbterd/adapters/base.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import os
from pathlib import Path
from types import NotImplementedType
from typing import List, Tuple

import click
Expand Down Expand Up @@ -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"""
Expand All @@ -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]
55 changes: 33 additions & 22 deletions dbterd/api/__init__.py
Original file line number Diff line number Diff line change
@@ -1,45 +1,54 @@
import logging
from pathlib import Path
from typing import List, Tuple

from click import Command, Context

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:
"""
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
Expand Down Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion dbterd/cli/params.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)
Expand Down
2 changes: 1 addition & 1 deletion dbterd/default.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@ def default_target() -> str:
return "dbml"


def deafult_algo() -> str:
def default_algo() -> str:
return "test_relationship"


Expand Down
20 changes: 20 additions & 0 deletions samples/dbtresto/erd.py
Original file line number Diff line number Diff line change
@@ -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)
15 changes: 14 additions & 1 deletion tests/unit/adapters/test_base.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
from pathlib import Path
from types import NotImplementedType
from unittest import mock

import click
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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")
Expand Down
36 changes: 36 additions & 0 deletions tests/unit/api/test_api.py
Original file line number Diff line number Diff line change
@@ -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"
12 changes: 0 additions & 12 deletions tests/unit/test_api.py

This file was deleted.

4 changes: 2 additions & 2 deletions tests/unit/test_default.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

0 comments on commit dc98d01

Please sign in to comment.