instrument NADE commands using metric spans (#1844)

* draft PR for feedback on instrumentation

* remove changes on loggers

* extract enum value for proper key to message dict

* write custom decorator for spans, add tests

* rename spans as logical steps, remove extra param from console phase, un-parametrize tests, merge test files to keep with covention

* adjust tests for new spans

* time.monotonic() -> time.perf_counter() for proper ordering on windows

* add span for getting snowsight url since it is a blind spot

* add event sharing counters from merge conflicts to new file

* add span for create_or_upgrade_app to address blind spot

* adjust span names for consistency

* add docstring for entity action span decorator

* fix spans bundle test

* sort spans deterministically by creation order

* only one sort, remove counter

* update comment for span decorator to reflect new method name

.github/CODEOWNERS

@@ -7,7 +7,7 @@
 src/snowflake/cli/api/metrics.py @snowflakedb/snowcli @snowflakedb/nade
 src/snowflake/cli/api/project/definition_conversion.py @snowflakedb/snowcli @snowflakedb/nade
 src/snowflake/cli/api/utils/ @snowflakedb/snowcli @snowflakedb/nade
-tests/api/metrics/ @snowflakedb/snowcli @snowflakedb/nade
+tests/api/test_metrics.py @snowflakedb/snowcli @snowflakedb/nade
 tests_common/__init__.py @snowflakedb/snowcli @snowflakedb/nade
 tests_common/conftest.py @snowflakedb/snowcli @snowflakedb/nade
 tests_common/deflake.py @snowflakedb/snowcli @snowflakedb/nade

src/snowflake/cli/_app/telemetry.py

@@ -67,6 +67,10 @@ class CLITelemetryField(Enum):
     CONFIG_FEATURE_FLAGS = "config_feature_flags"
     # Metrics
     COUNTERS = "counters"
+    SPANS = "spans"
+    COMPLETED_SPANS = "completed_spans"
+    NUM_SPANS_PAST_DEPTH_LIMIT = "num_spans_past_depth_limit"
+    NUM_SPANS_PAST_TOTAL_LIMIT = "num_spans_past_total_limit"
     # Information
     EVENT = "event"
     ERROR_MSG = "error_msg"
@@ -129,9 +133,12 @@ def _get_command_metrics() -> TelemetryDict:
     cli_context = get_cli_context()
 
     return {
-        CLITelemetryField.COUNTERS: {
-            **cli_context.metrics.counters,
-        }
+        CLITelemetryField.COUNTERS: cli_context.metrics.counters,
+        CLITelemetryField.SPANS: {
+            CLITelemetryField.COMPLETED_SPANS.value: cli_context.metrics.completed_spans,
+            CLITelemetryField.NUM_SPANS_PAST_DEPTH_LIMIT.value: cli_context.metrics.num_spans_past_depth_limit,
+            CLITelemetryField.NUM_SPANS_PAST_TOTAL_LIMIT.value: cli_context.metrics.num_spans_past_total_limit,
+        },
     }
 
 

src/snowflake/cli/_plugins/nativeapp/artifacts.py

@@ -22,6 +22,7 @@
 from typing import Any, Callable, Dict, Iterable, Iterator, List, Optional, Tuple, Union
 
 from click.exceptions import ClickException
+from snowflake.cli.api.cli_global_context import span
 from snowflake.cli.api.constants import DEFAULT_SIZE_LIMIT_MB
 from snowflake.cli.api.project.schemas.v1.native_app.path_mapping import PathMapping
 from snowflake.cli.api.project.util import to_identifier
@@ -649,6 +650,7 @@ def resolve_without_follow(path: Path) -> Path:
     return Path(os.path.abspath(path))
 
 
+@span("bundle")
 def build_bundle(
     project_root: Path,
     deploy_root: Path,

src/snowflake/cli/_plugins/nativeapp/codegen/compiler.py

@@ -88,7 +88,10 @@ def compile_artifacts(self):
         if not self._should_invoke_processors():
             return
 
-        with cc.phase("Invoking artifact processors"):
+        with (
+            cc.phase("Invoking artifact processors"),
+            get_cli_context().metrics.span("artifact_processors"),
+        ):
             if self._bundle_ctx.generated_root.exists():
                 raise ClickException(
                     f"Path {self._bundle_ctx.generated_root} already exists. Please choose a different name for your generated directory in the project definition file."

src/snowflake/cli/_plugins/nativeapp/codegen/snowpark/python_processor.py

@@ -48,7 +48,7 @@
     NativeAppExtensionFunction,
 )
 from snowflake.cli._plugins.stage.diff import to_stage_path
-from snowflake.cli.api.cli_global_context import get_cli_context
+from snowflake.cli.api.cli_global_context import get_cli_context, span
 from snowflake.cli.api.console import cli_console as cc
 from snowflake.cli.api.metrics import CLICounterField
 from snowflake.cli.api.project.schemas.v1.native_app.path_mapping import (
@@ -167,6 +167,7 @@ class SnowparkAnnotationProcessor(ArtifactProcessor):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
+    @span("snowpark_processor")
     def process(
         self,
         artifact_to_process: PathMapping,

src/snowflake/cli/_plugins/nativeapp/codegen/templates/templates_processor.py

@@ -23,7 +23,7 @@
     ArtifactProcessor,
 )
 from snowflake.cli._plugins.nativeapp.exceptions import InvalidTemplateInFileError
-from snowflake.cli.api.cli_global_context import get_cli_context
+from snowflake.cli.api.cli_global_context import get_cli_context, span
 from snowflake.cli.api.console import cli_console as cc
 from snowflake.cli.api.metrics import CLICounterField
 from snowflake.cli.api.project.schemas.v1.native_app.path_mapping import (
@@ -91,6 +91,7 @@ def expand_templates_in_file(
                 if expanded_template != file.contents:
                     file.edited_contents = expanded_template
 
+    @span("templates_processor")
     def process(
         self,
         artifact_to_process: PathMapping,

src/snowflake/cli/_plugins/nativeapp/entities/application.py

@@ -55,9 +55,13 @@
 from snowflake.cli._plugins.nativeapp.sf_facade import get_snowflake_facade
 from snowflake.cli._plugins.nativeapp.utils import needs_confirmation
 from snowflake.cli._plugins.workspace.context import ActionContext
-from snowflake.cli.api.cli_global_context import get_cli_context
+from snowflake.cli.api.cli_global_context import get_cli_context, span
 from snowflake.cli.api.console.abc import AbstractConsole
-from snowflake.cli.api.entities.common import EntityBase, get_sql_executor
+from snowflake.cli.api.entities.common import (
+    EntityBase,
+    attach_spans_to_entity_actions,
+    get_sql_executor,
+)
 from snowflake.cli.api.entities.utils import (
     drop_generic_object,
     execute_post_deploy_hooks,
@@ -321,6 +325,7 @@ def append_test_resource_suffix_to_identifier(
         return with_suffix
 
 
+@attach_spans_to_entity_actions(entity_name="app")
 class ApplicationEntity(EntityBase[ApplicationEntityModel]):
     """
     A Native App application object, created from an application package.
@@ -629,6 +634,7 @@ def get_objects_owned_by_application(self) -> List[ApplicationOwnedObject]:
             ).fetchall()
             return [{"name": row[1], "type": row[2]} for row in results]
 
+    @span("update_app_object")
     def create_or_upgrade_app(
         self,
         package: ApplicationPackageEntity,
@@ -1043,6 +1049,7 @@ def stream_events(
         except KeyboardInterrupt:
             return
 
+    @span("get_snowsight_url_for_app")
     def get_snowsight_url(self) -> str:
         """Returns the URL that can be used to visit this app via Snowsight."""
         name = identifier_for_url(self.name)

src/snowflake/cli/_plugins/nativeapp/entities/application_package.py

@@ -49,7 +49,12 @@
 from snowflake.cli._plugins.stage.diff import DiffResult
 from snowflake.cli._plugins.stage.manager import StageManager
 from snowflake.cli._plugins.workspace.context import ActionContext
-from snowflake.cli.api.entities.common import EntityBase, get_sql_executor
+from snowflake.cli.api.cli_global_context import span
+from snowflake.cli.api.entities.common import (
+    EntityBase,
+    attach_spans_to_entity_actions,
+    get_sql_executor,
+)
 from snowflake.cli.api.entities.utils import (
     drop_generic_object,
     execute_post_deploy_hooks,
@@ -157,6 +162,7 @@ def validate_source_stage(cls, input_value: str):
         return input_value
 
 
+@attach_spans_to_entity_actions(entity_name="app_pkg")
 class ApplicationPackageEntity(EntityBase[ApplicationPackageEntityModel]):
     """
     A Native App application package.
@@ -900,6 +906,7 @@ def validate_setup_script(
             if validation_result["status"] == "FAIL":
                 raise SetupScriptFailedValidation()
 
+    @span("validate_setup_script")
     def get_validation_result(
         self, use_scratch_stage: bool, interactive: bool, force: bool
     ):

src/snowflake/cli/api/cli_global_context.py

@@ -17,6 +17,7 @@
 from contextlib import contextmanager
 from contextvars import ContextVar
 from dataclasses import dataclass, field, replace
+from functools import wraps
 from pathlib import Path
 from typing import TYPE_CHECKING, Iterator
 
@@ -157,7 +158,7 @@ def enable_tracebacks(self) -> bool:
         return self._manager.enable_tracebacks
 
     @property
-    def metrics(self):
+    def metrics(self) -> CLIMetrics:
         return self._manager.metrics
 
     @property
@@ -213,6 +214,26 @@ def get_cli_context() -> _CliGlobalContextAccess:
     return _CliGlobalContextAccess(get_cli_context_manager())
 
 
+def span(span_name: str):
+    """
+    Decorator to start a command metrics span that encompasses a whole function
+
+    Must be used instead of directly calling @get_cli_context().metrics.span(span_name)
+    as a decorator to ensure that the cli context is grabbed at run time instead of at
+    module load time, which would not reflect forking
+    """
+
+    def decorator(func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            with get_cli_context().metrics.span(span_name):
+                return func(*args, **kwargs)
+
+        return wrapper
+
+    return decorator
+
+
 @contextmanager
 def fork_cli_context(
     connection_overrides: dict | None = None,

src/snowflake/cli/api/entities/common.py

@@ -2,6 +2,7 @@
 from typing import Generic, Type, TypeVar, get_args
 
 from snowflake.cli._plugins.workspace.context import ActionContext, WorkspaceContext
+from snowflake.cli.api.cli_global_context import span
 from snowflake.cli.api.sql_execution import SqlExecutor
 
 
@@ -20,6 +21,31 @@ class EntityActions(str, Enum):
 T = TypeVar("T")
 
 
+def attach_spans_to_entity_actions(entity_name: str):
+    """
+    Class decorator for EntityBase subclasses to automatically wrap
+    every implemented entity action method with a metrics span
+
+    Args:
+        entity_name (str): Custom name for entity type to be displayed in metrics
+    """
+
+    def decorator(cls: type[EntityBase]) -> type[EntityBase]:
+        for attr_name, attr_value in vars(cls).items():
+            is_entity_action = attr_name in [
+                enum_member for enum_member in EntityActions
+            ]
+
+            if is_entity_action and callable(attr_value):
+                attr_name_without_action_prefix = attr_name.partition("_")[2]
+                span_name = f"action.{entity_name}.{attr_name_without_action_prefix}"
+                action_with_span = span(span_name)(attr_value)
+                setattr(cls, attr_name, action_with_span)
+        return cls
+
+    return decorator
+
+
 class EntityBase(Generic[T]):
     """
     Base class for the fully-featured entity classes.

src/snowflake/cli/api/entities/utils.py

@@ -23,7 +23,7 @@
     to_stage_path,
 )
 from snowflake.cli._plugins.stage.utils import print_diff_to_console
-from snowflake.cli.api.cli_global_context import get_cli_context
+from snowflake.cli.api.cli_global_context import get_cli_context, span
 from snowflake.cli.api.console.abc import AbstractConsole
 from snowflake.cli.api.entities.common import get_sql_executor
 from snowflake.cli.api.errno import (
@@ -76,6 +76,7 @@ def _get_stage_paths_to_sync(
     return stage_paths
 
 
+@span("sync_deploy_root_with_stage")
 def sync_deploy_root_with_stage(
     console: AbstractConsole,
     deploy_root: Path,
@@ -212,7 +213,10 @@ def execute_post_deploy_hooks(
 
     get_cli_context().metrics.set_counter(CLICounterField.POST_DEPLOY_SCRIPTS, 1)
 
-    with console.phase(f"Executing {deployed_object_type} post-deploy actions"):
+    with (
+        console.phase(f"Executing {deployed_object_type} post-deploy actions"),
+        get_cli_context().metrics.span("post_deploy_hooks"),
+    ):
         sql_scripts_paths = []
         display_paths = []
         for hook in post_deploy_hooks:

src/snowflake/cli/api/metrics.py

@@ -116,10 +116,8 @@ class for holding metrics span data and encapsulating related operations
     children: Set[CLIMetricsSpan] = field(init=False, default_factory=set)
 
     # private state
-    # start time of the step from the monotonic clock in order to calculate execution time
-    _monotonic_start: float = field(
-        init=False, default_factory=lambda: time.monotonic()
-    )
+    # start time of the step from a performance counter in order to calculate execution time
+    _start_time: float = field(init=False, default_factory=time.perf_counter)
 
     def __hash__(self) -> int:
         return hash(self.span_id)
@@ -154,7 +152,7 @@ def finish(self, error: Optional[BaseException] = None) -> None:
         if error:
             self.error = error
 
-        self.execution_time = time.monotonic() - self._monotonic_start
+        self.execution_time = time.perf_counter() - self._start_time
 
     def to_dict(self) -> Dict:
         """
@@ -191,9 +189,10 @@ class CLIMetrics:
     _in_progress_spans: List[CLIMetricsSpan] = field(init=False, default_factory=list)
     # list of finished steps for telemetry to process
     _completed_spans: List[CLIMetricsSpan] = field(init=False, default_factory=list)
-    # monotonic clock time of when this class was initialized to approximate when the command first started executing
-    _monotonic_start: float = field(
-        init=False, default_factory=lambda: time.monotonic(), compare=False
+    # clock time of a performance counter when this class was initialized
+    # to approximate when the command first started executing
+    _start_time: float = field(
+        init=False, default_factory=time.perf_counter, compare=False
     )
 
     def clone(self) -> CLIMetrics:
@@ -223,7 +222,7 @@ def current_span(self) -> Optional[CLIMetricsSpan]:
         return self._in_progress_spans[-1] if len(self._in_progress_spans) > 0 else None
 
     @contextmanager
-    def start_span(self, name: str) -> Iterator[CLIMetricsSpan]:
+    def span(self, name: str) -> Iterator[CLIMetricsSpan]:
         """
         Starts a new span that tracks various metrics throughout its execution
 
@@ -236,7 +235,7 @@ def start_span(self, name: str) -> Iterator[CLIMetricsSpan]:
         """
         new_span = CLIMetricsSpan(
             name=name,
-            start_time=time.monotonic() - self._monotonic_start,
+            start_time=time.perf_counter() - self._start_time,
             parent=self.current_span,
         )
 
@@ -275,34 +274,29 @@ def num_spans_past_total_limit(self) -> int:
     @property
     def completed_spans(self) -> List[Dict]:
         """
-        Returns the completed spans tracked throughout a command, sorted by start time, for reporting telemetry
+        Returns the completed spans tracked throughout a command for reporting telemetry
 
         Ensures that the spans we send are within the configured limits and marks
         certain spans as trimmed if their children would bypass the limits we set
         """
         # take spans breadth-first within the depth and total limits
         # since we care more about the big picture than granularity
-        spans_to_report = set(
-            nsmallest(
-                n=self.SPAN_TOTAL_LIMIT,
-                iterable=(
-                    span
-                    for span in self._completed_spans
-                    if span.span_depth <= self.SPAN_DEPTH_LIMIT
-                ),
-                key=lambda span: (span.span_depth, span.start_time),
-            )
+        spans_to_report = nsmallest(
+            n=self.SPAN_TOTAL_LIMIT,
+            iterable=(
+                span
+                for span in self._completed_spans
+                if span.span_depth <= self.SPAN_DEPTH_LIMIT
+            ),
+            key=lambda span: (span.span_depth, span.start_time, span.execution_time),
         )
 
-        # sort by start time to make reading the payload easier
-        sorted_spans_to_report = sorted(
-            spans_to_report, key=lambda span: span.start_time
-        )
+        spans_to_report_set = set(spans_to_report)
 
         return [
             {
                 **span.to_dict(),
-                CLIMetricsSpan.TRIMMED_KEY: not span.children <= spans_to_report,
+                CLIMetricsSpan.TRIMMED_KEY: not span.children <= spans_to_report_set,
             }
-            for span in sorted_spans_to_report
+            for span in spans_to_report
         ]