Skip to content

Commit

Permalink
test tracking
Browse files Browse the repository at this point in the history
  • Loading branch information
MichelleArk committed Aug 7, 2023
1 parent 556b404 commit 2c34de3
Show file tree
Hide file tree
Showing 3 changed files with 37 additions and 1 deletion.
10 changes: 10 additions & 0 deletions core/dbt/plugins/manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
from dbt.exceptions import DbtRuntimeError
from dbt.plugins.contracts import PluginArtifacts
from dbt.plugins.manifest import PluginNodes
import dbt.tracking


def dbt_hook(func):
Expand Down Expand Up @@ -119,5 +120,14 @@ 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
14 changes: 14 additions & 0 deletions core/dbt/tracking.py
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@
RPC_REQUEST_SPEC = "iglu:com.dbt/rpc_request/jsonschema/1-0-1"
RUNNABLE_TIMING = "iglu:com.dbt/runnable/jsonschema/1-0-0"
RUN_MODEL_SPEC = "iglu:com.dbt/run_model/jsonschema/1-0-3"
PLUGIN_GET_NODES = "iglu:com.dbt/plugin_get_nodes/jsonschema/1-0-0"


class TimeoutEmitter(Emitter):
Expand Down Expand Up @@ -409,6 +410,19 @@ def track_partial_parser(options):
)


def track_plugin_get_nodes(options):
context = [SelfDescribingJson(PLUGIN_GET_NODES, options)]
assert active_user is not None, "Cannot track plugin node info when active user is None"

track(
active_user,
category="dbt",
action="plugin_get_nodes",
label=get_invocation_id(),
context=context,
)


def track_runnable_timing(options):
context = [SelfDescribingJson(RUNNABLE_TIMING, options)]
assert active_user is not None, "Cannot track runnable info when active user is None"
Expand Down
14 changes: 13 additions & 1 deletion tests/unit/test_plugin_manager.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import pytest
from unittest import mock

from dbt.exceptions import DbtRuntimeError
from dbt.plugins import PluginManager, dbtPlugin, dbt_hook
Expand Down Expand Up @@ -92,10 +93,21 @@ def test_plugin_manager_init_multiple_hooks(self, get_nodes_plugin, get_artifact
assert len(pm.hooks["get_manifest_artifacts"]) == 1
assert pm.hooks["get_manifest_artifacts"][0] == get_artifacts_plugin.get_manifest_artifacts

def test_get_nodes(self, get_nodes_plugins):
@mock.patch("dbt.tracking")
def test_get_nodes(self, tracking, get_nodes_plugins):
tracking.active_user = mock.Mock()
pm = PluginManager(plugins=get_nodes_plugins)

nodes = pm.get_nodes()

assert len(nodes.models) == 2
assert tracking.track_plugin_get_nodes.called_once_with(
{
"plugin_name": get_nodes_plugins[0].name,
"num_model_nodes": 2,
"num_model_packages": 1,
}
)

def test_get_manifest_artifact(self, get_artifacts_plugins):
pm = PluginManager(plugins=get_artifacts_plugins)
Expand Down

0 comments on commit 2c34de3

Please sign in to comment.