From dd336ee905c5bf1a741a5f0da756d3bb9c9f2d17 Mon Sep 17 00:00:00 2001
From: Marcin Wojtyczka <marcin.wojtyczka@databricks.com>
Date: Thu, 4 Sep 2025 12:15:58 +0200
Subject: [PATCH 01/20] added detailed telemetry

---
 demos/dqx_demo_pii_detection.py           |  2 +-
 docs/dqx/docs/guide/index.mdx             |  2 +-
 docs/dqx/docs/reference/engine.mdx        |  4 +-
 pyproject.toml                            |  2 +-
 src/databricks/labs/dqx/checks_storage.py | 10 ++++
 src/databricks/labs/dqx/engine.py         |  9 ++++
 src/databricks/labs/dqx/telemetry.py      | 56 +++++++++++++++++++++++
 7 files changed, 80 insertions(+), 5 deletions(-)
 create mode 100644 src/databricks/labs/dqx/telemetry.py

diff --git a/demos/dqx_demo_pii_detection.py b/demos/dqx_demo_pii_detection.py
index 43d7da5f0..ad722b124 100644
--- a/demos/dqx_demo_pii_detection.py
+++ b/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 ----------
 
diff --git a/docs/dqx/docs/guide/index.mdx b/docs/dqx/docs/guide/index.mdx
index 53fcd5f30..8a1ead03a 100644
--- a/docs/dqx/docs/guide/index.mdx
+++ b/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:**
diff --git a/docs/dqx/docs/reference/engine.mdx b/docs/dqx/docs/reference/engine.mdx
index d453ea020..76d238977 100644
--- a/docs/dqx/docs/reference/engine.mdx
+++ b/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)).
diff --git a/pyproject.toml b/pyproject.toml
index 2e0eff867..0dcf2f2f9 100644
--- a/pyproject.toml
+++ b/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"]
diff --git a/src/databricks/labs/dqx/checks_storage.py b/src/databricks/labs/dqx/checks_storage.py
index 65b65fc0b..3fc20dd04 100644
--- a/src/databricks/labs/dqx/checks_storage.py
+++ b/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 log_telemetry
 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
 
