dvt-core 1.11.0b4__py3-none-any.whl

dvt/plugins/manager.py

@@ -0,0 +1,164 @@
+import functools
+import importlib
+import pkgutil
+from types import ModuleType
+from typing import Callable, Dict, List, Mapping
+
+import dvt.tracking
+from dvt.contracts.graph.manifest import Manifest
+from dvt.plugins.contracts import PluginArtifacts
+from dvt.plugins.manifest import PluginNodes
+
+from dbt_common.exceptions import DbtRuntimeError
+from dbt_common.tests import test_caching_enabled
+
+
+def dbt_hook(func):
+    def inner(*args, **kwargs):
+        try:
+            return func(*args, **kwargs)
+        except Exception as e:
+            raise DbtRuntimeError(f"{func.__name__}: {e}")
+
+    setattr(inner, "is_dbt_hook", True)
+    return inner
+
+
+class dbtPlugin:
+    """
+    EXPERIMENTAL: dbtPlugin is the base class for creating plugins.
+    Its interface is **not** stable and will likely change between dbt-core versions.
+    """
+
+    def __init__(self, project_name: str) -> None:
+        self.project_name = project_name
+        try:
+            self.initialize()
+        except DbtRuntimeError as e:
+            # Remove the first line of DbtRuntimeError to avoid redundant "Runtime Error" line
+            raise DbtRuntimeError("\n".join(str(e).split("\n")[1:]))
+        except Exception as e:
+            raise DbtRuntimeError(str(e))
+
+    @property
+    def name(self) -> str:
+        return self.__class__.__name__
+
+    def initialize(self) -> None:
+        """
+        Initialize the plugin. This function may be overridden by subclasses that have
+        additional initialization steps.
+        """
+        pass
+
+    def get_nodes(self) -> PluginNodes:
+        """
+        Provide PluginNodes to dbt for injection into dbt's DAG.
+        Currently the only node types that are accepted are model nodes.
+        """
+        raise NotImplementedError(f"get_nodes hook not implemented for {self.name}")
+
+    def get_manifest_artifacts(self, manifest: Manifest) -> PluginArtifacts:
+        """
+        Given a manifest, provide PluginArtifacts derived for writing by core.
+        PluginArtifacts share the same lifecycle as the manifest.json file -- they
+        will either be written or not depending on whether the manifest is written.
+        """
+        raise NotImplementedError(f"get_manifest_artifacts hook not implemented for {self.name}")
+
+
+@functools.lru_cache(maxsize=None)
+def _get_dbt_modules() -> Mapping[str, ModuleType]:
+    # This is an expensive function, especially in the context of testing, when
+    # it is called repeatedly, so we break it out and cache the result globally.
+    return {
+        name: importlib.import_module(name)
+        for _, name, _ in pkgutil.iter_modules()
+        if name.startswith(PluginManager.PLUGIN_MODULE_PREFIX)
+    }
+
+
+_MODULES_CACHE = None
+
+
+class PluginManager:
+    PLUGIN_MODULE_PREFIX = "dbt_"
+    PLUGIN_ATTR_NAME = "plugins"
+
+    def __init__(self, plugins: List[dbtPlugin]) -> None:
+        self._plugins = plugins
+        self._valid_hook_names = set()
+        # default hook implementations from dbtPlugin
+        for hook_name in dir(dbtPlugin):
+            if not hook_name.startswith("_"):
+                self._valid_hook_names.add(hook_name)
+
+        self.hooks: Dict[str, List[Callable]] = {}
+        for plugin in self._plugins:
+            for hook_name in dir(plugin):
+                hook = getattr(plugin, hook_name)
+                if (
+                    callable(hook)
+                    and hasattr(hook, "is_dbt_hook")
+                    and hook_name in self._valid_hook_names
+                ):
+                    if hook_name in self.hooks:
+                        self.hooks[hook_name].append(hook)
+                    else:
+                        self.hooks[hook_name] = [hook]
+
+    @classmethod
+    def from_modules(cls, project_name: str) -> "PluginManager":
+
+        if test_caching_enabled():
+            global _MODULES_CACHE
+            if _MODULES_CACHE is None:
+                discovered_dbt_modules = cls.get_prefixed_modules()
+                _MODULES_CACHE = discovered_dbt_modules
+            else:
+                discovered_dbt_modules = _MODULES_CACHE
+        else:
+            discovered_dbt_modules = cls.get_prefixed_modules()
+
+        plugins = []
+        for name, module in discovered_dbt_modules.items():
+            if hasattr(module, cls.PLUGIN_ATTR_NAME):
+                available_plugins = getattr(module, cls.PLUGIN_ATTR_NAME, [])
+                for plugin_cls in available_plugins:
+                    assert issubclass(
+                        plugin_cls, dbtPlugin
+                    ), f"'plugin' in {name} must be subclass of dbtPlugin"
+                    plugin = plugin_cls(project_name=project_name)
+                    plugins.append(plugin)
+        return cls(plugins=plugins)
+
+    @classmethod
+    def get_prefixed_modules(cls):
+        return {
+            name: importlib.import_module(name)
+            for _, name, _ in pkgutil.iter_modules()
+            if name.startswith(cls.PLUGIN_MODULE_PREFIX)
+        }
+
+    def get_manifest_artifacts(self, manifest: Manifest) -> PluginArtifacts:
+        all_plugin_artifacts = {}
+        for hook_method in self.hooks.get("get_manifest_artifacts", []):
+            plugin_artifacts = hook_method(manifest)
+            all_plugin_artifacts.update(plugin_artifacts)
+        return all_plugin_artifacts
+
+    def get_nodes(self) -> PluginNodes:
+        all_plugin_nodes = PluginNodes()
+        for hook_method in self.hooks.get("get_nodes", []):
+            plugin_nodes = hook_method()
+            dbt.tracking.track_plugin_get_nodes(
+                {
+                    "plugin_name": hook_method.__self__.name,  # type: ignore
+                    "num_model_nodes": len(plugin_nodes.models),
+                    "num_model_packages": len(
+                        {model.package_name for model in plugin_nodes.models.values()}
+                    ),
+                }
+            )
+            all_plugin_nodes.update(plugin_nodes)
+        return all_plugin_nodes

