Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Feat/83 python api (cont) #94

Merged
merged 5 commits into from
Feb 24, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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
Loading