Adding detailed telemetry

demos/dqx_demo_pii_detection.py

@@ -3,7 +3,7 @@
 # MAGIC # Using DQX for PII Detection
 # MAGIC Increased regulation makes Databricks customers responsible for any Personally Identifiable Information (PII) stored in Unity Catalog. Companies need to be able to perform PII detection for data at-rest and in-transit to proactively quarantine or anonymize PII before persisting the data.
 # MAGIC
-# MAGIC DQX provides in-flight data quality monitoring for Spark `DataFrames`. You can apply checks, get row-level metadata, and quarantine failing records. Workloads can also use DQX's built-in functions to check `DataFrames` for PII.
+# MAGIC DQX provides in-transit data quality monitoring for Spark `DataFrames`. You can apply checks, get row-level metadata, and quarantine failing records. Workloads can also use DQX's built-in functions to check `DataFrames` for PII.
 
 # COMMAND ----------
 

docs/dqx/docs/guide/index.mdx

@@ -21,7 +21,7 @@ For more details, see the [Installation Guide](/docs/installation/).
 
 | Quality checking type | Integration with processing pipelines | Description                                                                                                                                                                                                    |
 |---------------------- | ------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
-| **In-flight**         | Code-level only                       | DQX allows data quality to be validated on the fly while the data is being processed, before it is written to storage. This requires DQX to be used as a library and integrated directly into user pipelines.  |
+| **In-transit**        | Code-level only                       | DQX allows data quality to be validated on the fly while the data is being processed, before it is written to storage. This requires DQX to be used as a library and integrated directly into user pipelines.  |
 | **At-rest**           | Code-level or No-code (Workflows)     | DQX enables data quality checking on existing data stored in tables. For no-code integration, DQX must first be installed in the workspace as a tool to deploy workflows.                                      |
 
 **Integration options:**

docs/dqx/docs/reference/engine.mdx

@@ -60,8 +60,8 @@ The following table outlines the available methods of the `DQEngine` and their f
 | `validate_checks`                              | Validates the provided quality checks to ensure they conform to the expected structure and types.                                                                                                                                 | `checks`: List of checks to validate; `custom_check_functions`: (optional) dictionary of custom check functions that can be used; `validate_custom_check_functions`: (optional) if set to True, validates custom check functions (defaults to True).                                                                                                                                                                                                                                                                                                                                                                                              | Yes                                       |
 | `get_invalid`                                  | Retrieves records from the DataFrame that violate data quality checks (records with warnings and errors).                                                                                                                         | `df`: Input DataFrame.                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                            | Yes                                       |
 | `get_valid`                                    | Retrieves records from the DataFrame that pass all data quality checks.                                                                                                                                                           | `df`: Input DataFrame.                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                            | Yes                                       |