dvt/plugins/manifest.py

@@ -0,0 +1,21 @@
+from dataclasses import dataclass, field
+from typing import Dict
+
+from dvt.artifacts.resources import NodeVersion  # noqa
+
+# all these are just exports, they need "noqa" so flake8 will not complain.
+from dvt.contracts.graph.manifest import Manifest  # noqa
+from dvt.contracts.graph.node_args import ModelNodeArgs
+from dvt.graph.graph import UniqueId  # noqa
+from dvt.node_types import AccessType, NodeType  # noqa
+
+
+@dataclass
+class PluginNodes:
+    models: Dict[str, ModelNodeArgs] = field(default_factory=dict)
+
+    def add_model(self, model_args: ModelNodeArgs) -> None:
+        self.models[model_args.unique_id] = model_args
+
+    def update(self, other: "PluginNodes") -> None:
+        self.models.update(other.models)

dvt/profiler.py

@@ -0,0 +1,20 @@
+from contextlib import contextmanager
+from cProfile import Profile
+from pstats import Stats
+from typing import Any, Generator
+
+
+@contextmanager
+def profiler(enable: bool, outfile: str) -> Generator[Any, None, None]:
+    try:
+        if enable:
+            profiler = Profile()
+            profiler.enable()
+
+        yield
+    finally:
+        if enable:
+            profiler.disable()
+            stats = Stats(profiler)
+            stats.sort_stats("tottime")
+            stats.dump_stats(str(outfile))

dvt/py.typed

@@ -0,0 +1 @@
+# dummy file, our types are defined inline

dvt/runners/__init__.py

@@ -0,0 +1,2 @@
+from .exposure_runner import ExposureRunner
+from .saved_query_runner import SavedQueryRunner

dvt/runners/exposure_runner.py

@@ -0,0 +1,7 @@
+from dvt.runners.no_op_runner import NoOpRunner
+
+
+class ExposureRunner(NoOpRunner):
+    @property
+    def description(self) -> str:
+        return f"exposure {self.node.name}"

dvt/runners/no_op_runner.py

@@ -0,0 +1,46 @@
+import threading
+
+from dvt.artifacts.schemas.results import RunStatus
+from dvt.artifacts.schemas.run import RunResult
+from dvt.contracts.graph.manifest import Manifest
+from dvt.events.types import LogNodeNoOpResult
+from dvt.task.base import BaseRunner
+
+from dbt_common.events.functions import fire_event
+
+
+class NoOpRunner(BaseRunner):
+    @property
+    def description(self) -> str:
+        raise NotImplementedError("description not implemented")
+
+    def before_execute(self) -> None:
+        pass
+
+    def compile(self, manifest: Manifest):
+        return self.node
+
+    def after_execute(self, result) -> None:
+        fire_event(
+            LogNodeNoOpResult(
+                description=self.description,
+                index=self.node_index,
+                total=self.num_nodes,
+                node_info=self.node.node_info,
+            )
+        )
+
+    def execute(self, compiled_node, manifest):
+        # no-op
+        return RunResult(
+            node=compiled_node,
+            status=RunStatus.NoOp,
+            timing=[],
+            thread_id=threading.current_thread().name,
+            execution_time=0,
+            message="NO-OP",
+            adapter_response={},
+            failures=0,
+            batch_results=None,
+            agate_table=None,
+        )

dvt/runners/saved_query_runner.py

@@ -0,0 +1,7 @@
+from dvt.runners.no_op_runner import NoOpRunner
+
+
+class SavedQueryRunner(NoOpRunner):
+    @property
+    def description(self) -> str:
+        return f"saved query {self.node.name}"

dvt/selected_resources.py

@@ -0,0 +1,8 @@
+from typing import Any, Set
+
+SELECTED_RESOURCES = []
+
+
+def set_selected_resources(selected_resources: Set[Any]) -> None:
+    global SELECTED_RESOURCES
+    SELECTED_RESOURCES = list(selected_resources)