diff --git a/CHANGELOG.md b/CHANGELOG.md index 173691d02..7e861e88c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,9 @@ ## dbt-databricks next +### Features + +- Add opt-in invocation telemetry for eligible commands via `connection_parameters.enable_dbt_telemetry` ([#1620](https://github.com/databricks/dbt-databricks/pull/1620)) + ### Fixes - Recreate materialized views when query schema drifts, honoring `on_configuration_change` ([#1621](https://github.com/databricks/dbt-databricks/pull/1621) resolves [#1359](https://github.com/databricks/dbt-databricks/issues/1359)) diff --git a/dbt/adapters/databricks/connections.py b/dbt/adapters/databricks/connections.py index b6825edd6..f4c706ab5 100644 --- a/dbt/adapters/databricks/connections.py +++ b/dbt/adapters/databricks/connections.py @@ -55,6 +55,7 @@ from dbt.adapters.databricks.logging import logger from dbt.adapters.databricks.python_models.run_tracking import PythonRunTracker from dbt.adapters.databricks.spog.decision import check_spog_preconditions +from dbt.adapters.databricks.telemetry import hooks as telemetry_hooks from dbt.adapters.databricks.utils import QueryTagsUtils, is_cluster_http_path, redact_credentials if TYPE_CHECKING: @@ -482,7 +483,9 @@ def open(cls, connection: Connection) -> Connection: creds: DatabricksCredentials = connection.credentials timeout = creds.connect_timeout - cls.credentials_manager = creds.authenticate() + # Avoid a manager overwritten by another concurrent open. + credentials_manager = creds.authenticate() + cls.credentials_manager = credentials_manager # SPOG decision matrix: collect every http_path in play (default + # per-compute) and validate them against the host's discovery probe. @@ -502,7 +505,7 @@ def open(cls, connection: Connection) -> Connection: merged_query_tags = QueryConfigUtils.get_merged_query_tags(query_header_context, creds) conn_args = SqlUtils.prepare_connection_arguments( - creds, cls.credentials_manager, databricks_connection.http_path, merged_query_tags + creds, credentials_manager, databricks_connection.http_path, merged_query_tags ) def connect() -> DatabricksHandle: @@ -518,6 +521,8 @@ def connect() -> DatabricksHandle: databricks_connection.capabilities = cls._get_capabilities_for_http_path( databricks_connection.http_path ) + + telemetry_hooks.on_connection_open(creds, credentials_manager) return conn else: raise DbtDatabaseError("Failed to create connection") diff --git a/dbt/adapters/databricks/handle.py b/dbt/adapters/databricks/handle.py index 3da799c9c..356fc932b 100644 --- a/dbt/adapters/databricks/handle.py +++ b/dbt/adapters/databricks/handle.py @@ -394,6 +394,8 @@ def prepare_connection_arguments( connection_parameters = creds.connection_parameters.copy() # type: ignore[union-attr] + connection_parameters.pop("enable_dbt_telemetry", None) + http_headers: list[tuple[str, str]] = list( creds.get_all_http_headers(connection_parameters.pop("http_headers", {})).items() ) diff --git a/dbt/adapters/databricks/impl.py b/dbt/adapters/databricks/impl.py index 6640f0d76..0fceaf3dd 100644 --- a/dbt/adapters/databricks/impl.py +++ b/dbt/adapters/databricks/impl.py @@ -101,6 +101,7 @@ sdk_supports_workspace_id, ) from dbt.adapters.databricks.spog.extract import extract_workspace_id +from dbt.adapters.databricks.telemetry import hooks as telemetry_hooks from dbt.adapters.databricks.utils import ( get_first_row, handle_missing_objects, @@ -295,6 +296,8 @@ def __init__(self, config: Any, mp_context: SpawnContext) -> None: self.get_behavior_flag_no_warn(USE_MANAGED_ICEBERG["name"]) ) + telemetry_hooks.on_adapter_init(self) + # Warehouses always meet capability cutoffs at parse time; clusters keep the # conservative False until a real connection is available. # `_parse_replacements_` is injected by AdapterMeta, so mypy can't resolve it here. @@ -902,6 +905,14 @@ def get_behavior_flag_no_warn(self, behavior_flag_name: str) -> bool: behavior_flag = getattr(self.behavior, behavior_flag_name) return behavior_flag.no_warn + def set_macro_resolver(self, macro_resolver: Any) -> None: + super().set_macro_resolver(macro_resolver) + telemetry_hooks.on_post_parse(self, macro_resolver) + + def cleanup_connections(self) -> None: + telemetry_hooks.on_run_end(self) + super().cleanup_connections() + @available.parse(lambda *a, **k: (None, None)) @record_function( DatabricksAdapterAddQueryRecord, diff --git a/dbt/adapters/databricks/telemetry/__init__.py b/dbt/adapters/databricks/telemetry/__init__.py new file mode 100644 index 000000000..5a36de849 --- /dev/null +++ b/dbt/adapters/databricks/telemetry/__init__.py @@ -0,0 +1,3 @@ +from dbt.adapters.databricks.telemetry.config import is_enabled + +__all__ = ["is_enabled"] diff --git a/dbt/adapters/databricks/telemetry/builder.py b/dbt/adapters/databricks/telemetry/builder.py new file mode 100644 index 000000000..6051efa83 --- /dev/null +++ b/dbt/adapters/databricks/telemetry/builder.py @@ -0,0 +1,400 @@ +from importlib.metadata import version as _pkg_version +from typing import Any, Callable, Optional + +from dbt.adapters.databricks.__version__ import version as _adapter_version +from dbt.adapters.databricks.credentials import DatabricksCredentials +from dbt.adapters.databricks.spog.extract import extract_workspace_id +from dbt.adapters.databricks.telemetry import models + +# Mirrored from impl.py to avoid an import cycle. +_BEHAVIOR_FLAGS = ( + "use_user_folder_for_python", + "use_materialization_v2", + "use_replace_on_for_insert_overwrite", + "use_managed_iceberg", + "use_concurrent_microbatch", + "use_describe_as_json_for_relation_metadata", +) + +_COMMAND_MAP = { + "run": models.DbtCommand.RUN, + "build": models.DbtCommand.BUILD, + "test": models.DbtCommand.TEST, + "seed": models.DbtCommand.SEED, + "snapshot": models.DbtCommand.SNAPSHOT, + "compile": models.DbtCommand.COMPILE, + "docs": models.DbtCommand.DOCS, + "clone": models.DbtCommand.CLONE, + "retry": models.DbtCommand.RETRY, + "show": models.DbtCommand.SHOW, + "list": models.DbtCommand.LIST, + "source": models.DbtCommand.SOURCE, + "run_operation": models.DbtCommand.RUN_OPERATION, +} + + +def classify_compute_type(http_path: Optional[str]) -> models.ComputeType: + if not http_path: + return models.ComputeType.TYPE_UNSPECIFIED + path = http_path.split("?", 1)[0] + if path.startswith(("/sql/1.0/warehouses/", "/sql/1.0/endpoints/")): + return models.ComputeType.SQL_WAREHOUSE + if path.startswith("/sql/protocolv1/"): + return models.ComputeType.ALL_PURPOSE_CLUSTER + return models.ComputeType.OTHER + + +def classify_auth_family(creds: DatabricksCredentials) -> models.AuthFamily: + if getattr(creds, "token", None): + return models.AuthFamily.PAT + if getattr(creds, "azure_client_id", None) and getattr(creds, "azure_client_secret", None): + return models.AuthFamily.AZURE_SERVICE_PRINCIPAL + if not getattr(creds, "client_secret", None): + return models.AuthFamily.OAUTH_U2M + # client_secret is ambiguous between M2M and legacy Azure. + return models.AuthFamily.LEGACY_CLIENT_SECRET_AMBIGUOUS + + +def classify_command(which: Optional[str]) -> models.DbtCommand: + if not which: + return models.DbtCommand.TYPE_UNSPECIFIED + token = str(which).strip().lower().replace("-", "_").split()[0] + return _COMMAND_MAP.get(token, models.DbtCommand.OTHER) + + +def classify_warn_error_policy(warn_error: Any, warn_error_options: Any) -> models.WarnErrorPolicy: + # The legacy boolean takes precedence. + if warn_error: + return models.WarnErrorPolicy.WARN_ERROR_ALL + if warn_error_options: + opts = warn_error_options + get = lambda name: ( # noqa: E731 - keeps object/dict compatibility together + opts.get(name) if isinstance(opts, dict) else getattr(opts, name, None) + ) + error = get("error") or get("include") or [] + warn = get("warn") or get("exclude") or [] + silence = get("silence") or [] + # Named overrides make `error: all` custom. + if error in ("all", "*") and not warn and not silence: + return models.WarnErrorPolicy.WARN_ERROR_ALL + has_policy = bool(error or warn or silence) + if has_policy: + return models.WarnErrorPolicy.WARN_ERROR_CUSTOM_POLICY + return models.WarnErrorPolicy.WARN_ERROR_DISABLED + + +def _resource_type(node: Any) -> str: + rt = getattr(node, "resource_type", None) + return str(getattr(rt, "value", rt)) + + +def _bump(counts: models.ResourceCounts, node: Any, resource_type: str) -> None: + if resource_type == "model": + counts.model_count += 1 + elif resource_type == "test": + counts.data_test_count += 1 + if getattr(node, "test_metadata", None) is not None: + counts.generic_data_test_count += 1 + elif resource_type == "seed": + counts.seed_count += 1 + elif resource_type == "snapshot": + counts.snapshot_count += 1 + elif resource_type == "source": + counts.source_count += 1 + elif resource_type == "function": + counts.function_count += 1 + elif resource_type == "exposure": + counts.exposure_count += 1 + elif resource_type == "saved_query": + counts.saved_query_count += 1 + elif resource_type == "unit_test": + counts.unit_test_count += 1 + else: + counts.other_count += 1 + + +def aggregate_manifest(manifest: Any) -> models.ManifestStats: + stats = models.ManifestStats() + project_name = None + metadata = getattr(manifest, "metadata", None) + if metadata is not None: + project_name = getattr(metadata, "project_name", None) + + collections = [ + "nodes", + "sources", + "exposures", + "metrics", + "semantic_models", + "saved_queries", + "functions", + "unit_tests", + ] + for collection in collections: + items = getattr(manifest, collection, None) + if not items: + continue + for node in items.values(): + resource_type = _resource_type(node) + _bump(stats.enabled_total, node, resource_type) + is_root = getattr(node, "package_name", None) == project_name + _bump( + stats.enabled_root_project if is_root else stats.enabled_installed_packages, + node, + resource_type, + ) + return stats + + +def ephemeral_resource_ids(manifest: Any) -> set[str]: + """Return ephemeral IDs for local counting only.""" + result = set() + for node in (getattr(manifest, "nodes", None) or {}).values(): + config = getattr(node, "config", None) + is_ephemeral = bool(getattr(node, "is_ephemeral_model", False)) or ( + getattr(config, "materialized", None) == "ephemeral" + ) + unique_id = getattr(node, "unique_id", None) + if is_ephemeral and unique_id: + result.add(str(unique_id)) + return result + + +def _get_flags() -> Any: + try: + from dbt.flags import get_flags + + return get_flags() + except Exception: + return None + + +def build_invocation_config(config: Any) -> models.InvocationConfig: + flags = _get_flags() + thread_count = getattr(config, "threads", None) or getattr(flags, "THREADS", None) or 0 + return models.InvocationConfig( + thread_count=int(thread_count), + dbt_command=classify_command(getattr(flags, "WHICH", None)), + full_refresh=bool(getattr(flags, "FULL_REFRESH", False)), + empty=bool(getattr(flags, "EMPTY", False)), + fail_fast=bool(getattr(flags, "FAIL_FAST", False)), + warn_error_policy=classify_warn_error_policy( + getattr(flags, "WARN_ERROR", None), getattr(flags, "WARN_ERROR_OPTIONS", None) + ), + ) + + +def build_connection_config(creds: DatabricksCredentials) -> models.ConnectionConfig: + http_path = getattr(creds, "http_path", None) + connection_parameters = getattr(creds, "connection_parameters", None) or {} + return models.ConnectionConfig( + default_compute_type=classify_compute_type(http_path), + configured_auth_family=classify_auth_family(creds), + named_compute_count=len(getattr(creds, "compute", None) or {}), + # Parse only the `o` parameter; discard its value. + spog_routing_configured=extract_workspace_id(http_path) is not None, + use_kernel=bool(connection_parameters.get("use_kernel")), + ) + + +def build_project_config(behavior_flag: Callable[[str], bool]) -> models.ProjectConfig: + values = {name: bool(behavior_flag(name)) for name in _BEHAVIOR_FLAGS} + return models.ProjectConfig(**values) + + +def build_post_parse_log( + manifest: Any, + config: Any, + creds: DatabricksCredentials, + behavior_flag: Callable[[str], bool], +) -> models.TelemetryLog: + invocation_id = _invocation_id(manifest) + payload = models.PostParsePayload( + invocation_config=build_invocation_config(config), + manifest_stats=aggregate_manifest(manifest), + connection_config=build_connection_config(creds), + project_config=build_project_config(behavior_flag), + ) + return models.TelemetryLog( + invocation_id=invocation_id, + adapter_version=_adapter_version, + dbt_core_version=_dbt_core_version(), + post_parse=payload, + ) + + +def _invocation_id(manifest: Any) -> str: + try: + from dbt_common.invocation import get_invocation_id + + invocation_id = get_invocation_id() + if invocation_id: + return str(invocation_id) + except Exception: + pass + metadata = getattr(manifest, "metadata", None) + return str(getattr(metadata, "invocation_id", "") or "") + + +def _dbt_core_version() -> str: + try: + return _pkg_version("dbt-core") + except Exception: + return "" + + +_STATUS_ATTR = { + "success": "success", + "error": "error", + "fail": "fail", + "warn": "warn", + "skipped": "skipped", + "partial_success": "partial_success", + "pass": "pass_", + "runtime_error": "runtime_error", + "no_op": "no_op", + "reused": "reused", +} +_STATUS_BUCKETS = tuple(dict.fromkeys(_STATUS_ATTR.values())) + +_RESOURCE_TYPE = { + "model": models.ResourceType.MODEL, + "test": models.ResourceType.DATA_TEST, + "unit_test": models.ResourceType.UNIT_TEST, + "seed": models.ResourceType.SEED, + "snapshot": models.ResourceType.SNAPSHOT, + "source": models.ResourceType.SOURCE, + "function": models.ResourceType.FUNCTION, + "exposure": models.ResourceType.EXPOSURE, + "saved_query": models.ResourceType.SAVED_QUERY, +} + +_AUXILIARY_TYPES = {"operation", "hook"} + + +def _norm(value: Any) -> str: + return str(value).strip().lower().replace("-", "_").replace(" ", "_") + + +def _set_total(counts: models.NodeStatusCounts) -> None: + counts.total = sum(getattr(counts, b) for b in _STATUS_BUCKETS) + + +def _bump_status(counts: models.NodeStatusCounts, status: Any) -> bool: + attr = _STATUS_ATTR.get(_norm(status)) + if attr is None: + return False + setattr(counts, attr, getattr(counts, attr) + 1) + return True + + +def _resource_from_uid(unique_id: Any) -> str: + return _norm(str(unique_id).split(".", 1)[0]) + + +def aggregate_node_results(results: list) -> tuple: + result_counts = models.NodeStatusCounts() + auxiliary = models.NodeStatusCounts() + by_type: dict = {} + unknown = 0 + for unique_id, status in results: + rtype = _resource_from_uid(unique_id) + if rtype in _AUXILIARY_TYPES: + _bump_status(auxiliary, status) + continue + if not _bump_status(result_counts, status): + continue + enum = _RESOURCE_TYPE.get(rtype) + if enum is None: + unknown += 1 + else: + _bump_status(by_type.setdefault(enum, models.NodeStatusCounts()), status) + _set_total(result_counts) + _set_total(auxiliary) + results_by_resource_type = [] + for enum, counts in by_type.items(): + _set_total(counts) + results_by_resource_type.append( + models.ResourceOutcomeStats(resource_type=enum, status_counts=counts) + ) + return result_counts, results_by_resource_type, auxiliary, unknown + + +def _classify_outcome( + exc_type: Optional[type], + has_failures: bool, + fail_fast_triggered: bool, + task_success: Optional[bool], +) -> tuple[models.InvocationStatus, models.TerminationReason]: + if exc_type is not None: + if issubclass(exc_type, (KeyboardInterrupt, SystemExit)): + return models.InvocationStatus.INTERRUPTED, models.TerminationReason.INTERRUPTED + try: + from dbt_common.exceptions import DbtBaseException, DbtInternalError + + if issubclass(exc_type, DbtBaseException) and not issubclass( + exc_type, DbtInternalError + ): + return models.InvocationStatus.HANDLED_ERROR, models.TerminationReason.TASK_ERROR + except Exception: + pass + return models.InvocationStatus.INTERNAL_ERROR, models.TerminationReason.INTERNAL_ERROR + if task_success is False or (task_success is None and has_failures): + reason = ( + models.TerminationReason.FAIL_FAST + if fail_fast_triggered + else models.TerminationReason.NORMAL + ) + return models.InvocationStatus.HANDLED_ERROR, reason + return models.InvocationStatus.SUCCESS, models.TerminationReason.NORMAL + + +def build_post_run_log( + invocation_id: str, + elapsed_ms: int, + exc_type: Optional[type], + results: list, + expected_result_resources: int, + coverage_complete: bool, + results_captured: bool, + selected_resources: Optional[int] = None, + fail_fast_triggered: bool = False, + task_success: Optional[bool] = None, +) -> models.TelemetryLog: + result_counts, by_type, auxiliary, unknown = aggregate_node_results(results) + has_failures = bool( + result_counts.error + or result_counts.fail + or result_counts.runtime_error + or result_counts.partial_success + ) + status, reason = _classify_outcome( + exc_type, + has_failures, + fail_fast_triggered, + task_success, + ) + aggregates_available = bool(results_captured) + return models.TelemetryLog( + invocation_id=invocation_id, + adapter_version=_adapter_version, + dbt_core_version=_dbt_core_version(), + event_type=models.EventType.POST_RUN, + post_run=models.PostRunPayload( + run_outcome=models.RunOutcome( + invocation_status=status, + termination_reason=reason, + invocation_duration_ms=elapsed_ms, + result_aggregates_available=aggregates_available, + expected_result_coverage_complete=( + coverage_complete if aggregates_available else None + ), + ), + selected_resources=selected_resources, + expected_result_resources=expected_result_resources, + result_counts=result_counts if aggregates_available else None, + results_by_resource_type=by_type if aggregates_available else None, + auxiliary_hook_results=auxiliary if aggregates_available else None, + unknown_resource_type_results=unknown if aggregates_available else None, + ), + ) diff --git a/dbt/adapters/databricks/telemetry/client.py b/dbt/adapters/databricks/telemetry/client.py new file mode 100644 index 000000000..80f98ebf5 --- /dev/null +++ b/dbt/adapters/databricks/telemetry/client.py @@ -0,0 +1,61 @@ +from typing import Any, Callable, Optional + +import requests + +from dbt.adapters.databricks.credentials import BearerAuth + +TELEMETRY_AUTHENTICATED_PATH = "/telemetry-ext" +TELEMETRY_UNAUTHENTICATED_PATH = "/telemetry-unauth" + +_TIMEOUT_SECONDS = 10 + +HeaderFactory = Callable[[], dict[str, str]] + + +def _normalize_host(host: str) -> str: + host = host.rstrip("/") + if not host.startswith(("http://", "https://")): + host = f"https://{host}" + return host + + +def send( + host: Optional[str], + body: dict[str, Any], + header_factory: Optional[HeaderFactory] = None, + workspace_id: Optional[int] = None, +) -> bool: + if not host: + return False + + try: + headers = {"Accept": "application/json", "Content-Type": "application/json"} + path = TELEMETRY_AUTHENTICATED_PATH + if header_factory is not None: + auth = BearerAuth(header_factory) + else: + auth = None + path = TELEMETRY_UNAUTHENTICATED_PATH + + if workspace_id is not None: + headers["x-databricks-org-id"] = str(workspace_id) + + url = _normalize_host(host) + path + + response = requests.post( + url, + json=body, + headers=headers, + auth=auth, + timeout=_TIMEOUT_SECONDS, + ) + + if response.status_code // 100 != 2: + return False + try: + ack = response.json() + except ValueError: + return False + return ack.get("numProtoSuccess", 0) >= 1 and not ack.get("errors") + except Exception: + return False diff --git a/dbt/adapters/databricks/telemetry/config.py b/dbt/adapters/databricks/telemetry/config.py new file mode 100644 index 000000000..61c78fdf6 --- /dev/null +++ b/dbt/adapters/databricks/telemetry/config.py @@ -0,0 +1,50 @@ +from typing import Optional + +from dbt.adapters.databricks.credentials import DatabricksCredentials + +ENABLE_FLAG = "enable_dbt_telemetry" +ELIGIBLE_COMMANDS = {"build", "run", "test", "seed", "snapshot"} + + +def is_enabled(credentials: Optional[DatabricksCredentials]) -> bool: + if credentials is None: + return False + params = credentials.connection_parameters or {} + value = params.get(ENABLE_FLAG, False) + if isinstance(value, bool): + return value + if isinstance(value, int): + return value == 1 + if isinstance(value, str): + return value.strip().lower() in {"1", "true", "yes", "on"} + return False + + +def is_eligible_command() -> bool: + try: + from dbt.flags import get_flags + + which = getattr(get_flags(), "WHICH", None) + command = str(which or "").strip().lower().replace("_", "-").split()[0] + return command in ELIGIBLE_COMMANDS + except Exception: + return False + + +def is_enabled_for_invocation(credentials: Optional[DatabricksCredentials]) -> bool: + return is_enabled(credentials) and is_eligible_command() + + +def has_reusable_transport(credentials: Optional[DatabricksCredentials]) -> bool: + """Kernel OAuth U2M credentials are not reusable.""" + if credentials is None: + return False + params = credentials.connection_parameters or {} + kernel_u2m = ( + bool(params.get("use_kernel")) + and credentials.auth_type == "oauth" + and not credentials.token + and not credentials.client_secret + and not credentials.azure_client_secret + ) + return not kernel_u2m diff --git a/dbt/adapters/databricks/telemetry/coordinator.py b/dbt/adapters/databricks/telemetry/coordinator.py new file mode 100644 index 000000000..b453594a9 --- /dev/null +++ b/dbt/adapters/databricks/telemetry/coordinator.py @@ -0,0 +1,395 @@ +import threading +import time +import uuid +from collections import Counter +from typing import Any, Callable, Optional + +from dbt.adapters.databricks.telemetry import client, encoder +from dbt.adapters.databricks.telemetry.models import TelemetryLog + +HeaderFactory = Callable[[], dict[str, str]] + + +class Transport: + """Reusable transport with a redacted representation.""" + + __slots__ = ("host", "header_factory", "workspace_id") + + def __init__( + self, + host: Optional[str], + header_factory: Optional[HeaderFactory], + workspace_id: Optional[Any], + ) -> None: + self.host = host + self.header_factory = header_factory + self.workspace_id = workspace_id + + def __repr__(self) -> str: # pragma: no cover - defensive redaction + return "" + + +class _InvocationState: + __slots__ = ( + "post_parse", + "post_run", + "transport", + "post_parse_event_id", + "post_run_event_id", + "post_parse_sent", + "post_parse_terminal", + "post_run_sent", + "sending", + "start_monotonic", + "node_results", + "hook_results", + "end_run_statuses", + "end_run_success", + "fail_fast_triggered", + "expected_result_count", + "ephemeral_ids", + "results_captured", + "closed", + ) + + def __init__(self) -> None: + self.post_parse: Optional[TelemetryLog] = None + self.post_run: Optional[TelemetryLog] = None + self.transport: Optional[Transport] = None + self.post_parse_event_id: str = str(uuid.uuid4()) + self.post_run_event_id: str = str(uuid.uuid4()) + self.post_parse_sent: bool = False + self.post_parse_terminal: bool = False + self.post_run_sent: bool = False + self.sending: bool = False + self.start_monotonic: float = time.monotonic() + self.node_results: list = [] + self.hook_results: list = [] + self.end_run_statuses: Optional[list] = None + self.end_run_success: Optional[bool] = None + self.fail_fast_triggered: bool = False + self.expected_result_count: int = 0 + self.ephemeral_ids: set[str] = set() + self.results_captured: bool = False + self.closed: bool = False + + +class Coordinator: + def __init__(self) -> None: + self._lock = threading.Lock() + self._states: dict[str, _InvocationState] = {} + self._threads_lock = threading.Lock() + self._threads: list[threading.Thread] = [] + + def _state(self, invocation_id: str) -> Optional[_InvocationState]: + state = self._states.get(invocation_id) + if state is None: + state = _InvocationState() + self._states[invocation_id] = state + elif state.closed: + return None + return state + + def set_post_parse(self, invocation_id: str, payload: TelemetryLog) -> None: + with self._lock: + state = self._state(invocation_id) + if state is None: + return + if state.post_parse is None: + state.post_parse = payload + self.send_if_ready(invocation_id) + + def set_post_run(self, invocation_id: str, payload: TelemetryLog) -> None: + with self._lock: + state = self._state(invocation_id) + if state is None: + return + if state.post_run is None: + state.post_run = payload + self.send_if_ready(invocation_id) + + def set_transport(self, invocation_id: str, transport: Transport) -> None: + with self._lock: + state = self._state(invocation_id) + if state is None: + return + if state.transport is None: + state.transport = transport + self.send_if_ready(invocation_id) + + def mark_start(self, invocation_id: str) -> None: + with self._lock: + self._state(invocation_id) + + def is_closed(self, invocation_id: str) -> bool: + with self._lock: + state = self._states.get(invocation_id) + return state is not None and state.closed + + def needs_post_parse(self, invocation_id: str) -> bool: + with self._lock: + state = self._states.get(invocation_id) + return state is None or (not state.closed and state.post_parse is None) + + def record_node_result(self, invocation_id: str, unique_id: str, status: str) -> None: + with self._lock: + state = self._states.get(invocation_id) + if state is None or state.closed: + return + # Error events may precede NodeFinished for the same node. + for index, (observed_id, _) in enumerate(state.node_results): + if str(observed_id) == str(unique_id): + state.node_results[index] = (unique_id, status) + break + else: + state.node_results.append((unique_id, status)) + state.results_captured = True + + def record_hook_result(self, invocation_id: str, status: str) -> None: + with self._lock: + state = self._states.get(invocation_id) + if state is None or state.closed: + return + state.hook_results.append(("operation", status)) + state.results_captured = True + + def record_end_run( + self, invocation_id: str, statuses: list, success: Optional[bool] = None + ) -> None: + with self._lock: + state = self._states.get(invocation_id) + if state is None or state.closed: + return + state.end_run_statuses = list(statuses) + state.end_run_success = success + state.results_captured = True + + def mark_fail_fast_triggered(self, invocation_id: str) -> None: + with self._lock: + state = self._states.get(invocation_id) + if state is None or state.closed or state.end_run_statuses is not None: + return + state.fail_fast_triggered = True + + def outcome_snapshot(self, invocation_id: str) -> tuple[Optional[bool], bool]: + with self._lock: + state = self._states.get(invocation_id) + if state is None: + return None, False + return state.end_run_success, state.fail_fast_triggered + + def record_expected_count(self, invocation_id: str, count: int) -> None: + with self._lock: + state = self._states.get(invocation_id) + if state is None or state.closed: + return + state.expected_result_count = count + + def record_ephemeral_ids(self, invocation_id: str, unique_ids: set[str]) -> None: + with self._lock: + state = self._state(invocation_id) + if state is not None: + state.ephemeral_ids = set(unique_ids) + + def result_snapshot(self, invocation_id: str) -> tuple: + with self._lock: + state = self._states.get(invocation_id) + if state is None: + return [], 0, 0, False, False + observed_ids = {str(uid) for uid, _ in state.node_results} + observed_ephemeral_count = len(observed_ids.intersection(state.ephemeral_ids)) + top_level_results = [ + result for result in state.node_results if str(result[0]) not in state.ephemeral_ids + ] + if state.end_run_statuses is not None: + # EndRunResult statuses have no unique IDs. + remaining = Counter(_status_key(status) for status in state.end_run_statuses) + typed_results = top_level_results + state.hook_results + for _, status in typed_results: + key = _status_key(status) + if remaining[key] > 0: + remaining[key] -= 1 + remaining_statuses = [ + status for status in state.end_run_statuses if _take_status(remaining, status) + ] + missing_expected = max(state.expected_result_count - len(top_level_results), 0) + synthesized = [(None, status) for status in remaining_statuses[:missing_expected]] + synthetic_ephemeral_count = max(len(remaining_statuses) - missing_expected, 0) + selected_count = ( + state.expected_result_count + + observed_ephemeral_count + + synthetic_ephemeral_count + ) + reconciled_results = top_level_results + synthesized + state.hook_results + return ( + reconciled_results, + selected_count, + state.expected_result_count, + len(top_level_results) + len(synthesized) >= state.expected_result_count, + True, + ) + # Missing ephemeral results make selection unknown. + partial_selected_count: Optional[int] + if not state.ephemeral_ids: + partial_selected_count = state.expected_result_count + elif observed_ephemeral_count == len(state.ephemeral_ids): + partial_selected_count = state.expected_result_count + observed_ephemeral_count + else: + partial_selected_count = None + return ( + top_level_results + state.hook_results, + partial_selected_count, + state.expected_result_count, + False, + False, + ) + + def elapsed_ms(self, invocation_id: str) -> int: + with self._lock: + state = self._states.get(invocation_id) + if state is None: + return 0 + elapsed_seconds = time.monotonic() - state.start_monotonic + return int(max(elapsed_seconds, 0.0) * 1000) + + def send_if_ready(self, invocation_id: str) -> None: + # Keep connection-open off the network path. + with self._lock: + if not self._ready_to_send(self._states.get(invocation_id)): + return + thread = threading.Thread( + target=self._drain, + args=(invocation_id,), + name="dbt-telemetry-send", + daemon=True, + ) + with self._threads_lock: + self._threads = [t for t in self._threads if t.is_alive()] + self._threads.append(thread) + thread.start() + + def _ready_to_send(self, state: Optional[_InvocationState]) -> bool: + if state is None or state.closed or state.sending: + return False + transport = state.transport + if transport is None or transport.header_factory is None: + return False + if state.post_parse is not None and not state.post_parse_sent: + return True + return ( + state.post_run is not None + and not state.post_run_sent + and (state.post_parse_terminal or state.post_parse is None) + ) + + def flush(self, timeout: Optional[float] = None) -> None: + if timeout is None: + timeout = float(client._TIMEOUT_SECONDS * 2) + deadline = time.monotonic() + timeout + while True: + with self._threads_lock: + self._threads = [t for t in self._threads if t.is_alive()] + pending = list(self._threads) + if not pending: + return + remaining = deadline - time.monotonic() + if remaining <= 0: + return + pending[0].join(remaining) + + def _drain(self, invocation_id: str) -> None: + # Serialize phases without holding the lock during sends. + while True: + with self._lock: + state = self._states.get(invocation_id) + if state is None or state.closed or state.sending: + return + transport = state.transport + if transport is None or transport.header_factory is None: + return + + phase = None + payload = None + event_id = None + if state.post_parse is not None and not state.post_parse_sent: + phase = "post_parse" + payload = state.post_parse + event_id = state.post_parse_event_id + state.post_parse_sent = True + elif ( + state.post_run is not None + and not state.post_run_sent + and (state.post_parse_terminal or state.post_parse is None) + ): + phase = "post_run" + payload = state.post_run + event_id = state.post_run_event_id + state.post_run_sent = True + if payload is None or event_id is None: + return + + state.sending = True + host = transport.host + header_factory = transport.header_factory + workspace_id = transport.workspace_id + + self._send(host, payload, event_id, header_factory, workspace_id) + + with self._lock: + state = self._states.get(invocation_id) + if state is None or state.closed: + return + state.sending = False + if phase == "post_parse": + state.post_parse_terminal = True + + def _send( + self, + host: Optional[str], + payload: TelemetryLog, + event_id: str, + header_factory: Optional[HeaderFactory], + workspace_id: Optional[Any], + ) -> None: + try: + body = encoder.encode_request(payload, event_id, workspace_id=workspace_id) + client.send(host, body, header_factory=header_factory, workspace_id=workspace_id) + except Exception: # pragma: no cover - best-effort + return + + def close(self, invocation_id: str) -> None: + with self._lock: + # Reject late callbacks and clear sensitive state. + state = self._states.get(invocation_id) + if state is None: + state = _InvocationState() + self._states[invocation_id] = state + state.closed = True + state.post_parse = None + state.post_run = None + state.transport = None + state.node_results.clear() + state.hook_results.clear() + state.end_run_statuses = None + state.end_run_success = None + state.fail_fast_triggered = False + state.ephemeral_ids.clear() + + +def _status_key(status: Any) -> str: + return str(status).strip().lower().replace("-", "_").replace(" ", "_") + + +def _take_status(remaining: Counter, status: Any) -> bool: + key = _status_key(status) + if remaining[key] <= 0: + return False + remaining[key] -= 1 + return True + + +_COORDINATOR = Coordinator() + + +def coordinator() -> Coordinator: + return _COORDINATOR diff --git a/dbt/adapters/databricks/telemetry/encoder.py b/dbt/adapters/databricks/telemetry/encoder.py new file mode 100644 index 000000000..8b45e1687 --- /dev/null +++ b/dbt/adapters/databricks/telemetry/encoder.py @@ -0,0 +1,62 @@ +import dataclasses +import json +import time +from enum import Enum +from typing import Any, Optional + +from dbt.adapters.databricks.telemetry.models import TelemetryLog + +DRIVER_NAME = "dbt-databricks" + + +def _proto_dict(log: TelemetryLog) -> dict[str, Any]: + def factory(items: list) -> dict: + out = {} + for k, v in items: + if v is None: + continue + out[k[:-1] if k.endswith("_") else k] = v.value if isinstance(v, Enum) else v + return out + + return dataclasses.asdict(log, dict_factory=factory) + + +def _coerce_workspace_id(workspace_id: Optional[Any]) -> Optional[int]: + try: + return int(workspace_id) if workspace_id is not None else None + except (TypeError, ValueError): + return None + + +def encode_frontend_log( + log: TelemetryLog, + frontend_log_event_id: str, + workspace_id: Optional[Any] = None, +) -> str: + entry = {"dbt_databricks_telemetry_log": _proto_dict(log)} + frontend_log: dict[str, Any] = { + "frontend_log_event_id": frontend_log_event_id, + "context": { + "client_context": { + "timestamp_millis": int(time.time() * 1000), + "user_agent": f"{DRIVER_NAME}/{log.adapter_version}", + } + }, + "entry": entry, + } + coerced = _coerce_workspace_id(workspace_id) + if coerced is not None: + frontend_log["workspace_id"] = coerced + return json.dumps(frontend_log) + + +def encode_request( + log: TelemetryLog, + frontend_log_event_id: str, + workspace_id: Optional[Any] = None, +) -> dict[str, Any]: + return { + "uploadTime": int(time.time() * 1000), + "items": [], + "protoLogs": [encode_frontend_log(log, frontend_log_event_id, workspace_id)], + } diff --git a/dbt/adapters/databricks/telemetry/hooks.py b/dbt/adapters/databricks/telemetry/hooks.py new file mode 100644 index 000000000..a4b28d351 --- /dev/null +++ b/dbt/adapters/databricks/telemetry/hooks.py @@ -0,0 +1,132 @@ +import sys +from typing import Any, Optional + +from dbt.adapters.databricks.credentials import DatabricksCredentials +from dbt.adapters.databricks.telemetry import builder, listener +from dbt.adapters.databricks.telemetry.config import ( + has_reusable_transport, + is_enabled_for_invocation, +) +from dbt.adapters.databricks.telemetry.coordinator import Transport, coordinator + + +def _current_invocation_id() -> Optional[str]: + try: + from dbt_common.invocation import get_invocation_id + + invocation_id = get_invocation_id() + return str(invocation_id) if invocation_id else None + except Exception: + return None + + +def on_adapter_init(adapter: Any) -> None: + try: + creds = getattr(getattr(adapter, "config", None), "credentials", None) + if not isinstance(creds, DatabricksCredentials) or not is_enabled_for_invocation(creds): + return + invocation_id = _current_invocation_id() + if not invocation_id: + return + coord = coordinator() + coord.mark_start(invocation_id) + if not listener.register(): + coord.close(invocation_id) + except Exception: # pragma: no cover - best-effort + return + + +def on_post_parse(adapter: Any, manifest: Any) -> None: + try: + config = getattr(adapter, "config", None) + creds = getattr(config, "credentials", None) + if not isinstance(creds, DatabricksCredentials) or not is_enabled_for_invocation(creds): + return + invocation_id = _current_invocation_id() + coord = coordinator() + if invocation_id and not coord.needs_post_parse(invocation_id): + return + log = builder.build_post_parse_log( + manifest=manifest, + config=config, + creds=creds, + behavior_flag=adapter.get_behavior_flag_no_warn, + ) + if not log.invocation_id: + return + coord.record_ephemeral_ids(log.invocation_id, builder.ephemeral_resource_ids(manifest)) + coord.set_post_parse(log.invocation_id, log) + except Exception: # pragma: no cover - best-effort + return + + +def on_connection_open( + credentials: Optional[DatabricksCredentials], + credentials_manager: Optional[Any], +) -> None: + try: + if ( + not is_enabled_for_invocation(credentials) + or not has_reusable_transport(credentials) + or credentials_manager is None + ): + return + invocation_id = _current_invocation_id() + if not invocation_id: + return + transport = Transport( + host=getattr(credentials_manager, "host", None), + header_factory=getattr(credentials_manager, "header_factory", None), + workspace_id=getattr(credentials_manager, "workspace_id", None), + ) + coordinator().set_transport(invocation_id, transport) + except Exception: # pragma: no cover - best-effort + return + + +def _finalize_post_run(invocation_id: str, exc_type: Optional[type]) -> None: + coord = coordinator() + if coord.is_closed(invocation_id): + return + results, selected, expected, coverage_complete, results_captured = coord.result_snapshot( + invocation_id + ) + task_success, fail_fast_triggered = coord.outcome_snapshot(invocation_id) + log = builder.build_post_run_log( + invocation_id, + coord.elapsed_ms(invocation_id), + exc_type, + results, + expected, + coverage_complete, + results_captured, + selected_resources=selected, + fail_fast_triggered=fail_fast_triggered, + task_success=task_success, + ) + coord.set_post_run(invocation_id, log) + coord.flush() + coord.close(invocation_id) + + +def on_end_run_result(invocation_id: str) -> None: + try: + _finalize_post_run(invocation_id, None) + except Exception: # pragma: no cover - best-effort + return + + +def on_run_end(adapter: Any) -> None: + try: + config = getattr(adapter, "config", None) + creds = getattr(config, "credentials", None) + if not isinstance(creds, DatabricksCredentials) or not is_enabled_for_invocation(creds): + return + invocation_id = _current_invocation_id() + if not invocation_id: + return + exc_type = sys.exc_info()[0] + if exc_type is not None: + _finalize_post_run(invocation_id, exc_type) + except Exception: # pragma: no cover - best-effort + return diff --git a/dbt/adapters/databricks/telemetry/listener.py b/dbt/adapters/databricks/telemetry/listener.py new file mode 100644 index 000000000..7c9fed265 --- /dev/null +++ b/dbt/adapters/databricks/telemetry/listener.py @@ -0,0 +1,81 @@ +from typing import Any + +from dbt.adapters.databricks.telemetry.coordinator import coordinator + + +def _current_invocation_id() -> str: + from dbt_common.invocation import get_invocation_id + + return str(get_invocation_id() or "") + + +def _fail_fast_enabled() -> bool: + try: + from dbt.flags import get_flags + + return bool(getattr(get_flags(), "FAIL_FAST", False)) + except Exception: + return False + + +def _on_event(msg: Any) -> None: + try: + name = msg.info.name + if name not in ( + "NodeFinished", + "EndRunResult", + "LogHookEndLine", + "ConcurrencyLine", + "GenericExceptionOnRun", + "SkippingDetails", + "RunResultFailure", + "RunResultError", + ): + return + invocation_id = _current_invocation_id() + if not invocation_id: + return + coord = coordinator() + if name == "EndRunResult": + coord.record_end_run( + invocation_id, + [r.status for r in msg.data.results], + success=getattr(msg.data, "success", None), + ) + from dbt.adapters.databricks.telemetry import hooks + + hooks.on_end_run_result(invocation_id) + elif name == "NodeFinished": + info = msg.data.node_info + coord.record_node_result(invocation_id, info.unique_id, info.node_status) + elif name == "LogHookEndLine": + coord.record_hook_result(invocation_id, msg.data.status) + elif name == "ConcurrencyLine": + coord.record_expected_count(invocation_id, msg.data.node_count) + elif name == "GenericExceptionOnRun": + # No NodeFinished event follows. + unique_id = msg.data.unique_id or msg.data.node_info.unique_id + if unique_id: + coord.record_node_result(invocation_id, unique_id, "error") + elif name == "SkippingDetails": + unique_id = msg.data.node_info.unique_id + if unique_id: + coord.record_node_result(invocation_id, unique_id, "skipped") + elif name in ("RunResultFailure", "RunResultError") and _fail_fast_enabled(): + coord.mark_fail_fast_triggered(invocation_id) + except Exception: # pragma: no cover - best-effort + return + + +def register() -> bool: + try: + from dbt_common.events.event_manager_client import ( + add_callback_to_manager, + get_event_manager, + ) + + if _on_event not in get_event_manager().callbacks: + add_callback_to_manager(_on_event) + return True + except Exception: # pragma: no cover - best-effort + return False diff --git a/dbt/adapters/databricks/telemetry/models.py b/dbt/adapters/databricks/telemetry/models.py new file mode 100644 index 000000000..f84549c00 --- /dev/null +++ b/dbt/adapters/databricks/telemetry/models.py @@ -0,0 +1,192 @@ +from dataclasses import dataclass, field +from enum import Enum +from typing import Optional + + +class EventType(Enum): + TYPE_UNSPECIFIED = "TYPE_UNSPECIFIED" + POST_PARSE = "POST_PARSE" + POST_RUN = "POST_RUN" + + +class ComputeType(Enum): + TYPE_UNSPECIFIED = "TYPE_UNSPECIFIED" + SQL_WAREHOUSE = "SQL_WAREHOUSE" + ALL_PURPOSE_CLUSTER = "ALL_PURPOSE_CLUSTER" + OTHER = "OTHER" + + +class AuthFamily(Enum): + TYPE_UNSPECIFIED = "TYPE_UNSPECIFIED" + PAT = "PAT" + OAUTH_U2M = "OAUTH_U2M" + OAUTH_M2M = "OAUTH_M2M" + AZURE_SERVICE_PRINCIPAL = "AZURE_SERVICE_PRINCIPAL" + LEGACY_CLIENT_SECRET_AMBIGUOUS = "LEGACY_CLIENT_SECRET_AMBIGUOUS" + OTHER = "OTHER" + + +class DbtCommand(Enum): + TYPE_UNSPECIFIED = "TYPE_UNSPECIFIED" + RUN = "RUN" + BUILD = "BUILD" + TEST = "TEST" + SEED = "SEED" + SNAPSHOT = "SNAPSHOT" + COMPILE = "COMPILE" + DOCS = "DOCS" + CLONE = "CLONE" + RETRY = "RETRY" + SHOW = "SHOW" + LIST = "LIST" + SOURCE = "SOURCE" + RUN_OPERATION = "RUN_OPERATION" + OTHER = "OTHER" + + +class WarnErrorPolicy(Enum): + TYPE_UNSPECIFIED = "TYPE_UNSPECIFIED" + WARN_ERROR_DISABLED = "WARN_ERROR_DISABLED" + WARN_ERROR_ALL = "WARN_ERROR_ALL" + WARN_ERROR_CUSTOM_POLICY = "WARN_ERROR_CUSTOM_POLICY" + + +class ResourceType(Enum): + TYPE_UNSPECIFIED = "TYPE_UNSPECIFIED" + MODEL = "MODEL" + DATA_TEST = "DATA_TEST" + UNIT_TEST = "UNIT_TEST" + SEED = "SEED" + SNAPSHOT = "SNAPSHOT" + SOURCE = "SOURCE" + FUNCTION = "FUNCTION" + EXPOSURE = "EXPOSURE" + SAVED_QUERY = "SAVED_QUERY" + OTHER = "OTHER" + + +class InvocationStatus(Enum): + TYPE_UNSPECIFIED = "TYPE_UNSPECIFIED" + SUCCESS = "SUCCESS" + HANDLED_ERROR = "HANDLED_ERROR" + INTERRUPTED = "INTERRUPTED" + INTERNAL_ERROR = "INTERNAL_ERROR" + + +class TerminationReason(Enum): + TYPE_UNSPECIFIED = "TYPE_UNSPECIFIED" + NORMAL = "NORMAL" + FAIL_FAST = "FAIL_FAST" + INTERRUPTED = "INTERRUPTED" + TASK_ERROR = "TASK_ERROR" + INTERNAL_ERROR = "INTERNAL_ERROR" + + +@dataclass +class ResourceCounts: + model_count: int = 0 + data_test_count: int = 0 + generic_data_test_count: int = 0 + seed_count: int = 0 + snapshot_count: int = 0 + source_count: int = 0 + function_count: int = 0 + exposure_count: int = 0 + saved_query_count: int = 0 + other_count: int = 0 + unit_test_count: int = 0 + + +@dataclass +class ManifestStats: + enabled_total: ResourceCounts = field(default_factory=ResourceCounts) + enabled_root_project: ResourceCounts = field(default_factory=ResourceCounts) + enabled_installed_packages: ResourceCounts = field(default_factory=ResourceCounts) + + +@dataclass +class InvocationConfig: + thread_count: int = 0 + dbt_command: DbtCommand = DbtCommand.TYPE_UNSPECIFIED + full_refresh: bool = False + empty: bool = False + fail_fast: bool = False + warn_error_policy: WarnErrorPolicy = WarnErrorPolicy.WARN_ERROR_DISABLED + + +@dataclass +class ConnectionConfig: + default_compute_type: ComputeType = ComputeType.TYPE_UNSPECIFIED + configured_auth_family: AuthFamily = AuthFamily.TYPE_UNSPECIFIED + named_compute_count: int = 0 + spog_routing_configured: bool = False + use_kernel: bool = False + + +@dataclass +class ProjectConfig: + use_user_folder_for_python: bool = False + use_materialization_v2: bool = False + use_replace_on_for_insert_overwrite: bool = False + use_managed_iceberg: bool = False + use_concurrent_microbatch: bool = False + use_describe_as_json_for_relation_metadata: bool = False + + +@dataclass +class PostParsePayload: + invocation_config: InvocationConfig + manifest_stats: ManifestStats + connection_config: ConnectionConfig + project_config: ProjectConfig + + +@dataclass +class NodeStatusCounts: + total: int = 0 + success: int = 0 + error: int = 0 + fail: int = 0 + warn: int = 0 + skipped: int = 0 + partial_success: int = 0 + pass_: int = 0 # proto field: pass + runtime_error: int = 0 + no_op: int = 0 + reused: int = 0 + + +@dataclass +class ResourceOutcomeStats: + resource_type: ResourceType = ResourceType.TYPE_UNSPECIFIED + status_counts: NodeStatusCounts = field(default_factory=NodeStatusCounts) + + +@dataclass +class RunOutcome: + invocation_status: InvocationStatus = InvocationStatus.TYPE_UNSPECIFIED + termination_reason: TerminationReason = TerminationReason.TYPE_UNSPECIFIED + invocation_duration_ms: int = 0 + result_aggregates_available: bool = False + expected_result_coverage_complete: Optional[bool] = None + + +@dataclass +class PostRunPayload: + run_outcome: RunOutcome = field(default_factory=RunOutcome) + selected_resources: Optional[int] = None + expected_result_resources: int = 0 + result_counts: Optional[NodeStatusCounts] = None + results_by_resource_type: Optional[list[ResourceOutcomeStats]] = None + auxiliary_hook_results: Optional[NodeStatusCounts] = None + unknown_resource_type_results: Optional[int] = None + + +@dataclass +class TelemetryLog: + invocation_id: str + adapter_version: str + dbt_core_version: str + event_type: EventType = EventType.POST_PARSE + post_parse: Optional[PostParsePayload] = None + post_run: Optional[PostRunPayload] = None diff --git a/tests/unit/telemetry/__init__.py b/tests/unit/telemetry/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/unit/telemetry/test_builder.py b/tests/unit/telemetry/test_builder.py new file mode 100644 index 000000000..eaa8eaa37 --- /dev/null +++ b/tests/unit/telemetry/test_builder.py @@ -0,0 +1,319 @@ +from types import SimpleNamespace + +from dbt.adapters.databricks.telemetry import builder, models + + +def _creds(**kw): + base = dict( + token=None, + client_id=None, + client_secret=None, + azure_client_id=None, + azure_client_secret=None, + auth_type=None, + http_path="/sql/1.0/warehouses/x", + compute=None, + connection_parameters=None, + ) + base.update(kw) + return SimpleNamespace(**base) + + +def _node(resource_type, package_name="root", test_metadata=None): + return SimpleNamespace( + resource_type=resource_type, package_name=package_name, test_metadata=test_metadata + ) + + +class TestClassifyComputeType: + def test_sql_warehouse(self): + assert builder.classify_compute_type("/sql/1.0/warehouses/a") == ( + models.ComputeType.SQL_WAREHOUSE + ) + + def test_endpoints_form_is_warehouse(self): + assert builder.classify_compute_type("/sql/1.0/endpoints/a") == ( + models.ComputeType.SQL_WAREHOUSE + ) + + def test_all_purpose_cluster(self): + assert builder.classify_compute_type("/sql/protocolv1/o/1/2") == ( + models.ComputeType.ALL_PURPOSE_CLUSTER + ) + + def test_query_string_ignored(self): + assert builder.classify_compute_type("/sql/1.0/warehouses/a?o=9") == ( + models.ComputeType.SQL_WAREHOUSE + ) + + def test_unknown_form_is_other(self): + assert builder.classify_compute_type("/unknown") == models.ComputeType.OTHER + + def test_missing_is_unspecified(self): + assert builder.classify_compute_type(None) == models.ComputeType.TYPE_UNSPECIFIED + + +class TestClassifyAuthFamily: + def test_token_is_pat(self): + assert builder.classify_auth_family(_creds(token="dapi")) == models.AuthFamily.PAT + + def test_azure_service_principal(self): + assert ( + builder.classify_auth_family(_creds(azure_client_id="a", azure_client_secret="b")) + == models.AuthFamily.AZURE_SERVICE_PRINCIPAL + ) + + def test_no_secret_is_u2m(self): + assert builder.classify_auth_family(_creds(auth_type="oauth")) == ( + models.AuthFamily.OAUTH_U2M + ) + + def test_client_secret_is_ambiguous(self): + assert builder.classify_auth_family(_creds(client_id="c", client_secret="s")) == ( + models.AuthFamily.LEGACY_CLIENT_SECRET_AMBIGUOUS + ) + + +class TestClassifyCommand: + def test_known(self): + assert builder.classify_command("run") == models.DbtCommand.RUN + assert builder.classify_command("build") == models.DbtCommand.BUILD + + def test_normalized_forms(self): + assert builder.classify_command("run-operation") == models.DbtCommand.RUN_OPERATION + assert builder.classify_command("source freshness") == models.DbtCommand.SOURCE + assert builder.classify_command("docs generate") == models.DbtCommand.DOCS + + def test_unknown_is_other(self): + assert builder.classify_command("parse") == models.DbtCommand.OTHER + + def test_missing_is_unspecified(self): + assert builder.classify_command(None) == models.DbtCommand.TYPE_UNSPECIFIED + + +class TestClassifyWarnErrorPolicy: + def test_disabled(self): + assert builder.classify_warn_error_policy(None, None) == ( + models.WarnErrorPolicy.WARN_ERROR_DISABLED + ) + + def test_all(self): + assert builder.classify_warn_error_policy(True, None) == ( + models.WarnErrorPolicy.WARN_ERROR_ALL + ) + + def test_error_all_without_overrides_is_all(self): + options = SimpleNamespace(error="all", warn=[], silence=[]) + assert builder.classify_warn_error_policy(False, options) == ( + models.WarnErrorPolicy.WARN_ERROR_ALL + ) + + def test_error_all_with_named_override_is_custom(self): + options = SimpleNamespace(error="all", warn=["SomeWarning"], silence=[]) + assert builder.classify_warn_error_policy(False, options) == ( + models.WarnErrorPolicy.WARN_ERROR_CUSTOM_POLICY + ) + + def test_legacy_warn_error_takes_precedence(self): + options = SimpleNamespace(error=[], warn=[], silence=["SomeWarning"]) + assert builder.classify_warn_error_policy(True, options) == ( + models.WarnErrorPolicy.WARN_ERROR_ALL + ) + + def test_custom_policy(self): + assert builder.classify_warn_error_policy(False, {"include": ["X"]}) == ( + models.WarnErrorPolicy.WARN_ERROR_CUSTOM_POLICY + ) + + +class TestBuildConnectionConfig: + def test_derives_all_fields(self): + cc = builder.build_connection_config( + _creds( + client_id="c", + client_secret="s", + http_path="/sql/protocolv1/o/1/2?o=42", + compute={"a": {}, "b": {}}, + connection_parameters={"use_kernel": True}, + ) + ) + assert cc.default_compute_type == models.ComputeType.ALL_PURPOSE_CLUSTER + assert cc.configured_auth_family == models.AuthFamily.LEGACY_CLIENT_SECRET_AMBIGUOUS + assert cc.named_compute_count == 2 + assert cc.spog_routing_configured is True + assert cc.use_kernel is True + + def test_spog_parameter_is_parsed_not_substring_matched(self): + cc = builder.build_connection_config( + _creds(token="dapi", http_path="/sql/1.0/warehouses/w?foo=x&o=42") + ) + assert cc.spog_routing_configured is True + + cc = builder.build_connection_config( + _creds(token="dapi", http_path="/sql/1.0/warehouses/w?foo=?o=42") + ) + assert cc.spog_routing_configured is False + + +class TestAggregateManifest: + def _manifest(self): + return SimpleNamespace( + metadata=SimpleNamespace(project_name="root", invocation_id="inv-1"), + nodes={ + "m1": _node("model", "root"), + "m2": _node("model", "dep_pkg"), + "t_generic": _node("test", "root", test_metadata={"name": "not_null"}), + "t_singular": _node("test", "root"), + "sd": _node("seed", "root"), + "sn": _node("snapshot", "root"), + "op": _node("operation", "root"), + }, + sources={"s1": _node("source", "root")}, + exposures={"e1": _node("exposure", "root")}, + metrics={"me1": _node("metric", "root")}, + saved_queries={"sq1": _node("saved_query", "root")}, + functions={}, + semantic_models={"sem1": _node("semantic_model", "root")}, + unit_tests={"ut2": _node("unit_test", "root")}, + ) + + def test_total_counts(self): + ms = builder.aggregate_manifest(self._manifest()) + assert ms.enabled_total.model_count == 2 + assert ms.enabled_total.data_test_count == 2 + assert ms.enabled_total.generic_data_test_count == 1 + assert ms.enabled_total.seed_count == 1 + assert ms.enabled_total.snapshot_count == 1 + assert ms.enabled_total.source_count == 1 + assert ms.enabled_total.exposure_count == 1 + assert ms.enabled_total.saved_query_count == 1 + assert ms.enabled_total.other_count == 3 + assert ms.enabled_total.unit_test_count == 1 + + def test_root_vs_installed_split(self): + ms = builder.aggregate_manifest(self._manifest()) + assert ms.enabled_root_project.model_count == 1 + assert ms.enabled_installed_packages.model_count == 1 + + +class TestBuildPostRunLog: + def test_success_when_no_exception(self): + log = builder.build_post_run_log("inv", 250, None, [], 0, True, True) + assert log.event_type == models.EventType.POST_RUN + outcome = log.post_run.run_outcome + assert outcome.invocation_status == models.InvocationStatus.SUCCESS + assert outcome.termination_reason == models.TerminationReason.NORMAL + assert outcome.invocation_duration_ms == 250 + assert outcome.result_aggregates_available is True + + def test_handled_error_when_failures(self): + outcome = builder.build_post_run_log( + "inv", 1, None, [("model.p.m1", "error")], 1, True, True + ).post_run.run_outcome + assert outcome.invocation_status == models.InvocationStatus.HANDLED_ERROR + assert outcome.termination_reason == models.TerminationReason.NORMAL + + def test_fail_fast_termination_when_observed(self): + outcome = builder.build_post_run_log( + "inv", + 1, + None, + [("model.p.m1", "error"), ("model.p.m2", "skipped")], + 2, + True, + True, + selected_resources=2, + fail_fast_triggered=True, + task_success=False, + ).post_run.run_outcome + assert outcome.invocation_status == models.InvocationStatus.HANDLED_ERROR + assert outcome.termination_reason == models.TerminationReason.FAIL_FAST + + def test_keyboard_interrupt(self): + outcome = builder.build_post_run_log( + "inv", 1, KeyboardInterrupt, [], 0, False, True + ).post_run.run_outcome + assert outcome.invocation_status == models.InvocationStatus.INTERRUPTED + assert outcome.termination_reason == models.TerminationReason.INTERRUPTED + + def test_authoritative_task_failure_includes_auxiliary_failures(self): + outcome = builder.build_post_run_log( + "inv", + 1, + None, + [("operation.p.h", "error"), ("model.p.m", "skipped")], + 1, + True, + True, + task_success=False, + ).post_run.run_outcome + assert outcome.invocation_status == models.InvocationStatus.HANDLED_ERROR + + def test_other_exception_is_internal_error(self): + outcome = builder.build_post_run_log( + "inv", 1, RuntimeError, [], 0, False, True + ).post_run.run_outcome + assert outcome.invocation_status == models.InvocationStatus.INTERNAL_ERROR + assert outcome.termination_reason == models.TerminationReason.INTERNAL_ERROR + + def test_aggregates_unavailable_when_not_captured(self): + post_run = builder.build_post_run_log("inv", 1, None, [], 0, False, False).post_run + outcome = post_run.run_outcome + assert outcome.result_aggregates_available is False + assert outcome.expected_result_coverage_complete is None + assert post_run.result_counts is None + assert post_run.results_by_resource_type is None + assert post_run.auxiliary_hook_results is None + assert post_run.unknown_resource_type_results is None + + +class TestAggregateNodeResults: + def test_result_counts_and_total(self): + rc, _, _, unknown = builder.aggregate_node_results( + [("model.p.a", "success"), ("model.p.b", "error"), ("test.p.t", "pass")] + ) + assert rc.total == 3 + assert rc.success == 1 and rc.error == 1 and rc.pass_ == 1 + assert unknown == 0 + + def test_results_by_resource_type(self): + _, by_type, _, _ = builder.aggregate_node_results( + [("model.p.a", "success"), ("model.p.b", "success"), ("test.p.t", "pass")] + ) + by = {r.resource_type: r.status_counts for r in by_type} + assert by[models.ResourceType.MODEL].success == 2 + assert by[models.ResourceType.MODEL].total == 2 + assert by[models.ResourceType.DATA_TEST].pass_ == 1 + + def test_operations_are_auxiliary(self): + rc, by_type, aux, _ = builder.aggregate_node_results( + [("operation.p.hook", "success"), ("model.p.a", "success")] + ) + assert aux.total == 1 and aux.success == 1 + assert rc.total == 1 + assert [r.resource_type for r in by_type] == [models.ResourceType.MODEL] + + def test_unknown_resource_type(self): + rc, by_type, _, unknown = builder.aggregate_node_results( + [("analysis.p.a", "success"), ("model.p.m", "success")] + ) + assert unknown == 1 + assert rc.total == 2 + assert [r.resource_type for r in by_type] == [models.ResourceType.MODEL] + + def test_unavailable_resource_type_is_unknown(self): + rc, _, _, unknown = builder.aggregate_node_results([(None, "skipped")]) + assert rc.total == 1 + assert unknown == 1 + + def test_invariant_known_plus_unknown_equals_total(self): + rc, by_type, _, unknown = builder.aggregate_node_results( + [ + ("model.p.m", "success"), + ("test.p.t", "pass"), + ("analysis.p.a", "fail"), + ("operation.p.h", "success"), + ] + ) + known_total = sum(r.status_counts.total for r in by_type) + assert known_total + unknown == rc.total diff --git a/tests/unit/telemetry/test_config.py b/tests/unit/telemetry/test_config.py new file mode 100644 index 000000000..8dcdfdb17 --- /dev/null +++ b/tests/unit/telemetry/test_config.py @@ -0,0 +1,53 @@ +from types import SimpleNamespace + +from dbt.adapters.databricks.telemetry import config + + +def _creds(connection_parameters): + return SimpleNamespace( + connection_parameters=connection_parameters, + auth_type=None, + token=None, + client_secret=None, + azure_client_secret=None, + ) + + +class TestOptIn: + def test_defaults_off(self): + assert config.is_enabled(_creds({})) is False + assert config.is_enabled(_creds(None)) is False + + def test_explicit_opt_in(self): + assert config.is_enabled(_creds({"enable_dbt_telemetry": True})) is True + + +class TestCommandEligibility: + def test_warehouse_graph_commands_are_eligible(self, monkeypatch): + from dbt import flags + + monkeypatch.setattr(flags, "get_flags", lambda: SimpleNamespace(WHICH="build")) + assert config.is_eligible_command() is True + + def test_parse_only_and_unwired_task_shapes_are_ineligible(self, monkeypatch): + from dbt import flags + + for command in ("compile", "source freshness", "run-operation"): + monkeypatch.setattr( + flags, + "get_flags", + lambda command=command: SimpleNamespace(WHICH=command), + ) + assert config.is_eligible_command() is False + + +class TestTransportEligibility: + def test_kernel_u2m_is_ineligible(self): + creds = _creds({"use_kernel": True}) + creds.auth_type = "oauth" + assert config.has_reusable_transport(creds) is False + + def test_kernel_pat_is_eligible(self): + creds = _creds({"use_kernel": True}) + creds.token = "token" + assert config.has_reusable_transport(creds) is True diff --git a/tests/unit/telemetry/test_coordinator.py b/tests/unit/telemetry/test_coordinator.py new file mode 100644 index 000000000..b931651f2 --- /dev/null +++ b/tests/unit/telemetry/test_coordinator.py @@ -0,0 +1,281 @@ +import json +import threading + +from dbt.adapters.databricks.telemetry import coordinator as coord_mod +from dbt.adapters.databricks.telemetry import models + + +def _log(invocation_id="inv-1"): + return models.TelemetryLog( + invocation_id=invocation_id, + adapter_version="1.2.3", + dbt_core_version="1.12.0", + post_parse=models.PostParsePayload( + invocation_config=models.InvocationConfig(), + manifest_stats=models.ManifestStats(), + connection_config=models.ConnectionConfig(), + project_config=models.ProjectConfig(), + ), + ) + + +def _run_log(invocation_id="inv-1"): + return models.TelemetryLog( + invocation_id=invocation_id, + adapter_version="1.2.3", + dbt_core_version="1.12.0", + event_type=models.EventType.POST_RUN, + post_run=models.PostRunPayload(), + ) + + +def _transport(header_factory=lambda: {"Authorization": "Bearer x"}, workspace_id="42"): + return coord_mod.Transport( + host="https://h", header_factory=header_factory, workspace_id=workspace_id + ) + + +class _Capture: + def __init__(self): + self.calls = [] + + def __call__(self, host, body, header_factory=None, workspace_id=None): + self.calls.append((host, body, header_factory, workspace_id)) + return True + + +class TestTransportOpacity: + def test_repr_is_redacted(self): + t = _transport() + assert "redacted" in repr(t) + assert "Bearer" not in repr(t) + + def test_not_a_dataclass_and_no_dict(self): + t = _transport() + assert not hasattr(t, "__dict__") + + +class TestSendOrdering: + def test_parse_then_transport_sends_once(self, monkeypatch): + capture = _Capture() + monkeypatch.setattr(coord_mod.client, "send", capture) + c = coord_mod.Coordinator() + c.set_post_parse("inv-1", _log()) + assert capture.calls == [] + c.set_transport("inv-1", _transport()) + c.flush() + assert len(capture.calls) == 1 + + def test_transport_then_parse_sends_once(self, monkeypatch): + capture = _Capture() + monkeypatch.setattr(coord_mod.client, "send", capture) + c = coord_mod.Coordinator() + c.set_transport("inv-1", _transport()) + assert capture.calls == [] + c.set_post_parse("inv-1", _log()) + c.flush() + assert len(capture.calls) == 1 + + def test_repeated_parse_is_idempotent(self, monkeypatch): + capture = _Capture() + monkeypatch.setattr(coord_mod.client, "send", capture) + c = coord_mod.Coordinator() + c.set_transport("inv-1", _transport()) + c.set_post_parse("inv-1", _log()) + c.set_post_parse("inv-1", _log()) + c.flush() + assert len(capture.calls) == 1 + + def test_transport_without_auth_headers_does_not_send(self, monkeypatch): + capture = _Capture() + monkeypatch.setattr(coord_mod.client, "send", capture) + c = coord_mod.Coordinator() + c.set_post_parse("inv-1", _log()) + c.set_transport("inv-1", _transport(header_factory=None)) + c.flush() + assert capture.calls == [] + + def test_transport_does_not_block_on_slow_send(self, monkeypatch): + in_send = threading.Event() + release = threading.Event() + + def slow_send(host, body, header_factory=None, workspace_id=None): + in_send.set() + assert release.wait(timeout=2) + return True + + monkeypatch.setattr(coord_mod.client, "send", slow_send) + c = coord_mod.Coordinator() + c.set_post_parse("inv-1", _log()) + c.set_transport("inv-1", _transport()) + assert in_send.wait(timeout=2) + release.set() + c.flush(timeout=2) + + +class TestIsolationAndClose: + def test_distinct_invocations_do_not_cross_pair(self, monkeypatch): + capture = _Capture() + monkeypatch.setattr(coord_mod.client, "send", capture) + c = coord_mod.Coordinator() + c.set_post_parse("inv-A", _log("inv-A")) + c.set_transport("inv-B", _transport()) + c.flush() + assert capture.calls == [] + + def test_closed_invocation_rejects_late_callbacks(self, monkeypatch): + capture = _Capture() + monkeypatch.setattr(coord_mod.client, "send", capture) + c = coord_mod.Coordinator() + c.close("inv-1") + c.set_post_parse("inv-1", _log()) + c.set_transport("inv-1", _transport()) + c.flush() + assert capture.calls == [] + + +class TestPostRun: + def _entry(self, call): + return json.loads(call[1]["protoLogs"][0])["entry"]["dbt_databricks_telemetry_log"] + + def test_both_phases_send_post_parse_first(self, monkeypatch): + capture = _Capture() + monkeypatch.setattr(coord_mod.client, "send", capture) + c = coord_mod.Coordinator() + c.set_post_parse("inv-1", _log()) + c.set_post_run("inv-1", _run_log()) + c.set_transport("inv-1", _transport()) + c.flush() + assert len(capture.calls) == 2 + assert self._entry(capture.calls[0])["event_type"] == "POST_PARSE" + assert self._entry(capture.calls[1])["event_type"] == "POST_RUN" + + def test_close_waits_for_post_run_behind_slow_post_parse(self, monkeypatch): + started = threading.Event() + release = threading.Event() + phases = [] + + def blocking_send(host, body, header_factory=None, workspace_id=None): + entry = json.loads(body["protoLogs"][0])["entry"]["dbt_databricks_telemetry_log"] + phases.append(entry["event_type"]) + if entry["event_type"] == "POST_PARSE": + started.set() + assert release.wait(timeout=2) + return True + + monkeypatch.setattr(coord_mod.client, "send", blocking_send) + monkeypatch.setattr(coord_mod.client, "_TIMEOUT_SECONDS", 0.25) + c = coord_mod.Coordinator() + c.set_post_parse("inv-1", _log()) + c.set_transport("inv-1", _transport()) + assert started.wait(timeout=2) + c.set_post_run("inv-1", _run_log()) + + timer = threading.Timer(0.35, release.set) + timer.start() + c.flush() + c.close("inv-1") + timer.join() + + assert phases == ["POST_PARSE", "POST_RUN"] + + def test_phases_use_distinct_event_ids(self, monkeypatch): + capture = _Capture() + monkeypatch.setattr(coord_mod.client, "send", capture) + c = coord_mod.Coordinator() + c.set_transport("inv-1", _transport()) + c.set_post_parse("inv-1", _log()) + c.set_post_run("inv-1", _run_log()) + c.flush() + ids = { + json.loads(call[1]["protoLogs"][0])["frontend_log_event_id"] for call in capture.calls + } + assert len(ids) == 2 + + +class TestResultCapture: + def test_node_results_fallback_snapshot(self): + c = coord_mod.Coordinator() + c.mark_start("inv-1") + c.record_expected_count("inv-1", 2) + c.record_node_result("inv-1", "model.p.m1", "success") + c.record_node_result("inv-1", "test.p.t", "pass") + results, selected, expected, coverage_complete, captured = c.result_snapshot("inv-1") + assert captured is False + assert selected == 2 + assert expected == 2 + assert coverage_complete is False + assert results == [("model.p.m1", "success"), ("test.p.t", "pass")] + + def test_partial_snapshot_omits_unprovable_selected_ephemeral_count(self): + c = coord_mod.Coordinator() + c.mark_start("inv-1") + c.record_expected_count("inv-1", 1) + c.record_ephemeral_ids("inv-1", {"model.p.ephemeral1", "model.p.ephemeral2"}) + c.record_node_result("inv-1", "model.p.ephemeral1", "success") + + _, selected, _, _, captured = c.result_snapshot("inv-1") + + assert selected is None + assert captured is False + + def test_end_run_is_authoritative_and_complete(self): + c = coord_mod.Coordinator() + c.mark_start("inv-1") + c.record_expected_count("inv-1", 2) + c.record_node_result("inv-1", "model.p.m1", "success") + c.record_end_run("inv-1", ["success", "skipped"]) + results, _, _, coverage_complete, _ = c.result_snapshot("inv-1") + assert coverage_complete is True + assert results == [("model.p.m1", "success"), (None, "skipped")] + + def test_hooks_are_auxiliary_and_subtracted_from_end_statuses(self): + c = coord_mod.Coordinator() + c.mark_start("inv-1") + c.record_expected_count("inv-1", 1) + c.record_node_result("inv-1", "model.p.m1", "success") + c.record_hook_result("inv-1", "success") + c.record_end_run("inv-1", ["success", "success"]) + + results, selected, expected, coverage_complete, _ = c.result_snapshot("inv-1") + + assert results == [("model.p.m1", "success"), ("operation", "success")] + assert selected == expected == 1 + assert coverage_complete is True + + def test_fail_fast_synthetic_ephemeral_counts_as_selected_not_result(self): + c = coord_mod.Coordinator() + c.mark_start("inv-1") + c.record_expected_count("inv-1", 2) + c.record_ephemeral_ids("inv-1", {"model.p.ephemeral"}) + c.record_node_result("inv-1", "model.p.m1", "error") + c.record_end_run("inv-1", ["error", "skipped", "skipped"]) + + results, selected, expected, coverage_complete, _ = c.result_snapshot("inv-1") + + assert results == [("model.p.m1", "error"), (None, "skipped")] + assert selected == 3 + assert expected == 2 + assert coverage_complete is True + + def test_record_after_close_ignored(self): + c = coord_mod.Coordinator() + c.mark_start("inv-1") + c.close("inv-1") + c.record_node_result("inv-1", "model.p.m1", "success") + results, _, _, _, _ = c.result_snapshot("inv-1") + assert results == [] + + def test_ephemeral_is_selected_but_not_a_top_level_result(self): + c = coord_mod.Coordinator() + c.mark_start("inv-1") + c.record_expected_count("inv-1", 1) + c.record_ephemeral_ids("inv-1", {"model.p.ephemeral"}) + c.record_node_result("inv-1", "model.p.ephemeral", "success") + c.record_node_result("inv-1", "model.p.table", "success") + results, selected, expected, coverage_complete, captured = c.result_snapshot("inv-1") + assert results == [("model.p.table", "success")] + assert selected == 2 + assert expected == 1 + assert coverage_complete is False + assert captured is False diff --git a/tests/unit/telemetry/test_encoder.py b/tests/unit/telemetry/test_encoder.py new file mode 100644 index 000000000..22211799b --- /dev/null +++ b/tests/unit/telemetry/test_encoder.py @@ -0,0 +1,212 @@ +import json + +from dbt.adapters.databricks.telemetry import encoder, models + + +def _log(): + return models.TelemetryLog( + invocation_id="inv-1", + adapter_version="1.2.3", + dbt_core_version="1.12.0", + post_parse=models.PostParsePayload( + invocation_config=models.InvocationConfig( + thread_count=4, dbt_command=models.DbtCommand.RUN + ), + manifest_stats=models.ManifestStats(enabled_total=models.ResourceCounts(model_count=3)), + connection_config=models.ConnectionConfig( + default_compute_type=models.ComputeType.SQL_WAREHOUSE, + configured_auth_family=models.AuthFamily.PAT, + ), + project_config=models.ProjectConfig(use_materialization_v2=True), + ), + ) + + +def _post_run_log(): + return models.TelemetryLog( + invocation_id="inv-2", + adapter_version="1.2.3", + dbt_core_version="1.12.0", + event_type=models.EventType.POST_RUN, + post_run=models.PostRunPayload( + run_outcome=models.RunOutcome( + invocation_status=models.InvocationStatus.HANDLED_ERROR, + termination_reason=models.TerminationReason.NORMAL, + invocation_duration_ms=1234, + result_aggregates_available=True, + expected_result_coverage_complete=True, + ), + selected_resources=10, + result_counts=models.NodeStatusCounts(total=8, success=6, fail=1, pass_=5), + results_by_resource_type=[ + models.ResourceOutcomeStats( + resource_type=models.ResourceType.MODEL, + status_counts=models.NodeStatusCounts(total=6, success=6), + ) + ], + auxiliary_hook_results=models.NodeStatusCounts(), + unknown_resource_type_results=1, + ), + ) + + +class TestEncoder: + def test_uses_dedicated_entry_field(self): + fe = json.loads(encoder.encode_request(_log(), "evt-1")["protoLogs"][0]) + assert "dbt_databricks_telemetry_log" in fe["entry"] + + def test_post_parse_contains_every_proto_field(self): + entry = json.loads(encoder.encode_request(_log(), "evt-1")["protoLogs"][0])["entry"][ + "dbt_databricks_telemetry_log" + ] + assert set(entry) == { + "invocation_id", + "event_type", + "adapter_version", + "dbt_core_version", + "post_parse", + } + assert set(entry["post_parse"]) == { + "invocation_config", + "manifest_stats", + "connection_config", + "project_config", + } + post_parse = entry["post_parse"] + assert set(post_parse["invocation_config"]) == { + "thread_count", + "dbt_command", + "full_refresh", + "empty", + "fail_fast", + "warn_error_policy", + } + assert set(post_parse["connection_config"]) == { + "default_compute_type", + "configured_auth_family", + "named_compute_count", + "spog_routing_configured", + "use_kernel", + } + assert set(post_parse["project_config"]) == { + "use_user_folder_for_python", + "use_materialization_v2", + "use_replace_on_for_insert_overwrite", + "use_managed_iceberg", + "use_concurrent_microbatch", + "use_describe_as_json_for_relation_metadata", + } + assert set(post_parse["manifest_stats"]) == { + "enabled_total", + "enabled_root_project", + "enabled_installed_packages", + } + assert set(post_parse["manifest_stats"]["enabled_total"]) == { + "model_count", + "data_test_count", + "generic_data_test_count", + "seed_count", + "snapshot_count", + "source_count", + "function_count", + "exposure_count", + "saved_query_count", + "other_count", + "unit_test_count", + } + + def test_numeric_workspace_id_is_coerced(self): + fe = json.loads(encoder.encode_request(_log(), "e", workspace_id="42")["protoLogs"][0]) + assert fe["workspace_id"] == 42 + + def test_non_numeric_workspace_id_is_omitted(self): + fe = json.loads(encoder.encode_request(_log(), "e", workspace_id="abc")["protoLogs"][0]) + assert "workspace_id" not in fe + + def test_post_parse_omits_unset_phase(self): + entry = json.loads(encoder.encode_request(_log(), "e")["protoLogs"][0])["entry"][ + "dbt_databricks_telemetry_log" + ] + assert "post_parse" in entry + assert "post_run" not in entry + + +class TestPostRunEncoder: + def _entry(self): + fe = json.loads(encoder.encode_request(_post_run_log(), "e")["protoLogs"][0]) + return fe["entry"]["dbt_databricks_telemetry_log"] + + def test_only_post_run_phase(self): + entry = self._entry() + assert entry["event_type"] == "POST_RUN" + assert "post_run" in entry + assert "post_parse" not in entry + + def test_pass_field_uses_proto_name(self): + rc = self._entry()["post_run"]["result_counts"] + assert rc["pass"] == 5 + assert "pass_" not in rc + + def test_outcome_and_resource_enum_names(self): + post_run = self._entry()["post_run"] + assert post_run["run_outcome"]["invocation_status"] == "HANDLED_ERROR" + assert post_run["run_outcome"]["invocation_duration_ms"] == 1234 + assert post_run["results_by_resource_type"][0]["resource_type"] == "MODEL" + + def test_post_run_contains_every_proto_field(self): + post_run = self._entry()["post_run"] + assert set(post_run) == { + "run_outcome", + "selected_resources", + "expected_result_resources", + "result_counts", + "results_by_resource_type", + "auxiliary_hook_results", + "unknown_resource_type_results", + } + assert set(post_run["run_outcome"]) == { + "invocation_status", + "termination_reason", + "invocation_duration_ms", + "result_aggregates_available", + "expected_result_coverage_complete", + } + assert set(post_run["result_counts"]) == { + "total", + "success", + "error", + "fail", + "warn", + "skipped", + "partial_success", + "pass", + "runtime_error", + "no_op", + "reused", + } + assert set(post_run["results_by_resource_type"][0]) == { + "resource_type", + "status_counts", + } + + def test_unavailable_aggregates_are_omitted(self): + log = models.TelemetryLog( + invocation_id="inv-3", + adapter_version="1.2.3", + dbt_core_version="1.12.0", + event_type=models.EventType.POST_RUN, + post_run=models.PostRunPayload( + run_outcome=models.RunOutcome(result_aggregates_available=False), + expected_result_resources=2, + ), + ) + entry = json.loads(encoder.encode_request(log, "e")["protoLogs"][0])["entry"][ + "dbt_databricks_telemetry_log" + ]["post_run"] + assert set(entry) == {"run_outcome", "expected_result_resources"} + assert set(entry["run_outcome"]) == { + "invocation_status", + "termination_reason", + "invocation_duration_ms", + "result_aggregates_available", + } diff --git a/tests/unit/telemetry/test_hooks.py b/tests/unit/telemetry/test_hooks.py new file mode 100644 index 000000000..bd257c34c --- /dev/null +++ b/tests/unit/telemetry/test_hooks.py @@ -0,0 +1,35 @@ +from types import SimpleNamespace +from unittest.mock import Mock + +from dbt.adapters.databricks.telemetry import hooks + + +def test_post_parse_is_not_rebuilt(monkeypatch): + coord = Mock() + coord.needs_post_parse.return_value = False + build = Mock() + monkeypatch.setattr(hooks, "coordinator", lambda: coord) + monkeypatch.setattr(hooks, "_current_invocation_id", lambda: "inv-1") + monkeypatch.setattr(hooks, "DatabricksCredentials", object) + monkeypatch.setattr(hooks, "is_enabled_for_invocation", lambda _: True) + monkeypatch.setattr(hooks.builder, "build_post_parse_log", build) + adapter = SimpleNamespace(config=SimpleNamespace(credentials=SimpleNamespace())) + + hooks.on_post_parse(adapter, SimpleNamespace()) + + coord.needs_post_parse.assert_called_once_with("inv-1") + build.assert_not_called() + + +def test_closed_invocation_is_not_finalized_again(monkeypatch): + coord = Mock() + coord.is_closed.return_value = True + build = Mock() + monkeypatch.setattr(hooks, "coordinator", lambda: coord) + monkeypatch.setattr(hooks.builder, "build_post_run_log", build) + + hooks._finalize_post_run("inv-1", None) + + coord.is_closed.assert_called_once_with("inv-1") + coord.result_snapshot.assert_not_called() + build.assert_not_called() diff --git a/tests/unit/telemetry/test_listener.py b/tests/unit/telemetry/test_listener.py new file mode 100644 index 000000000..ad0dbd975 --- /dev/null +++ b/tests/unit/telemetry/test_listener.py @@ -0,0 +1,54 @@ +from types import SimpleNamespace +from unittest.mock import Mock + +from dbt.adapters.databricks.telemetry import listener + + +def _message(name, data): + return SimpleNamespace(info=SimpleNamespace(name=name), data=data) + + +def test_end_run_records_authoritative_results_then_finalizes(monkeypatch): + coord = Mock() + monkeypatch.setattr(listener, "coordinator", lambda: coord) + monkeypatch.setattr(listener, "_current_invocation_id", lambda: "inv-1") + + from dbt.adapters.databricks.telemetry import hooks + + finalize = Mock() + monkeypatch.setattr(hooks, "on_end_run_result", finalize) + results = [SimpleNamespace(status="success")] + + listener._on_event(_message("EndRunResult", SimpleNamespace(results=results, success=True))) + + coord.record_end_run.assert_called_once_with("inv-1", ["success"], success=True) + finalize.assert_called_once_with("inv-1") + + +def test_generic_exception_records_typed_error(monkeypatch): + coord = Mock() + monkeypatch.setattr(listener, "coordinator", lambda: coord) + monkeypatch.setattr(listener, "_current_invocation_id", lambda: "inv-1") + + listener._on_event( + _message( + "GenericExceptionOnRun", + SimpleNamespace( + unique_id="model.p.m", + node_info=SimpleNamespace(unique_id="model.p.m"), + ), + ) + ) + + coord.record_node_result.assert_called_once_with("inv-1", "model.p.m", "error") + + +def test_pre_end_failure_marks_fail_fast_triggered(monkeypatch): + coord = Mock() + monkeypatch.setattr(listener, "coordinator", lambda: coord) + monkeypatch.setattr(listener, "_current_invocation_id", lambda: "inv-1") + monkeypatch.setattr(listener, "_fail_fast_enabled", lambda: True) + + listener._on_event(_message("RunResultFailure", SimpleNamespace())) + + coord.mark_fail_fast_triggered.assert_called_once_with("inv-1")