-| `load_checks`                                  | Loads quality rules (checks) from storage backend. Multiple storage backends are supported including tables, files or workspace files, installation-managed sources where the location is inferred automatically from run config  | `config`: Configuration for loading checks from a storage backend, i.e. `FileChecksStorageConfig`: file in a local filesystem (YAML or JSON), or workspace files if invoked from Databricks notebook or job; `WorkspaceFileChecksStorageConfig`: file in a workspace (YAML or JSON) using absolute paths; `VolumeFileChecksStorageConfig`: file in a Unity Catalog Volume (YAML or JSON); `TableChecksStorageConfig`: a table; `InstallationChecksStorageConfig`: installation-managed storage backend, using the `checks_location` field from the run configuration. See more details below.                                                     | Yes (only with `FileChecksStorageConfig`) |
-| `save_checks`                                  | Saves quality rules (checks) to storage backend. Multiple storage backends are supported including tables, files or workspace files, installation-managed targets where the location is inferred automatically from run config    | `checks`: List of checks defined as dictionary; `config`: Configuration for saving checks in a storage backend, i.e. `FileChecksStorageConfig`: file in a local filesystem (YAML or JSON), or workspace files if invoked from Databricks notebook or job; `WorkspaceFileChecksStorageConfig`: file in a workspace (YAML or JSON); `VolumeFileChecksStorageConfig`: file in a Unity Catalog Volume (YAML or JSON); `TableChecksStorageConfig`: a table; `InstallationChecksStorageConfig`: storage defined in the installation context, using the `checks_location` field from the run configuration. See more details below.                      | Yes (only with `FileChecksStorageConfig`) |
+| `load_checks`                                  | Loads quality rules (checks) from storage backend. Multiple storage backends are supported including tables, files or workspace files, installation-managed sources where the location is inferred automatically from run config. | `config`: Configuration for loading checks from a storage backend, i.e. `FileChecksStorageConfig`: file in a local filesystem (YAML or JSON), or workspace files if invoked from Databricks notebook or job; `WorkspaceFileChecksStorageConfig`: file in a workspace (YAML or JSON) using absolute paths; `VolumeFileChecksStorageConfig`: file in a Unity Catalog Volume (YAML or JSON); `TableChecksStorageConfig`: a table; `InstallationChecksStorageConfig`: installation-managed storage backend, using the `checks_location` field from the run configuration. See more details below.                                                     | Yes (only with `FileChecksStorageConfig`) |
+| `save_checks`                                  | Saves quality rules (checks) to storage backend. Multiple storage backends are supported including tables, files or workspace files, installation-managed targets where the location is inferred automatically from run config.   | `checks`: List of checks defined as dictionary; `config`: Configuration for saving checks in a storage backend, i.e. `FileChecksStorageConfig`: file in a local filesystem (YAML or JSON), or workspace files if invoked from Databricks notebook or job; `WorkspaceFileChecksStorageConfig`: file in a workspace (YAML or JSON); `VolumeFileChecksStorageConfig`: file in a Unity Catalog Volume (YAML or JSON); `TableChecksStorageConfig`: a table; `InstallationChecksStorageConfig`: storage defined in the installation context, using the `checks_location` field from the run configuration. See more details below.                      | Yes (only with `FileChecksStorageConfig`) |
 | `save_results_in_table`                        | Save quality checking results in delta table(s).                                                                                                                                                                                  | `output_df`: (optional) Dataframe containing the output data; `quarantine_df`: (optional) Dataframe containing the output data; `output_config`: `OutputConfig` object with the table name, output mode, and options for the output data; `quarantine_config`: `OutputConfig` object with the table name, output mode, and options for the quarantine data - if provided, data will be split; `run_config_name`: Name of the run config to use; `assume_user`: If True, assume user installation.                                                                                                                                                 | No                                        |
 
 The 'Supports local execution' in the above table indicates which methods can be used for local testing without a Databricks workspace (see the usage in [local testing section](/docs/reference/testing/#local-testing-with-dqengine)).

pyproject.toml

@@ -172,7 +172,7 @@ extend-exclude = 'demos/|tests/e2e/notebooks'
 cache-dir = ".venv/ruff-cache"
 target-version = "py310"
 line-length = 120
-exclude = ["demos/*", "tests/e2e/notebooks"]
+exclude = ["demos/*", "tests/e2e/notebooks", "ui/*"]
 
 [tool.ruff.lint.isort]
 known-first-party = ["databricks.labs.dqx"]

src/databricks/labs/dqx/checks_storage.py

@@ -28,6 +28,7 @@
     get_file_deserializer,
 )
 from databricks.labs.dqx.config_loader import RunConfigLoader
+from databricks.labs.dqx.telemetry import telemetry_logger
 from databricks.labs.dqx.utils import TABLE_PATTERN
 from databricks.labs.dqx.checks_serializer import FILE_SERIALIZERS
 
@@ -69,6 +70,7 @@ def __init__(self, ws: WorkspaceClient, spark: SparkSession):
         self.ws = ws
         self.spark = spark
 
+    @telemetry_logger("load_checks", "table")
     def load(self, config: TableChecksStorageConfig) -> list[dict]:
         """
         Load checks (dq rules) from a Delta table in the workspace.
@@ -85,6 +87,7 @@ def load(self, config: TableChecksStorageConfig) -> list[dict]:
         rules_df = self.spark.read.table(config.location)
         return serialize_checks_from_dataframe(rules_df, run_config_name=config.run_config_name) or []
 
+    @telemetry_logger("save_checks", "table")
     def save(self, checks: list[dict], config: TableChecksStorageConfig) -> None:
         """
         Save checks to a Delta table in the workspace.
@@ -111,6 +114,7 @@ class WorkspaceFileChecksStorageHandler(ChecksStorageHandler[WorkspaceFileChecks
     def __init__(self, ws: WorkspaceClient):
         self.ws = ws
 
+    @telemetry_logger("load_checks", "workspace_file")
     def load(self, config: WorkspaceFileChecksStorageConfig) -> list[dict]:
         """Load checks (dq rules) from a file (json or yaml) in the workspace.
         This does not require installation of DQX in the workspace.
@@ -137,6 +141,7 @@ def load(self, config: WorkspaceFileChecksStorageConfig) -> list[dict]:
         except (yaml.YAMLError, json.JSONDecodeError) as e:
             raise ValueError(f"Invalid checks in file: {file_path}: {e}") from e
 
+    @telemetry_logger("save_checks", "workspace_file")
     def save(self, checks: list[dict], config: WorkspaceFileChecksStorageConfig) -> None:
         """Save checks (dq rules) to yaml file in the workspace.
         This does not require installation of DQX in the workspace.
@@ -217,11 +222,13 @@ class InstallationChecksStorageHandler(ChecksStorageHandler[InstallationChecksSt
     """
 
     def __init__(self, ws: WorkspaceClient, spark: SparkSession, run_config_loader: RunConfigLoader | None = None):
+        self.ws = ws
         self._run_config_loader = run_config_loader or RunConfigLoader(ws)
         self.workspace_file_handler = WorkspaceFileChecksStorageHandler(ws)
         self.table_handler = TableChecksStorageHandler(ws, spark)
         self.volume_handler = VolumeFileChecksStorageHandler(ws)
 
+    @telemetry_logger("load_checks", "installation")
     def load(self, config: InstallationChecksStorageConfig) -> list[dict]:
         """
         Load checks (dq rules) from the installation configuration.
@@ -238,6 +245,7 @@ def load(self, config: InstallationChecksStorageConfig) -> list[dict]:
         handler, config = self._get_storage_handler_and_config(config)
         return handler.load(config)
 
+    @telemetry_logger("save_checks", "installation")
     def save(self, checks: list[dict], config: InstallationChecksStorageConfig) -> None:
         """
         Save checks (dq rules) to yaml file or table in the installation folder.
@@ -285,6 +293,7 @@ class VolumeFileChecksStorageHandler(ChecksStorageHandler[VolumeFileChecksStorag
     def __init__(self, ws: WorkspaceClient):
         self.ws = ws
 
+    @telemetry_logger("load_checks", "volume")
     def load(self, config: VolumeFileChecksStorageConfig) -> list[dict]:
         """Load checks (dq rules) from a file (json or yaml) in a Unity Catalog volume.
 
@@ -316,6 +325,7 @@ def load(self, config: VolumeFileChecksStorageConfig) -> list[dict]:
         except (yaml.YAMLError, json.JSONDecodeError) as e:
             raise ValueError(f"Invalid checks in file: {file_path}: {e}") from e
 
+    @telemetry_logger("save_checks", "volume")
     def save(self, checks: list[dict], config: VolumeFileChecksStorageConfig) -> None:
         """Save checks (dq rules) to yaml file in a Unity Catalog volume.
         This does not require installation of DQX in a Unity Catalog volume.

src/databricks/labs/dqx/engine.py

@@ -35,6 +35,7 @@
 from databricks.labs.dqx.checks_validator import ChecksValidator, ChecksValidationStatus
 from databricks.labs.dqx.schema import dq_result_schema
 from databricks.labs.dqx.utils import read_input_data, save_dataframe_as_table
+from databricks.labs.dqx.telemetry import telemetry_logger, log_telemetry
 from databricks.sdk import WorkspaceClient
 
 logger = logging.getLogger(__name__)
@@ -336,6 +337,7 @@ def _create_results_array(
                 run_time=self.run_time,
                 ref_dfs=ref_dfs,
             )
+            log_telemetry(self.ws, "check", check.check_func.__name__)
             result = manager.process()
             check_conditions.append(result.condition)
             # The DataFrame should contain any new columns added by the dataset-level checks
@@ -382,6 +384,7 @@ def __init__(
             checks_handler_factory or ChecksStorageHandlerFactory(self.ws, self.spark)
         )
 
+    @telemetry_logger("engine", "apply_checks")
     def apply_checks(
         self, df: DataFrame, checks: list[DQRule], ref_dfs: dict[str, DataFrame] | None = None
     ) -> DataFrame:
@@ -397,6 +400,7 @@ def apply_checks(
         """
         return self._engine.apply_checks(df, checks, ref_dfs)
 
+    @telemetry_logger("engine", "apply_checks_and_split")
     def apply_checks_and_split(
         self, df: DataFrame, checks: list[DQRule], ref_dfs: dict[str, DataFrame] | None = None
     ) -> tuple[DataFrame, DataFrame]:
@@ -414,6 +418,7 @@ def apply_checks_and_split(
         """
         return self._engine.apply_checks_and_split(df, checks, ref_dfs)
 
+    @telemetry_logger("engine", "apply_checks_by_metadata")
     def apply_checks_by_metadata(
         self,
         df: DataFrame,
@@ -438,6 +443,7 @@ def apply_checks_by_metadata(
         """
         return self._engine.apply_checks_by_metadata(df, checks, custom_check_functions, ref_dfs)
 
+    @telemetry_logger("engine", "apply_checks_by_metadata_and_split")
     def apply_checks_by_metadata_and_split(
         self,
         df: DataFrame,
@@ -463,6 +469,7 @@ def apply_checks_by_metadata_and_split(
         """
         return self._engine.apply_checks_by_metadata_and_split(df, checks, custom_check_functions, ref_dfs)
 
+    @telemetry_logger("engine", "apply_checks_and_save_in_table")
     def apply_checks_and_save_in_table(
         self,
         checks: list[DQRule],
@@ -500,6 +507,7 @@ def apply_checks_and_save_in_table(
             checked_df = self.apply_checks(df, checks, ref_dfs)
             save_dataframe_as_table(checked_df, output_config)
 
+    @telemetry_logger("engine", "apply_checks_by_metadata_and_save_in_table")
     def apply_checks_by_metadata_and_save_in_table(
         self,
         checks: list[dict],
@@ -590,6 +598,7 @@ def get_valid(self, df: DataFrame) -> DataFrame:
         """
         return self._engine.get_valid(df)
 
+    @telemetry_logger("engine", "save_results_in_table")
     def save_results_in_table(
         self,
         output_df: DataFrame | None = None,

src/databricks/labs/dqx/telemetry.py

@@ -0,0 +1,58 @@
+import logging
+from collections.abc import Callable
+from databricks.sdk import WorkspaceClient
+from databricks.sdk.config import with_user_agent_extra
+from databricks.sdk.errors import DatabricksError
+
+
+logger = logging.getLogger(__name__)
+
+
+def log_telemetry(ws: WorkspaceClient, key: str, value: str) -> None:
+    """
+    Trace specific telemetry information in the Databricks workspace by setting user agent extra info.
+
+    Args:
+        ws: WorkspaceClient
+        key: telemetry key to log
+        value: telemetry value to log
+    """
+    with_user_agent_extra(key, value)
+    try:
+        ws.current_user.me()
+    except DatabricksError as e:
+        # support local execution
+        logger.debug(f"Databricks workspace is not available: {e}")
+
+
+def telemetry_logger(key: str, value: str) -> Callable:
+    """
+    Decorator to log telemetry for method calls.
+
+    Usage: @telemetry_logger("telemetry_key", "telemetry_value")
+
+    Args:
+        key: Telemetry key to log
+        value: Telemetry value to log
+    """
+
+    def decorator(func: Callable) -> Callable:
+        def wrapper(self, *args, **kwargs):
+            """
+            Expecting the workspace client be available in the calling class as 'ws' or 'workspace_client' attribute.
+            """
+            if hasattr(self, 'ws'):  # requires workspace client to be set
+                log_telemetry(self.ws, key, value)
+            elif hasattr(self, 'workspace_client'):  # requires workspace client to be set
+                log_telemetry(self.workspace_client, key, value)
+            else:
+                raise AttributeError(
+                    f"Workspace client not found on {self.__class__.__name__}. "
+                    f"Telemetry decorator requires 'ws' or 'workspace_client' attribute. "
+                    f"Make sure your class has a workspace client defined."
+                )
+            return func(self, *args, **kwargs)
+
+        return wrapper
+
+    return decorator