dbterd.api

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