从源码看dbt-loom 是如何开发dbt plugin的

以前简单介绍过dbt plugin 有一套明确的格式，实际上与dbt adapter 套路是类似的，以下从源码角度看看dbt-loom 的设计

dbt-loom 参考架构

从此图可以看出，dbt-loom 可以从其他地方获取dbt 的元数据（dbt cloud，本地文件，对象存储），之后基于此插件结合dbt-core 可以进行元数据的合并，之后执行的时候创建对应的模型到数据仓库中

代码说明

• 代码结构

dbt_loom

├── __init__.py

├── clients

│   ├── az_blob.py

│   ├── dbt_cloud.py

│   ├── gcs.py

│   └── s3.py

├── config.py

└── manifests.py

• python 包命名

应该是dbt- 格式的，这个是plugin的要求，pyproject.toml 如下

[tool.poetry]

name = "dbt-loom"

version = "0.5.1"

• 插件入口
  默认需要提供一个plugins 的属性
  init.py 中

plugins = [dbtLoom]

• dbtLoom dbtPlugin 实现

可以看到实现了initialize，get_nodes，但是没有实现get_manifest_artifacts

class dbtLoom(dbtPlugin):

    """

    dbtLoom is a dbt plugin that loads manifest files, parses a DAG from the manifest,

    and injects public nodes from imported manifest.

    """

 

    def __init__(self, project_name: str):

        configuration_path = Path(

            os.environ.get("DBT_LOOM_CONFIG", "dbt_loom.config.yml")

        )

        # 此处是dbt-loom 自己定义的一个ManifestLoader，有点类似dbt-core 的ManifestLoader

        self._manifest_loader = ManifestLoader()

 

        self.config: Optional[dbtLoomConfig] = self.read_config(configuration_path)

        self.models: Dict[str, ModelNodeArgs] = {}

 

        import dbt.contracts.graph.manifest

 

        fire_event(

            Note(

                msg="dbt-loom: Patching ref protection methods to support dbt-loom dependencies."

            )

        )

        dbt.contracts.graph.manifest.Manifest.is_invalid_protected_ref = (  # type: ignore

            self.dependency_wrapper(

                dbt.contracts.graph.manifest.Manifest.is_invalid_protected_ref

            )

        )

        dbt.contracts.graph.manifest.Manifest.is_invalid_private_ref = (  # type: ignore

            self.dependency_wrapper(

                dbt.contracts.graph.manifest.Manifest.is_invalid_private_ref

            )

        )

 

        super().__init__(project_name)

 

    def dependency_wrapper(self, function) -> Callable:

        def outer_function(inner_self, node, target_model, dependencies) -> bool:

            if self.config is not None:

                for manifest in self.config.manifests:

                    dependencies[manifest.name] = LoomRunnableConfig()

 

            return function(inner_self, node, target_model, dependencies)

 

        return outer_function

 

    def read_config(self, path: Path) -> Optional[dbtLoomConfig]:

        """Read the dbt-loom configuration file."""

        if not path.exists():

            return None

 

        with open(path) as file:

            config_content = file.read()

 

        config_content = self.replace_env_variables(config_content)

 

        return dbtLoomConfig(**yaml.load(config_content, yaml.SafeLoader))

 

    @staticmethod

    def replace_env_variables(config_str: str) -> str:

        """Replace environment variable placeholders in the configuration string."""

        pattern = r"\$(\w+)|\$\{([^}]+)\}"

        return re.sub(

            pattern,

            lambda match: os.environ.get(

                match.group(1) if match.group(1) is not None else match.group(2), ""

            ),

            config_str,

        )

   # initialize 方法，此处核心是进行外部元数据的load

    def initialize(self) -> None:

        """Initialize the plugin"""

 

        if self.models != {} or not self.config:

            return

 

        for manifest_reference in self.config.manifests:

            fire_event(

                Note(

                    msg=f"dbt-loom: Loading manifest for `{manifest_reference.name}`"

                    f" from `{manifest_reference.type.value}`"

                )

            )

 

            manifest = self._manifest_loader.load(manifest_reference)

            if manifest is None:

                continue

 

            selected_nodes = identify_node_subgraph(manifest)

            self.models.update(convert_model_nodes_to_model_node_args(selected_nodes))

  # dbt_hook 装饰器

    @dbt_hook

    def get_nodes(self) -> PluginNodes:

        """

        Inject PluginNodes to dbt for injection into dbt's DAG.

        """

        fire_event(Note(msg="dbt-loom: Injecting nodes"))

       # 基于加载的元数据构建自己的node

        return PluginNodes(models=self.models)

loader 处理
核心是读取配置中不同的路径（s3，api，本地）具体可以查看ManifestLoader

说明

dbt-loom 实际上为我们提供了一个比较完整的dbt plugin 开发说明，实际上我们基于此可以完成一些比较有意思的功能，以上是基于配置的元数据，实际上这些元数据也可以是其他数据库的，这样就可以基于数据库自动生成模型的元数据了，比较类似dbt codegen 包的一些能力

参考资料

dbt_loom/manifests.py
https://github.com/nicholasyager/dbt-loom
https://nicholasyager.com/2023/08/dbt_plugin_api.html