From 94e72cb5d2cec40cff422b969146fce8c39520c3 Mon Sep 17 00:00:00 2001 From: Michelle Ark Date: Fri, 11 Aug 2023 14:44:04 -0400 Subject: [PATCH] track plugin.get_nodes (#8336) (#8362) (cherry picked from commit 83ff38ab24a820d0bfca411ac51b67eeb2ef878d) --- .../Under the Hood-20230809-094834.yaml | 6 ++++++ core/dbt/plugins/manager.py | 10 ++++++++++ core/dbt/tracking.py | 14 ++++++++++++++ tests/unit/test_plugin_manager.py | 15 ++++++++++++++- 4 files changed, 44 insertions(+), 1 deletion(-) create mode 100644 .changes/unreleased/Under the Hood-20230809-094834.yaml diff --git a/.changes/unreleased/Under the Hood-20230809-094834.yaml b/.changes/unreleased/Under the Hood-20230809-094834.yaml new file mode 100644 index 00000000000..64883952ee9 --- /dev/null +++ b/.changes/unreleased/Under the Hood-20230809-094834.yaml @@ -0,0 +1,6 @@ +kind: Under the Hood +body: add tracking for plugin.get_nodes calls +time: 2023-08-09T09:48:34.819445-04:00 +custom: + Author: michelleark + Issue: "8344" diff --git a/core/dbt/plugins/manager.py b/core/dbt/plugins/manager.py index 3a52a5c5180..0cb784f4cac 100644 --- a/core/dbt/plugins/manager.py +++ b/core/dbt/plugins/manager.py @@ -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): @@ -116,5 +117,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 diff --git a/core/dbt/tracking.py b/core/dbt/tracking.py index f306d9f0b00..f7c0a2a7d19 100644 --- a/core/dbt/tracking.py +++ b/core/dbt/tracking.py @@ -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): @@ -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" diff --git a/tests/unit/test_plugin_manager.py b/tests/unit/test_plugin_manager.py index db68fe95c30..f86530ca94a 100644 --- a/tests/unit/test_plugin_manager.py +++ b/tests/unit/test_plugin_manager.py @@ -1,4 +1,6 @@ import pytest +from unittest import mock + from dbt.plugins import PluginManager, dbtPlugin, dbt_hook from dbt.plugins.manifest import PluginNodes, ModelNodeArgs from dbt.plugins.contracts import PluginArtifacts, PluginArtifact @@ -71,10 +73,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)