+    @log_telemetry("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 []
 
+    @log_telemetry("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
 
+    @log_telemetry("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
 
+    @log_telemetry("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)
 
+    @log_telemetry("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)
 
+    @log_telemetry("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
 
+    @log_telemetry("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
 
+    @log_telemetry("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.
diff --git a/src/databricks/labs/dqx/engine.py b/src/databricks/labs/dqx/engine.py
index 5734f99ac..f27945047 100644
--- a/src/databricks/labs/dqx/engine.py
+++ b/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 log_telemetry, trace
 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,
             )
+            trace(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)
         )
 
+    @log_telemetry("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)
 
+    @log_telemetry("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)
 
+    @log_telemetry("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)
 
+    @log_telemetry("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)
 
+    @log_telemetry("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)
 
+    @log_telemetry("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)
 
+    @log_telemetry("engine", "save_results_in_table")
     def save_results_in_table(
         self,
         output_df: DataFrame | None = None,
diff --git a/src/databricks/labs/dqx/telemetry.py b/src/databricks/labs/dqx/telemetry.py
new file mode 100644
index 000000000..809eeb3f2
--- /dev/null
+++ b/src/databricks/labs/dqx/telemetry.py
@@ -0,0 +1,56 @@
+import logging
+from typing 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 trace(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: key to log
+        value: 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 log_telemetry(key: str, value: str) -> Callable:
+    """
+    Decorator to automatically log telemetry for method calls.
+
+    Usage: @log_telemetry("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):
+            if hasattr(self, 'ws'):  # requires workspace client to be set
+                trace(self.ws, key, value)
+            elif hasattr(self, 'workspace_client'):  # requires workspace client to be set
+                trace(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

From 6a81a0fc12e97045f3957d95949c935058392631 Mon Sep 17 00:00:00 2001
From: Marcin Wojtyczka <marcin.wojtyczka@databricks.com>
Date: Thu, 4 Sep 2025 14:29:23 +0200
Subject: [PATCH 02/20] refactor

---
 src/databricks/labs/dqx/checks_storage.py | 18 +++++++++---------
 src/databricks/labs/dqx/engine.py         | 18 +++++++++---------
 src/databricks/labs/dqx/telemetry.py      | 22 ++++++++++++----------
 3 files changed, 30 insertions(+), 28 deletions(-)

diff --git a/src/databricks/labs/dqx/checks_storage.py b/src/databricks/labs/dqx/checks_storage.py
index 3fc20dd04..180d81e61 100644
--- a/src/databricks/labs/dqx/checks_storage.py
+++ b/src/databricks/labs/dqx/checks_storage.py
@@ -28,7 +28,7 @@
     get_file_deserializer,
 )
 from databricks.labs.dqx.config_loader import RunConfigLoader
-from databricks.labs.dqx.telemetry import log_telemetry
+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
 
@@ -70,7 +70,7 @@ def __init__(self, ws: WorkspaceClient, spark: SparkSession):
         self.ws = ws
         self.spark = spark
 
-    @log_telemetry("load_checks", "table")
+    @telemetry_logger("load_checks", "table")
     def load(self, config: TableChecksStorageConfig) -> list[dict]:
         """
         Load checks (dq rules) from a Delta table in the workspace.
@@ -87,7 +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 []
 
-    @log_telemetry("save_checks", "table")
+    @telemetry_logger("save_checks", "table")
     def save(self, checks: list[dict], config: TableChecksStorageConfig) -> None:
         """
         Save checks to a Delta table in the workspace.
@@ -114,7 +114,7 @@ class WorkspaceFileChecksStorageHandler(ChecksStorageHandler[WorkspaceFileChecks
     def __init__(self, ws: WorkspaceClient):
         self.ws = ws
 
-    @log_telemetry("load_checks", "workspace_file")
+    @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.
@@ -141,7 +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
 
-    @log_telemetry("save_checks", "workspace_file")
+    @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.
@@ -228,7 +228,7 @@ def __init__(self, ws: WorkspaceClient, spark: SparkSession, run_config_loader:
         self.table_handler = TableChecksStorageHandler(ws, spark)
         self.volume_handler = VolumeFileChecksStorageHandler(ws)
 
-    @log_telemetry("load_checks", "installation")
+    @telemetry_logger("load_checks", "installation")
     def load(self, config: InstallationChecksStorageConfig) -> list[dict]:
         """
         Load checks (dq rules) from the installation configuration.
@@ -245,7 +245,7 @@ def load(self, config: InstallationChecksStorageConfig) -> list[dict]:
         handler, config = self._get_storage_handler_and_config(config)
         return handler.load(config)
 
-    @log_telemetry("save_checks", "installation")
+    @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.
@@ -293,7 +293,7 @@ class VolumeFileChecksStorageHandler(ChecksStorageHandler[VolumeFileChecksStorag
     def __init__(self, ws: WorkspaceClient):
         self.ws = ws
 
-    @log_telemetry("load_checks", "volume")
+    @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.
 
@@ -325,7 +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
 
-    @log_telemetry("save_checks", "volume")
+    @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.
diff --git a/src/databricks/labs/dqx/engine.py b/src/databricks/labs/dqx/engine.py
index f27945047..3bb14b341 100644
--- a/src/databricks/labs/dqx/engine.py
+++ b/src/databricks/labs/dqx/engine.py
@@ -35,7 +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 log_telemetry, trace
+from databricks.labs.dqx.telemetry import telemetry_logger, log_telemetry
 from databricks.sdk import WorkspaceClient
 
 logger = logging.getLogger(__name__)
@@ -337,7 +337,7 @@ def _create_results_array(
                 run_time=self.run_time,
                 ref_dfs=ref_dfs,
             )
-            trace(self.ws, "check", check.check_func.__name__)
+            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
@@ -384,7 +384,7 @@ def __init__(
             checks_handler_factory or ChecksStorageHandlerFactory(self.ws, self.spark)
         )
 
-    @log_telemetry("engine", "apply_checks")
+    @telemetry_logger("engine", "apply_checks")
     def apply_checks(
         self, df: DataFrame, checks: list[DQRule], ref_dfs: dict[str, DataFrame] | None = None
     ) -> DataFrame:
@@ -400,7 +400,7 @@ def apply_checks(
         """
         return self._engine.apply_checks(df, checks, ref_dfs)
 
-    @log_telemetry("engine", "apply_checks_and_split")
+    @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]:
@@ -418,7 +418,7 @@ def apply_checks_and_split(
         """
         return self._engine.apply_checks_and_split(df, checks, ref_dfs)
 
-    @log_telemetry("engine", "apply_checks_by_metadata")
+    @telemetry_logger("engine", "apply_checks_by_metadata")
     def apply_checks_by_metadata(
         self,
         df: DataFrame,
@@ -443,7 +443,7 @@ def apply_checks_by_metadata(
         """
         return self._engine.apply_checks_by_metadata(df, checks, custom_check_functions, ref_dfs)
 
-    @log_telemetry("engine", "apply_checks_by_metadata_and_split")
+    @telemetry_logger("engine", "apply_checks_by_metadata_and_split")
     def apply_checks_by_metadata_and_split(
         self,
         df: DataFrame,
@@ -469,7 +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)
 
-    @log_telemetry("engine", "apply_checks_and_save_in_table")
+    @telemetry_logger("engine", "apply_checks_and_save_in_table")
     def apply_checks_and_save_in_table(
         self,
         checks: list[DQRule],
@@ -507,7 +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)
 
-    @log_telemetry("engine", "apply_checks_by_metadata_and_save_in_table")
+    @telemetry_logger("engine", "apply_checks_by_metadata_and_save_in_table")
     def apply_checks_by_metadata_and_save_in_table(
         self,
         checks: list[dict],
@@ -598,7 +598,7 @@ def get_valid(self, df: DataFrame) -> DataFrame:
         """
         return self._engine.get_valid(df)
 
-    @log_telemetry("engine", "save_results_in_table")
+    @telemetry_logger("engine", "save_results_in_table")
     def save_results_in_table(
         self,
         output_df: DataFrame | None = None,
diff --git a/src/databricks/labs/dqx/telemetry.py b/src/databricks/labs/dqx/telemetry.py
index 809eeb3f2..91775ab26 100644
--- a/src/databricks/labs/dqx/telemetry.py
+++ b/src/databricks/labs/dqx/telemetry.py
@@ -1,6 +1,5 @@
 import logging
-from typing import Callable
-
+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
@@ -9,14 +8,14 @@
 logger = logging.getLogger(__name__)
 
 
-def trace(ws: WorkspaceClient, key: str, value: str) -> None:
+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: key to log
-        value: value to log
+        key: telemetry key to log
+        value: telemetry value to log
     """
     with_user_agent_extra(key, value)
     try:
@@ -26,11 +25,11 @@ def trace(ws: WorkspaceClient, key: str, value: str) -> None:
         logger.debug(f"Databricks workspace is not available: {e}")
 
 
-def log_telemetry(key: str, value: str) -> Callable:
+def telemetry_logger(key: str, value: str) -> Callable:
     """
-    Decorator to automatically log telemetry for method calls.
+    Decorator to log telemetry for method calls.
 
-    Usage: @log_telemetry("telemetry_key", "telemetry_value")
+    Usage: @telemetry_logger("telemetry_key", "telemetry_value")
 
     Args:
         key: Telemetry key to log
@@ -39,10 +38,13 @@ def log_telemetry(key: str, value: str) -> Callable:
 
     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
-                trace(self.ws, key, value)
+                log_telemetry(self.ws, key, value)
             elif hasattr(self, 'workspace_client'):  # requires workspace client to be set
-                trace(self.workspace_client, key, value)
+                log_telemetry(self.workspace_client, key, value)
             else:
                 raise AttributeError(
                     f"Workspace client not found on {self.__class__.__name__}. "

From 998645b5bd526e0ba2c81ec13e8b10bd46d8c123 Mon Sep 17 00:00:00 2001
From: Marcin Wojtyczka <marcin.wojtyczka@databricks.com>
Date: Thu, 4 Sep 2025 15:33:20 +0200
Subject: [PATCH 03/20] refactor

---
 src/databricks/labs/dqx/telemetry.py | 22 ++++++++++------------
 1 file changed, 10 insertions(+), 12 deletions(-)

diff --git a/src/databricks/labs/dqx/telemetry.py b/src/databricks/labs/dqx/telemetry.py
index 91775ab26..d4a05eee3 100644
--- a/src/databricks/labs/dqx/telemetry.py
+++ b/src/databricks/labs/dqx/telemetry.py
@@ -25,31 +25,29 @@ def log_telemetry(ws: WorkspaceClient, key: str, value: str) -> None:
         logger.debug(f"Databricks workspace is not available: {e}")
 
 
-def telemetry_logger(key: str, value: str) -> Callable:
+def telemetry_logger(key: str, value: str, workspace_client_attr: str = "ws") -> Callable:
     """
     Decorator to log telemetry for method calls.
 
-    Usage: @telemetry_logger("telemetry_key", "telemetry_value")
+    Usage: 
+        @telemetry_logger("telemetry_key", "telemetry_value")  # Uses "ws" by default
+        @telemetry_logger("telemetry_key", "telemetry_value", "my_client")  # Custom attribute
 
     Args:
         key: Telemetry key to log
         value: Telemetry value to log
+        workspace_client_attr: Name of the workspace client attribute on the class (defaults to "ws")
     """
 
     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)
+            if hasattr(self, workspace_client_attr):
+                workspace_client = getattr(self, workspace_client_attr)
+                log_telemetry(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."
+                    f"Workspace client attribute '{workspace_client_attr}' not found on {self.__class__.__name__}. "
+                    f"Make sure your class has the specified workspace client attribute."
                 )
             return func(self, *args, **kwargs)
 

From 140fc64b091d325cc426516c459b128d181da066 Mon Sep 17 00:00:00 2001
From: Marcin Wojtyczka <marcin.wojtyczka@databricks.com>
Date: Thu, 4 Sep 2025 16:16:08 +0200
Subject: [PATCH 04/20] refactor

---
 src/databricks/labs/dqx/telemetry.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/databricks/labs/dqx/telemetry.py b/src/databricks/labs/dqx/telemetry.py
index d4a05eee3..c254a11b9 100644
--- a/src/databricks/labs/dqx/telemetry.py
+++ b/src/databricks/labs/dqx/telemetry.py
@@ -29,9 +29,9 @@ def telemetry_logger(key: str, value: str, workspace_client_attr: str = "ws") ->
     """
     Decorator to log telemetry for method calls.
 
-    Usage: 
-        @telemetry_logger("telemetry_key", "telemetry_value")  # Uses "ws" by default
-        @telemetry_logger("telemetry_key", "telemetry_value", "my_client")  # Custom attribute
+    Usage:
+        @telemetry_logger("telemetry_key", "telemetry_value")  # Uses "ws" attribute for workspace client by default
+        @telemetry_logger("telemetry_key", "telemetry_value", "my_ws_client")  # Custom attribute
 
     Args:
         key: Telemetry key to log

From b0b85b90b89fba94062e249809d50586b3472abb Mon Sep 17 00:00:00 2001
From: Marcin Wojtyczka <marcin.wojtyczka@databricks.com>
Date: Mon, 15 Sep 2025 09:35:47 +0200
Subject: [PATCH 05/20] updated api call to support group assigned clusters

---
 src/databricks/labs/dqx/base.py      | 2 +-
 src/databricks/labs/dqx/telemetry.py | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/databricks/labs/dqx/base.py b/src/databricks/labs/dqx/base.py
index 4bf3b88e9..ede1b32ba 100644
--- a/src/databricks/labs/dqx/base.py
+++ b/src/databricks/labs/dqx/base.py
@@ -45,7 +45,7 @@ def _verify_workspace_client(ws: WorkspaceClient) -> WorkspaceClient:
             setattr(ws.config, '_product_info', ('dqx', __version__))
 
         # make sure Databricks workspace is accessible
-        ws.current_user.me()
+        ws.get_workspace_id()
         return ws
 
 
diff --git a/src/databricks/labs/dqx/telemetry.py b/src/databricks/labs/dqx/telemetry.py
index c254a11b9..d3b523139 100644
--- a/src/databricks/labs/dqx/telemetry.py
+++ b/src/databricks/labs/dqx/telemetry.py
@@ -19,7 +19,7 @@ def log_telemetry(ws: WorkspaceClient, key: str, value: str) -> None:
     """
     with_user_agent_extra(key, value)
     try:
-        ws.current_user.me()
+        ws.get_workspace_id()
     except DatabricksError as e:
         # support local execution
         logger.debug(f"Databricks workspace is not available: {e}")

From df541b1bad65605625b373a9fda5aba0460d7cd8 Mon Sep 17 00:00:00 2001
From: Marcin Wojtyczka <marcin.wojtyczka@databricks.com>
Date: Mon, 15 Sep 2025 11:41:57 +0200
Subject: [PATCH 06/20] updated logging logic

---
 src/databricks/labs/dqx/telemetry.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/src/databricks/labs/dqx/telemetry.py b/src/databricks/labs/dqx/telemetry.py
index d3b523139..c8689ec06 100644
--- a/src/databricks/labs/dqx/telemetry.py
+++ b/src/databricks/labs/dqx/telemetry.py
@@ -1,7 +1,6 @@
 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
 
 
@@ -17,7 +16,12 @@ def log_telemetry(ws: WorkspaceClient, key: str, value: str) -> None:
         key: telemetry key to log
         value: telemetry value to log
     """
-    with_user_agent_extra(key, value)
+    new_config = ws.config.copy().with_user_agent_extra(key, value)
+    logger.debug(f"Added User-Agent extra {key}={value}")
+
+    # Recreate the WorkspaceClient from the same type to preserve type information
+    ws = type(ws)(config=new_config)
+
     try:
         ws.get_workspace_id()
     except DatabricksError as e:

From eed9fc1169f8c2b52b3166b239489d4b28fdd6ca Mon Sep 17 00:00:00 2001
From: Marcin Wojtyczka <marcin.wojtyczka@databricks.com>
Date: Mon, 15 Sep 2025 11:56:37 +0200
Subject: [PATCH 07/20] support local execution

---
 docs/dqx/docs/reference/testing.mdx  | 2 +-
 src/databricks/labs/dqx/telemetry.py | 5 +++++
 2 files changed, 6 insertions(+), 1 deletion(-)

diff --git a/docs/dqx/docs/reference/testing.mdx b/docs/dqx/docs/reference/testing.mdx
index f469420ef..ddfe81c08 100644
--- a/docs/dqx/docs/reference/testing.mdx
+++ b/docs/dqx/docs/reference/testing.mdx
@@ -139,7 +139,7 @@ We strongly recommend following the standard testing procedure outlined above, w
 
     def test_dq():
       spark = SparkSession.builder.master("local[*]").getOrCreate() # create spark local session
-      ws = MagicMock(spec=WorkspaceClient, **{"current_user.me.return_value": None}) # mock the workspace client
+      ws = MagicMock(spec=WorkspaceClient, **{"get_workspace_id.return_value": 0}) # mock the workspace client
 
       schema = "a: int, b: int, c: int"
       expected_schema = schema + f", _errors: {dq_result_schema.simpleString()}, _warnings: {dq_result_schema.simpleString()}"
diff --git a/src/databricks/labs/dqx/telemetry.py b/src/databricks/labs/dqx/telemetry.py
index c8689ec06..80ff3c3b9 100644
--- a/src/databricks/labs/dqx/telemetry.py
+++ b/src/databricks/labs/dqx/telemetry.py
@@ -16,6 +16,11 @@ def log_telemetry(ws: WorkspaceClient, key: str, value: str) -> None:
         key: telemetry key to log
         value: telemetry value to log
     """
+    if not hasattr(ws, "config"):
+        # support local execution
+        logger.warning("Workspace client is not configured.")
+        return
+
     new_config = ws.config.copy().with_user_agent_extra(key, value)
     logger.debug(f"Added User-Agent extra {key}={value}")
 

From 7bd2d16ab3a6f6f1ca835db4591c33de0ccdf92f Mon Sep 17 00:00:00 2001
From: Marcin Wojtyczka <marcin.wojtyczka@databricks.com>
Date: Mon, 15 Sep 2025 11:56:51 +0200
Subject: [PATCH 08/20] support local execution

---
 src/databricks/labs/dqx/telemetry.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/databricks/labs/dqx/telemetry.py b/src/databricks/labs/dqx/telemetry.py
index 80ff3c3b9..e087ed84c 100644
--- a/src/databricks/labs/dqx/telemetry.py
+++ b/src/databricks/labs/dqx/telemetry.py
@@ -18,7 +18,7 @@ def log_telemetry(ws: WorkspaceClient, key: str, value: str) -> None:
     """
     if not hasattr(ws, "config"):
         # support local execution
-        logger.warning("Workspace client is not configured.")
+        logger.debug("Workspace client is not configured.")
         return
 
     new_config = ws.config.copy().with_user_agent_extra(key, value)

From bfa68ecd7b277a433fb1152b03de533d43bcddd5 Mon Sep 17 00:00:00 2001
From: Marcin Wojtyczka <marcin.wojtyczka@databricks.com>
Date: Mon, 15 Sep 2025 13:13:22 +0200
Subject: [PATCH 09/20] fmt

---
 docs/dqx/docs/reference/quality_checks.mdx | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/docs/dqx/docs/reference/quality_checks.mdx b/docs/dqx/docs/reference/quality_checks.mdx
index dc1066fa4..2b71d4544 100644
--- a/docs/dqx/docs/reference/quality_checks.mdx
+++ b/docs/dqx/docs/reference/quality_checks.mdx
@@ -2650,8 +2650,8 @@ This is because the metadata format relies on string representations that can't
 DQX provides both built-in PII detection and support for custom PII detection implementations. PII detection checks can be installed and run as extras:
 
 ```bash
-   pip install databricks-labs-dqx[pii]
-   ```
+pip install databricks-labs-dqx[pii]
+```
 
 <Admonition type="info" title="PII Detection Guarantees">
 PII detection using natural language models is non-deterministic. DQX cannot guarantee detection of all PII in an input dataset. While DQX makes a best-effort to detect PII in your data, a broader system should be designed to ensure data is cleansed of PII.

From 9a1566589429709aa7d65ae32af094e3dbc54bfa Mon Sep 17 00:00:00 2001
From: Marcin Wojtyczka <marcin.wojtyczka@databricks.com>
Date: Mon, 15 Sep 2025 16:26:14 +0200
Subject: [PATCH 10/20] enforce workspace client validation

---
 src/databricks/labs/dqx/base.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/databricks/labs/dqx/base.py b/src/databricks/labs/dqx/base.py
index ede1b32ba..62e6c9b38 100644
--- a/src/databricks/labs/dqx/base.py
+++ b/src/databricks/labs/dqx/base.py
@@ -12,7 +12,7 @@
 
 class DQEngineBase(abc.ABC):
     def __init__(self, workspace_client: WorkspaceClient):
-        self._workspace_client = workspace_client
+        self._workspace_client = self._verify_workspace_client(workspace_client)
         self._spark = SparkSession.builder.getOrCreate()
 
     @cached_property
@@ -22,7 +22,7 @@ def ws(self) -> WorkspaceClient:
         Ensures workspace connectivity and sets the product info used for
         telemetry so that requests are attributed to *dqx*.
         """
-        return self._verify_workspace_client(self._workspace_client)
+        return self._workspace_client
 
     @cached_property
     def spark(self) -> SparkSession:

From a76450534db2ecb487065ac614ff0c4d99e62c03 Mon Sep 17 00:00:00 2001
From: Marcin Wojtyczka <marcin.wojtyczka@databricks.com>
Date: Mon, 15 Sep 2025 17:19:17 +0200
Subject: [PATCH 11/20] enforce workspace client validation

---
 docs/dqx/docs/reference/testing.mdx  | 2 +-
 src/databricks/labs/dqx/base.py      | 3 ++-
 src/databricks/labs/dqx/telemetry.py | 3 ++-
 3 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/docs/dqx/docs/reference/testing.mdx b/docs/dqx/docs/reference/testing.mdx
index ddfe81c08..75471c0d2 100644
--- a/docs/dqx/docs/reference/testing.mdx
+++ b/docs/dqx/docs/reference/testing.mdx
@@ -139,7 +139,7 @@ We strongly recommend following the standard testing procedure outlined above, w
 
     def test_dq():
       spark = SparkSession.builder.master("local[*]").getOrCreate() # create spark local session
-      ws = MagicMock(spec=WorkspaceClient, **{"get_workspace_id.return_value": 0}) # mock the workspace client
+      ws = MagicMock(spec=WorkspaceClient, **{"clusters.select_spark_version.return_value": "mocked"}) # mock the workspace client
 
       schema = "a: int, b: int, c: int"
       expected_schema = schema + f", _errors: {dq_result_schema.simpleString()}, _warnings: {dq_result_schema.simpleString()}"
diff --git a/src/databricks/labs/dqx/base.py b/src/databricks/labs/dqx/base.py
index 62e6c9b38..ffbbbf9f8 100644
--- a/src/databricks/labs/dqx/base.py
+++ b/src/databricks/labs/dqx/base.py
@@ -45,7 +45,8 @@ def _verify_workspace_client(ws: WorkspaceClient) -> WorkspaceClient:
             setattr(ws.config, '_product_info', ('dqx', __version__))
 
         # make sure Databricks workspace is accessible
-        ws.get_workspace_id()
+        # use api that works on all workspaces and clusters including group assigned clusters
+        ws.clusters.select_spark_version(latest=True, long_term_support=True)
         return ws
 
 
diff --git a/src/databricks/labs/dqx/telemetry.py b/src/databricks/labs/dqx/telemetry.py
index e087ed84c..beab5ff77 100644
--- a/src/databricks/labs/dqx/telemetry.py
+++ b/src/databricks/labs/dqx/telemetry.py
@@ -28,7 +28,8 @@ def log_telemetry(ws: WorkspaceClient, key: str, value: str) -> None:
     ws = type(ws)(config=new_config)
 
     try:
-        ws.get_workspace_id()
+        # use api that works on all workspaces and clusters including group assigned clusters
+        ws.clusters.select_spark_version(latest=True, long_term_support=True)
     except DatabricksError as e:
         # support local execution
         logger.debug(f"Databricks workspace is not available: {e}")

From 91e15d1f36abd9cef43e3b00d053ce55144638ea Mon Sep 17 00:00:00 2001
From: Marcin Wojtyczka <marcin.wojtyczka@databricks.com>
Date: Mon, 15 Sep 2025 18:03:23 +0200
Subject: [PATCH 12/20] refactor, added unit tests

---
 docs/dqx/docs/reference/testing.mdx          |  4 ++--
 src/databricks/labs/dqx/base.py              | 12 +---------
 src/databricks/labs/dqx/profiler/profiler.py |  7 +++++-
 src/databricks/labs/dqx/telemetry.py         |  5 ----
 tests/conftest.py                            |  2 +-
 tests/unit/test_engine.py                    | 25 ++++++++++++++++++++
 6 files changed, 35 insertions(+), 20 deletions(-)
 create mode 100644 tests/unit/test_engine.py

diff --git a/docs/dqx/docs/reference/testing.mdx b/docs/dqx/docs/reference/testing.mdx
index 75471c0d2..87b3f0336 100644
--- a/docs/dqx/docs/reference/testing.mdx
+++ b/docs/dqx/docs/reference/testing.mdx
@@ -139,7 +139,7 @@ We strongly recommend following the standard testing procedure outlined above, w
 
     def test_dq():
       spark = SparkSession.builder.master("local[*]").getOrCreate() # create spark local session
-      ws = MagicMock(spec=WorkspaceClient, **{"clusters.select_spark_version.return_value": "mocked"}) # mock the workspace client
+      ws = MagicMock(spec=WorkspaceClient) # mock the workspace client
 
       schema = "a: int, b: int, c: int"
       expected_schema = schema + f", _errors: {dq_result_schema.simpleString()}, _warnings: {dq_result_schema.simpleString()}"
@@ -150,7 +150,7 @@ We strongly recommend following the standard testing procedure outlined above, w
         DQRowRule(name="col_b_is_null_or_empty", criticality="error", check_func=is_not_null_and_not_empty, column="b"),
       ]
 
-      dq_engine = DQEngine(ws)
+      dq_engine = DQEngine(spark=spark, workspace_client=ws)
       df = dq_engine.apply_checks(test_df, checks)
 
       expected_df = spark.createDataFrame(
diff --git a/src/databricks/labs/dqx/base.py b/src/databricks/labs/dqx/base.py
index ffbbbf9f8..3d7ec05ce 100644
--- a/src/databricks/labs/dqx/base.py
+++ b/src/databricks/labs/dqx/base.py
@@ -2,7 +2,7 @@
 from collections.abc import Callable
 from functools import cached_property
 from typing import final
-from pyspark.sql import DataFrame, SparkSession
+from pyspark.sql import DataFrame
 
 from databricks.labs.dqx.checks_validator import ChecksValidationStatus
 from databricks.labs.dqx.rule import DQRule
@@ -13,7 +13,6 @@
 class DQEngineBase(abc.ABC):
     def __init__(self, workspace_client: WorkspaceClient):
         self._workspace_client = self._verify_workspace_client(workspace_client)
-        self._spark = SparkSession.builder.getOrCreate()
 
     @cached_property
     def ws(self) -> WorkspaceClient:
@@ -24,15 +23,6 @@ def ws(self) -> WorkspaceClient:
         """
         return self._workspace_client
 
-    @cached_property
-    def spark(self) -> SparkSession:
-        """Return the *SparkSession* associated with this engine.
-
-        The session is created during initialization using
-        *SparkSession.builder.getOrCreate()*.
-        """
-        return self._spark
-
     @staticmethod
     @final
     def _verify_workspace_client(ws: WorkspaceClient) -> WorkspaceClient:
diff --git a/src/databricks/labs/dqx/profiler/profiler.py b/src/databricks/labs/dqx/profiler/profiler.py
index f61bbf146..502699344 100644
--- a/src/databricks/labs/dqx/profiler/profiler.py
+++ b/src/databricks/labs/dqx/profiler/profiler.py
@@ -13,7 +13,8 @@
 
 import pyspark.sql.functions as F
 import pyspark.sql.types as T
-from pyspark.sql import DataFrame
+from pyspark.sql import DataFrame, SparkSession
+from databricks.sdk import WorkspaceClient
 
 from databricks.labs.blueprint.limiter import rate_limited
 from databricks.labs.dqx.base import DQEngineBase
@@ -34,6 +35,10 @@ class DQProfile:
 class DQProfiler(DQEngineBase):
     """Data Quality Profiler class to profile input data."""
 
+    def __init__(self, workspace_client: WorkspaceClient, spark: SparkSession | None = None):
+        super().__init__(workspace_client=workspace_client)
+        self.spark = SparkSession.builder.getOrCreate() if spark is None else spark
+
     default_profile_options = {
         "round": True,  # round the min/max values
         "max_in_count": 10,  # generate is_in if we have less than 1 percent of distinct values
diff --git a/src/databricks/labs/dqx/telemetry.py b/src/databricks/labs/dqx/telemetry.py
index beab5ff77..741a2a2a8 100644
--- a/src/databricks/labs/dqx/telemetry.py
+++ b/src/databricks/labs/dqx/telemetry.py
@@ -16,11 +16,6 @@ def log_telemetry(ws: WorkspaceClient, key: str, value: str) -> None:
         key: telemetry key to log
         value: telemetry value to log
     """
-    if not hasattr(ws, "config"):
-        # support local execution
-        logger.debug("Workspace client is not configured.")
-        return
-
     new_config = ws.config.copy().with_user_agent_extra(key, value)
     logger.debug(f"Added User-Agent extra {key}={value}")
 
diff --git a/tests/conftest.py b/tests/conftest.py
index 8f044c429..7070e709d 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -22,7 +22,7 @@
 
 @pytest.fixture
 def debug_env_name():
-    return "ws2"  # Specify the name of the debug environment from ~/.databricks/debug-env.json
+    return "ws"  # Specify the name of the debug environment from ~/.databricks/debug-env.json
 
 
 @pytest.fixture
diff --git a/tests/unit/test_engine.py b/tests/unit/test_engine.py
new file mode 100644
index 000000000..7be6e97d5
--- /dev/null
+++ b/tests/unit/test_engine.py
@@ -0,0 +1,25 @@
+from unittest.mock import create_autospec
+from unittest.mock import MagicMock
+import pytest
+from pyspark.sql import SparkSession
+from databricks.sdk.errors import DatabricksError
+from databricks.sdk import WorkspaceClient
+from databricks.labs.dqx.engine import DQEngine, DQEngineCore
+
+
+def test_engine_creation():
+    spark_mock = create_autospec(SparkSession)
+    ws = create_autospec(WorkspaceClient)
+    assert DQEngine(spark=spark_mock, workspace_client=ws)
+    assert DQEngineCore(spark=spark_mock, workspace_client=ws)
+
+
+def test_engine_creation_no_workspace_connection():
+    spark_mock = create_autospec(SparkSession)
+    ws = MagicMock(spec=WorkspaceClient)
+    ws.clusters.select_spark_version.side_effect = DatabricksError()
+
+    with pytest.raises(DatabricksError):
+        DQEngine(spark=spark_mock, workspace_client=ws)
+    with pytest.raises(DatabricksError):
+        DQEngineCore(spark=spark_mock, workspace_client=ws)

From 2d6fece5bdd747d744a00048a90394dac8161245 Mon Sep 17 00:00:00 2001
From: Marcin Wojtyczka <marcin.wojtyczka@databricks.com>
Date: Mon, 15 Sep 2025 19:02:24 +0200
Subject: [PATCH 13/20] made spark connect optional to support local execution
 and testing

---
 docs/dqx/docs/reference/testing.mdx  | 47 ++++++++++++++--------------
 src/databricks/labs/dqx/base.py      |  2 +-
 src/databricks/labs/dqx/telemetry.py |  2 +-
 src/databricks/labs/dqx/utils.py     | 24 +++++++++++---
 tests/unit/test_load_checks.py       |  4 +--
 5 files changed, 46 insertions(+), 33 deletions(-)

diff --git a/docs/dqx/docs/reference/testing.mdx b/docs/dqx/docs/reference/testing.mdx
index 87b3f0336..cb00db904 100644
--- a/docs/dqx/docs/reference/testing.mdx
+++ b/docs/dqx/docs/reference/testing.mdx
@@ -114,14 +114,21 @@ To run the integration tests on serverless compute, add the `DATABRICKS_SERVERLE
 ```
 When `DATABRICKS_SERVERLESS_COMPUTE_ID` is set, the `DATABRICKS_CLUSTER_ID` is ignored, and tests run on serverless compute.
 
-### Local testing with DQEngine
+### Local execution and testing with DQEngine
 
-If workspace-level access is unavailable in your testing environment, you can perform local testing by installing the latest `pyspark` package and mocking the workspace client.
-Below is an example test.
+If workspace-level access is unavailable in your environment, you can execute DQX locally without workspace connection by installing the latest `pyspark` package and mocking the workspace client:
+```
+# user latest dqx version
+pip install databricks-labs-dqx pyspark==3.5.0
+
+# or use local build
+hatch build
+pip install dist/databricks_labs_dqx-0.9.2-py3-none-any.whl pyspark==3.5.0
+```
 
 <Admonition type="tip" title="Usage tips">
 **This approach should be treated as experimental!** It does not offer the same level of testing as the standard approach, and it is only applicable to selected methods (see [here](/docs/reference/engine/#dqx-engine-methods)).
-We strongly recommend following the standard testing procedure outlined above, which includes proper initialization of the workspace client.
+We strongly recommend following the standard execution procedure outlined above, which includes proper initialization of the workspace client.
 </Admonition>
 
 <Tabs>
@@ -130,33 +137,25 @@ We strongly recommend following the standard testing procedure outlined above, w
     from unittest.mock import MagicMock
     from databricks.sdk import WorkspaceClient
     from pyspark.sql import SparkSession
-    from chispa.dataframe_comparer import assert_df_equality
-    from databricks.labs.dqx.check_funcs import is_not_null_and_not_empty
+    from databricks.labs.dqx.check_funcs import is_not_null
     from databricks.labs.dqx.engine import DQEngine
     from databricks.labs.dqx.rule import DQRowRule
-    from databricks.labs.dqx.schema import dq_result_schema
-
 
-    def test_dq():
-      spark = SparkSession.builder.master("local[*]").getOrCreate() # create spark local session
-      ws = MagicMock(spec=WorkspaceClient) # mock the workspace client
+    spark = SparkSession.builder.master("local[*]").getOrCreate() # create spark local session
+    ws = MagicMock(spec=WorkspaceClient) # mock the workspace client
 
-      schema = "a: int, b: int, c: int"
-      expected_schema = schema + f", _errors: {dq_result_schema.simpleString()}, _warnings: {dq_result_schema.simpleString()}"
-      test_df = spark.createDataFrame([[1, None, 3]], schema)
+    schema = "a: int, b: int, c: int"
+    test_df = spark.createDataFrame([[1, None, 3]], schema)
 
-      checks = [
-        DQRowRule(name="col_a_is_null_or_empty", criticality="warn", check_func=is_not_null_and_not_empty, column="a"),
-        DQRowRule(name="col_b_is_null_or_empty", criticality="error", check_func=is_not_null_and_not_empty, column="b"),
-      ]
+    checks = [
+      DQRowRule(name="col_a_is_null", criticality="warn", check_func=is_not_null, column="a"),
+      DQRowRule(name="col_b_is_null", criticality="error", check_func=is_not_null, column="b"),
+    ]
 
-      dq_engine = DQEngine(spark=spark, workspace_client=ws)
-      df = dq_engine.apply_checks(test_df, checks)
+    dq_engine = DQEngine(spark=spark, workspace_client=ws)
+    df = dq_engine.apply_checks(test_df, checks)
 
-      expected_df = spark.createDataFrame(
-        [[1, None, 3, {"b_is_null_or_empty": "Column b is null or empty"}, None]], expected_schema
-      )
-      assert_df_equality(df, expected_df)
+    df.show(truncate=False)
     ```
   </TabItem>
 </Tabs>
diff --git a/src/databricks/labs/dqx/base.py b/src/databricks/labs/dqx/base.py
index 3d7ec05ce..3da955be2 100644
--- a/src/databricks/labs/dqx/base.py
+++ b/src/databricks/labs/dqx/base.py
@@ -36,7 +36,7 @@ def _verify_workspace_client(ws: WorkspaceClient) -> WorkspaceClient:
 
         # make sure Databricks workspace is accessible
         # use api that works on all workspaces and clusters including group assigned clusters
-        ws.clusters.select_spark_version(latest=True, long_term_support=True)
+        ws.clusters.select_spark_version()
         return ws
 
 
diff --git a/src/databricks/labs/dqx/telemetry.py b/src/databricks/labs/dqx/telemetry.py
index 741a2a2a8..71fe3ee04 100644
--- a/src/databricks/labs/dqx/telemetry.py
+++ b/src/databricks/labs/dqx/telemetry.py
@@ -24,7 +24,7 @@ def log_telemetry(ws: WorkspaceClient, key: str, value: str) -> None:
 
     try:
         # use api that works on all workspaces and clusters including group assigned clusters
-        ws.clusters.select_spark_version(latest=True, long_term_support=True)
+        ws.clusters.select_spark_version()
     except DatabricksError as e:
         # support local execution
         logger.debug(f"Databricks workspace is not available: {e}")
diff --git a/src/databricks/labs/dqx/utils.py b/src/databricks/labs/dqx/utils.py
index a6424fc17..395143719 100644
--- a/src/databricks/labs/dqx/utils.py
+++ b/src/databricks/labs/dqx/utils.py
@@ -6,7 +6,16 @@
 
 from pyspark.sql import Column, SparkSession
 from pyspark.sql.dataframe import DataFrame
-from pyspark.sql.connect.column import Column as ConnectColumn
+
+# Import spark connect column if available (e.g. integration testing)
+try:
+    from pyspark.sql.connect.column import Column as ConnectColumn
+
+    _HAS_SPARK_CONNECT = True
+except ImportError:
+    ConnectColumn = None  # type: ignore
+    _HAS_SPARK_CONNECT = False
+
 from databricks.labs.dqx.config import InputConfig, OutputConfig
 
 logger = logging.getLogger(__name__)
@@ -21,7 +30,7 @@
 
 
 def get_column_name_or_alias(
-    column: str | Column | ConnectColumn, normalize: bool = False, allow_simple_expressions_only: bool = False
+    column: "str | Column | ConnectColumn", normalize: bool = False, allow_simple_expressions_only: bool = False
 ) -> str:
     """
     Extracts the column alias or name from a PySpark Column or ConnectColumn expression.
@@ -34,7 +43,7 @@ def get_column_name_or_alias(
     - Provides an optional normalization step for consistent naming.
 
     Args:
-        column: Column, ConnectColumn or string representing a column.
+        column: Column or string representing a column.
         normalize: If True, normalizes the column name (removes special characters, converts to lowercase).
         allow_simple_expressions_only: If True, raises an error if the column expression is not a simple expression.
             Complex PySpark expressions (e.g., conditionals, arithmetic, or nested transformations), cannot be fully
@@ -77,7 +86,7 @@ def get_columns_as_strings(columns: list[str | Column], allow_simple_expressions
     This function processes each column, ensuring that only valid column names are returned.
 
     Args:
-        columns: List of columns, ConnectColumns or strings representing columns.
+        columns: List of columns or strings representing columns.
         allow_simple_expressions_only: If True, raises an error if the column expression is not a simple expression.
 
     Returns:
@@ -135,7 +144,12 @@ def normalize_bound_args(val: Any) -> Any:
     if isinstance(val, (datetime.date, datetime.datetime)):
         return str(val)
 
-    if isinstance(val, (Column, ConnectColumn)):
+    if _HAS_SPARK_CONNECT and ConnectColumn is not None:
+        column_types: tuple[type[Any], ...] = (Column, ConnectColumn)
+    else:
+        column_types = (Column,)
+
+    if isinstance(val, column_types):
         col_str = get_column_name_or_alias(val, allow_simple_expressions_only=True)
         return col_str
     raise TypeError(f"Unsupported type for normalization: {type(val).__name__}")
diff --git a/tests/unit/test_load_checks.py b/tests/unit/test_load_checks.py
index d88256144..c5aaa10ed 100644
--- a/tests/unit/test_load_checks.py
+++ b/tests/unit/test_load_checks.py
@@ -63,7 +63,7 @@ def test_load_checks_from_local_file_exceptions(filename, expected_exception, ex
 
 
 def test_file_download_contents_none():
-    ws = create_autospec(WorkspaceClient)
+    ws = create_autospec(spec=WorkspaceClient)
     handler = VolumeFileChecksStorageHandler(ws)
     # Simulate file_download.contents being None
     ws.files.download.return_value.contents = None
@@ -73,7 +73,7 @@ def test_file_download_contents_none():
 
 def test_file_download_contents_read_none():
     # Simulate file_download.contents.read() returning None
-    ws = create_autospec(WorkspaceClient)
+    ws = create_autospec(spec=WorkspaceClient)
     handler = VolumeFileChecksStorageHandler(ws)
 
     mock_file_download = create_autospec(DownloadResponse, instance=True)

From 406a910b0f3522a606c6f69aa82d79acbeefb061 Mon Sep 17 00:00:00 2001
From: Marcin Wojtyczka <marcin.wojtyczka@databricks.com>
Date: Mon, 15 Sep 2025 19:09:56 +0200
Subject: [PATCH 14/20] fmt

---
 docs/dqx/docs/reference/engine.mdx | 4 ++--
 tests/unit/test_load_checks.py     | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/docs/dqx/docs/reference/engine.mdx b/docs/dqx/docs/reference/engine.mdx
index 76d238977..84ab7e441 100644
--- a/docs/dqx/docs/reference/engine.mdx
+++ b/docs/dqx/docs/reference/engine.mdx
@@ -41,7 +41,7 @@ spark = DatabricksSession.builder.getOrCreate()
 dq_engine = DQEngine(ws, spark)
 ```
 
-For local execution without a Databricks workspace, please refer to the [local testing section](/docs/reference/testing/#local-testing-with-dqengine).
+For local execution without a Databricks workspace, please refer to the [local testing section](/docs/reference/testing/#local-execution-and-testing-with-dqengine).
 
 ## DQX engine methods
 
@@ -64,7 +64,7 @@ The following table outlines the available methods of the `DQEngine` and their f
 | `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)).
+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-execution-and-testing-with-dqengine)).
 
 `InputConfig` support the following parameters:
 * `location`: The location of the input data source (e.g. table name or file path).
diff --git a/tests/unit/test_load_checks.py b/tests/unit/test_load_checks.py
index c5aaa10ed..d88256144 100644
--- a/tests/unit/test_load_checks.py
+++ b/tests/unit/test_load_checks.py
@@ -63,7 +63,7 @@ def test_load_checks_from_local_file_exceptions(filename, expected_exception, ex
 
 
 def test_file_download_contents_none():
-    ws = create_autospec(spec=WorkspaceClient)
+    ws = create_autospec(WorkspaceClient)
     handler = VolumeFileChecksStorageHandler(ws)
     # Simulate file_download.contents being None
     ws.files.download.return_value.contents = None
@@ -73,7 +73,7 @@ def test_file_download_contents_none():
 
 def test_file_download_contents_read_none():
     # Simulate file_download.contents.read() returning None
-    ws = create_autospec(spec=WorkspaceClient)
+    ws = create_autospec(WorkspaceClient)
     handler = VolumeFileChecksStorageHandler(ws)
 
     mock_file_download = create_autospec(DownloadResponse, instance=True)

From 0c822b0bfc605e33ddd353e74f70de079a584282 Mon Sep 17 00:00:00 2001
From: Marcin Wojtyczka <marcin.wojtyczka@databricks.com>
Date: Mon, 15 Sep 2025 19:35:35 +0200
Subject: [PATCH 15/20] updated docs, fmt

---
 docs/dqx/docs/reference/testing.mdx | 2 +-
 src/databricks/labs/dqx/utils.py    | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/docs/dqx/docs/reference/testing.mdx b/docs/dqx/docs/reference/testing.mdx
index cb00db904..f8117b11b 100644
--- a/docs/dqx/docs/reference/testing.mdx
+++ b/docs/dqx/docs/reference/testing.mdx
@@ -123,7 +123,7 @@ pip install databricks-labs-dqx pyspark==3.5.0
 
 # or use local build
 hatch build
-pip install dist/databricks_labs_dqx-0.9.2-py3-none-any.whl pyspark==3.5.0
+pip install dist/databricks_labs_dqx-0.9.2-py3-none-any.whl pyspark[sql]==3.5.6
 ```
 
 <Admonition type="tip" title="Usage tips">
diff --git a/src/databricks/labs/dqx/utils.py b/src/databricks/labs/dqx/utils.py
index 395143719..302863d2f 100644
--- a/src/databricks/labs/dqx/utils.py
+++ b/src/databricks/labs/dqx/utils.py
@@ -144,7 +144,7 @@ def normalize_bound_args(val: Any) -> Any:
     if isinstance(val, (datetime.date, datetime.datetime)):
         return str(val)
 
-    if _HAS_SPARK_CONNECT and ConnectColumn is not None:
+    if _HAS_SPARK_CONNECT:
         column_types: tuple[type[Any], ...] = (Column, ConnectColumn)
     else:
         column_types = (Column,)

From 743ad3e654b1b749612895cd8d15fb316530c701 Mon Sep 17 00:00:00 2001
From: Marcin Wojtyczka <marcin.wojtyczka@databricks.com>
Date: Mon, 15 Sep 2025 19:37:41 +0200
Subject: [PATCH 16/20] fmt

---
 src/databricks/labs/dqx/utils.py | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

diff --git a/src/databricks/labs/dqx/utils.py b/src/databricks/labs/dqx/utils.py
index 302863d2f..11282be47 100644
--- a/src/databricks/labs/dqx/utils.py
+++ b/src/databricks/labs/dqx/utils.py
@@ -10,11 +10,8 @@
 # Import spark connect column if available (e.g. integration testing)
 try:
     from pyspark.sql.connect.column import Column as ConnectColumn
-
-    _HAS_SPARK_CONNECT = True
 except ImportError:
     ConnectColumn = None  # type: ignore
-    _HAS_SPARK_CONNECT = False
 
 from databricks.labs.dqx.config import InputConfig, OutputConfig
 
@@ -144,7 +141,7 @@ def normalize_bound_args(val: Any) -> Any:
     if isinstance(val, (datetime.date, datetime.datetime)):
         return str(val)
 
-    if _HAS_SPARK_CONNECT:
+    if ConnectColumn is not None:
         column_types: tuple[type[Any], ...] = (Column, ConnectColumn)
     else:
         column_types = (Column,)

From 85917ad0c3cc6e8bf5f11dcb75759bf28294bf4f Mon Sep 17 00:00:00 2001
From: Marcin Wojtyczka <marcin.wojtyczka@databricks.com>
Date: Tue, 16 Sep 2025 08:52:34 +0200
Subject: [PATCH 17/20] fix copilot feedback

---
 src/databricks/labs/dqx/telemetry.py | 4 ++++
 src/databricks/labs/dqx/utils.py     | 8 +++++---
 2 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/src/databricks/labs/dqx/telemetry.py b/src/databricks/labs/dqx/telemetry.py
index 71fe3ee04..9336ef1fb 100644
--- a/src/databricks/labs/dqx/telemetry.py
+++ b/src/databricks/labs/dqx/telemetry.py
@@ -1,3 +1,4 @@
+import functools
 import logging
 from collections.abc import Callable
 from databricks.sdk import WorkspaceClient
@@ -33,6 +34,7 @@ def log_telemetry(ws: WorkspaceClient, key: str, value: str) -> None:
 def telemetry_logger(key: str, value: str, workspace_client_attr: str = "ws") -> Callable:
     """
     Decorator to log telemetry for method calls.
+    By default, it expects the decorated method to have "ws" attribute for workspace client.
 
     Usage:
         @telemetry_logger("telemetry_key", "telemetry_value")  # Uses "ws" attribute for workspace client by default
@@ -45,6 +47,8 @@ def telemetry_logger(key: str, value: str, workspace_client_attr: str = "ws") ->
     """
 
     def decorator(func: Callable) -> Callable:
+
+        @functools.wraps(func)  # preserve function metadata
         def wrapper(self, *args, **kwargs):
             if hasattr(self, workspace_client_attr):
                 workspace_client = getattr(self, workspace_client_attr)
diff --git a/src/databricks/labs/dqx/utils.py b/src/databricks/labs/dqx/utils.py
index 11282be47..fd81e54c9 100644
--- a/src/databricks/labs/dqx/utils.py
+++ b/src/databricks/labs/dqx/utils.py
@@ -7,7 +7,7 @@
 from pyspark.sql import Column, SparkSession
 from pyspark.sql.dataframe import DataFrame
 
-# Import spark connect column if available (e.g. integration testing)
+# Import spark connect column if spark session is created using spark connect
 try:
     from pyspark.sql.connect.column import Column as ConnectColumn
 except ImportError:
@@ -38,9 +38,10 @@ def get_column_name_or_alias(
     - Supports columns with one or multiple aliases.
     - Ensures the extracted expression is truncated to 255 characters.
     - Provides an optional normalization step for consistent naming.
+    - Supports ConnectColumn when PySpark Connect is available (falls back gracefully when not available).
 
     Args:
-        column: Column or string representing a column.
+        column: Column, ConnectColumn (if PySpark Connect available), or string representing a column.
         normalize: If True, normalizes the column name (removes special characters, converts to lowercase).
         allow_simple_expressions_only: If True, raises an error if the column expression is not a simple expression.
             Complex PySpark expressions (e.g., conditionals, arithmetic, or nested transformations), cannot be fully
@@ -81,9 +82,10 @@ def get_columns_as_strings(columns: list[str | Column], allow_simple_expressions
     Extracts column names from a list of PySpark Column or ConnectColumn expressions.
 
     This function processes each column, ensuring that only valid column names are returned.
+    Supports ConnectColumn when PySpark Connect is available (falls back gracefully when not available).
 
     Args:
-        columns: List of columns or strings representing columns.
+        columns: List of columns, ConnectColumns (if PySpark Connect available), or strings representing columns.
         allow_simple_expressions_only: If True, raises an error if the column expression is not a simple expression.
 
     Returns:

From 26a7c7f966412921bfb81b254f273dc315b0b76a Mon Sep 17 00:00:00 2001
From: Marcin Wojtyczka <marcin.wojtyczka@databricks.com>
Date: Tue, 16 Sep 2025 09:05:14 +0200
Subject: [PATCH 18/20] updated docs

---
 docs/dqx/docs/reference/testing.mdx | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docs/dqx/docs/reference/testing.mdx b/docs/dqx/docs/reference/testing.mdx
index f8117b11b..35fe0e8fb 100644
--- a/docs/dqx/docs/reference/testing.mdx
+++ b/docs/dqx/docs/reference/testing.mdx
@@ -123,7 +123,7 @@ pip install databricks-labs-dqx pyspark==3.5.0
 
 # or use local build
 hatch build
-pip install dist/databricks_labs_dqx-0.9.2-py3-none-any.whl pyspark[sql]==3.5.6
+pip install dist/databricks_labs_dqx-0.9.2-py3-none-any.whl pyspark==3.5.6
 ```
 
 <Admonition type="tip" title="Usage tips">

From e43727412d36926ddb86e47d248d8776651fc86a Mon Sep 17 00:00:00 2001
From: Marcin Wojtyczka <marcin.wojtyczka@databricks.com>
Date: Tue, 16 Sep 2025 09:57:16 +0200
Subject: [PATCH 19/20] updated docs

---
 docs/dqx/docs/reference/testing.mdx | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/docs/dqx/docs/reference/testing.mdx b/docs/dqx/docs/reference/testing.mdx
index 35fe0e8fb..2b6ca9875 100644
--- a/docs/dqx/docs/reference/testing.mdx
+++ b/docs/dqx/docs/reference/testing.mdx
@@ -126,6 +126,10 @@ hatch build
 pip install dist/databricks_labs_dqx-0.9.2-py3-none-any.whl pyspark==3.5.6
 ```
 
+<Admonition type="warning" title="Python requirements">
+You need a python version that is compatible with the pyspark version, e.g. Python 3.9, 3.10, 3.11, and 3.12 are compatible with PySpark 3.5, but Python 3.13 is not supported.
+</Admonition>
+
 <Admonition type="tip" title="Usage tips">
 **This approach should be treated as experimental!** It does not offer the same level of testing as the standard approach, and it is only applicable to selected methods (see [here](/docs/reference/engine/#dqx-engine-methods)).
 We strongly recommend following the standard execution procedure outlined above, which includes proper initialization of the workspace client.
@@ -159,5 +163,3 @@ We strongly recommend following the standard execution procedure outlined above,
     ```
   </TabItem>
 </Tabs>
-
-

From d98551c324a256acf4270a0d962343c633d98105 Mon Sep 17 00:00:00 2001
From: Marcin Wojtyczka <marcin.wojtyczka@databricks.com>
Date: Tue, 16 Sep 2025 10:27:44 +0200
Subject: [PATCH 20/20] added more telemetry

---
 src/databricks/labs/dqx/contexts/workflow_context.py | 4 +++-
 src/databricks/labs/dqx/profiler/dlt_generator.py    | 2 ++
 src/databricks/labs/dqx/profiler/generator.py        | 2 ++
 src/databricks/labs/dqx/profiler/profiler.py         | 4 ++++
 4 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/src/databricks/labs/dqx/contexts/workflow_context.py b/src/databricks/labs/dqx/contexts/workflow_context.py
index 4adb1f28a..eba262e11 100644
--- a/src/databricks/labs/dqx/contexts/workflow_context.py
+++ b/src/databricks/labs/dqx/contexts/workflow_context.py
@@ -13,6 +13,7 @@
 from databricks.labs.dqx.profiler.profiler import DQProfiler
 from databricks.labs.dqx.profiler.profiler_runner import ProfilerRunner
 from databricks.labs.dqx.quality_checker.quality_checker_runner import QualityCheckerRunner
+from databricks.labs.dqx.telemetry import log_telemetry
 
 
 class WorkflowContext(GlobalContext):
@@ -75,7 +76,7 @@ def profiler(self) -> ProfilerRunner:
         dq_engine = DQEngine(
             workspace_client=self.workspace_client, spark=self.spark, extra_params=self.config.extra_params
         )
-
+        log_telemetry(self.workspace_client, "workflow", "profiler")
         return ProfilerRunner(
             self.workspace_client,
             self.spark,
@@ -91,4 +92,5 @@ def quality_checker(self) -> QualityCheckerRunner:
         dq_engine = DQEngine(
             workspace_client=self.workspace_client, spark=self.spark, extra_params=self.config.extra_params
         )
+        log_telemetry(self.workspace_client, "workflow", "quality_checker")
         return QualityCheckerRunner(self.spark, dq_engine)
diff --git a/src/databricks/labs/dqx/profiler/dlt_generator.py b/src/databricks/labs/dqx/profiler/dlt_generator.py
index cc8dc49d4..c552fc062 100644
--- a/src/databricks/labs/dqx/profiler/dlt_generator.py
+++ b/src/databricks/labs/dqx/profiler/dlt_generator.py
@@ -5,12 +5,14 @@
 from databricks.labs.dqx.base import DQEngineBase
 from databricks.labs.dqx.profiler.common import val_to_str
 from databricks.labs.dqx.profiler.profiler import DQProfile
+from databricks.labs.dqx.telemetry import telemetry_logger
 
 __name_sanitize_re__ = re.compile(r"[^a-zA-Z0-9]+")
 logger = logging.getLogger(__name__)
 
 
 class DQDltGenerator(DQEngineBase):
+    @telemetry_logger("generator", "generate_dlt_rules")
     def generate_dlt_rules(
         self, rules: list[DQProfile], action: str | None = None, language: str = "SQL"
     ) -> list[str] | str | dict:
diff --git a/src/databricks/labs/dqx/profiler/generator.py b/src/databricks/labs/dqx/profiler/generator.py
index 0e0e0a4ca..5ca43681d 100644
--- a/src/databricks/labs/dqx/profiler/generator.py
+++ b/src/databricks/labs/dqx/profiler/generator.py
@@ -4,11 +4,13 @@
 from databricks.labs.dqx.engine import DQEngine
 from databricks.labs.dqx.profiler.common import val_maybe_to_str
 from databricks.labs.dqx.profiler.profiler import DQProfile
+from databricks.labs.dqx.telemetry import telemetry_logger
 
 logger = logging.getLogger(__name__)
 
 
 class DQGenerator(DQEngineBase):
+    @telemetry_logger("generator", "generate_dq_rules")
     def generate_dq_rules(self, profiles: list[DQProfile] | None = None, level: str = "error") -> list[dict]:
         """
         Generates a list of data quality rules based on the provided dq profiles.
diff --git a/src/databricks/labs/dqx/profiler/profiler.py b/src/databricks/labs/dqx/profiler/profiler.py
index 502699344..916fdbd33 100644
--- a/src/databricks/labs/dqx/profiler/profiler.py
+++ b/src/databricks/labs/dqx/profiler/profiler.py
@@ -19,6 +19,7 @@
 from databricks.labs.blueprint.limiter import rate_limited
 from databricks.labs.dqx.base import DQEngineBase
 from databricks.labs.dqx.config import InputConfig
+from databricks.labs.dqx.telemetry import telemetry_logger
 from databricks.labs.dqx.utils import read_input_data
 
 logger = logging.getLogger(__name__)
@@ -76,6 +77,7 @@ def get_columns_or_fields(columns: list[T.StructField]) -> list[T.StructField]:
         return out_columns
 
     # TODO: how to handle maps, arrays & structs?
+    @telemetry_logger("profiler", "profile")
     def profile(
         self, df: DataFrame, columns: list[str] | None = None, options: dict[str, Any] | None = None
     ) -> tuple[dict[str, Any], list[DQProfile]]:
@@ -110,6 +112,7 @@ def profile(
 
         return summary_stats, dq_rules
 
+    @telemetry_logger("profiler", "profile_table")
     def profile_table(
         self,
         table: str,
@@ -131,6 +134,7 @@ def profile_table(
         df = read_input_data(spark=self.spark, input_config=InputConfig(location=table))
         return self.profile(df=df, columns=columns, options=options)
 
+    @telemetry_logger("profiler", "profile_tables")
     def profile_tables(
         self,
         tables: list[str] | None = None,