dbterd.api

  1import logging
  2from pathlib import Path
  3
  4from click import Command, Context
  5
  6from dbterd import default
  7from dbterd.adapters.base import Executor
  8from dbterd.helpers.log import logger
  9
 10logger.setLevel(logging.WARNING)  # hide log
 11
 12
 13class DbtErd:
 14    """
 15    dbt ERD official API functions.
 16
 17
 18    **Usage**:
 19
 20    ## Get a whole ERD
 21
 22    ```python
 23    from dbterd.api import DbtErd
 24    erd = DbtErd().get_erd()
 25    ```
 26
 27    ## Get a whole ERD given all models attached to `my_exposure_name`
 28
 29    ```python
 30    from dbterd.api import DbtErd
 31    erd = DbtErd(select="exposure:my_exposure_name").get_erd()
 32    ```
 33    See the
 34    [Selection](https://dbterd.datnguyen.de/latest/nav/guide/cli-references.html#dbterd-run-select-s)
 35    page for more details.
 36
 37    ## Get a model (named `model.jaffle_shop.my_model`)'s ERD
 38
 39    ```python
 40    from dbterd.api import DbtErd
 41    erd = DbtErd().get_model_erd(s
 42        node_fqn="model.jaffle_shop.my_model"
 43    )
 44    ```
 45    """
 46
 47    def __init__(self, **kwargs) -> None:
 48        """Initialize the main Executor given similar input CLI parameters"""
 49
 50        self.params: dict = kwargs
 51        """
 52        Mimic CLI params with overriding `api = True`.\n
 53        Check
 54        [CLI reference page](https://dbterd.datnguyen.de/latest/nav/guide/cli-references.html)
 55        for more details of how the params are used.
 56        """
 57        self.__set_params_default_if_not_specified()
 58        ctx_command = self.params.get("api_context_command")
 59        self.executor: Executor = Executor(
 60            Context(Command(name=ctx_command if ctx_command else "run"))
 61        )
 62        """
 63        Mimic CLI's executor.\n
 64        The context command is `run` by default
 65        unless specifying a param named `api_context_command`
 66        """
 67
 68    def __set_params_default_if_not_specified(self) -> None:
 69        """Set base params' default value (mimic CLI behaviors where possible)"""
 70        self.params["api"] = True
 71
 72        self.params["select"] = self.params.get("select", [])
 73        self.params["exclude"] = self.params.get("exclude", [])
 74        self.params["resource_type"] = self.params.get(
 75            "resource_type", default.default_resource_types()
 76        )
 77        self.params["algo"] = self.params.get("algo", default.default_algo())
 78        self.params["entity_name_format"] = self.params.get(
 79            "entity_name_format", default.default_entity_name_format()
 80        )
 81        self.params["omit_columns"] = self.params.get("omit_columns", False)
 82        self.params["artifacts_dir"] = self.params.get("artifacts_dir", Path.cwd())
 83        self.params["target"] = self.params.get("target", default.default_target())
 84
 85    def get_erd(self) -> str:
 86        """Generate ERD code for a whole project
 87
 88        Usage:
 89        ```python
 90        from dbterd.api import DbtErd
 91        erd = DbtErd().get_erd()
 92        ```
 93
 94        Returns:
 95            str: ERD text
 96        """
 97        return self.executor.run(**self.params)
 98
 99    def get_model_erd(self, node_unique_id: str) -> str:
100        """Generate ERD code for a model.
101
102        Result contains the input model and 1 level relationship model(s) (if any).
103
104        Usage:
105
106            ```python
107            from dbterd.api import DbtErd
108            erd = DbtErd().get_model_erd(
109                node_unique_id="model.jaffle_shop.my_model"
110            )
111            ```
112
113        Args:
114            - node_unique_id (str): Manifest node unique ID
115
116        Returns:
117            str: ERD text
118        """
119        return self.executor.run(node_unique_id=node_unique_id, **self.params)