diff --git a/pyproject.toml b/pyproject.toml index 4590663d2..9e3089c9e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -37,9 +37,9 @@ classifiers = [ # also update fallback dependencies in workflow installer when changing dependencies! dependencies = [ "databricks-labs-blueprint>=0.9.1,<0.10", - "databricks-sdk~=0.71", - "databricks-labs-lsql>=0.5,<=0.16", + "databricks-sdk~=0.73", "sqlalchemy>=2.0,<3.0", + "PyYAML~=6.0.3", ] [project.optional-dependencies] @@ -111,7 +111,7 @@ dependencies = [ "pytest-xdist~=3.5.0", "pytest-benchmark~=5.1.0", "ruff~=0.3.4", - "types-PyYAML~=6.0.12", + "types-PyYAML~=6.0.3", "types-requests~=2.31.0", "databricks-connect~=15.4", "pyspark~=3.5.0", diff --git a/src/databricks/labs/dqx/dashboards/dashboard.lvdash.json b/src/databricks/labs/dqx/dashboards/dashboard.lvdash.json new file mode 100644 index 000000000..9848f60bb --- /dev/null +++ b/src/databricks/labs/dqx/dashboards/dashboard.lvdash.json @@ -0,0 +1,208 @@ +{ + "datasets": [ + { + "name": "00_1_dq_summary", + "displayName": "00_1_dq_summary", + "queryLines": [ + "/* --title 'Data Quality Summary' */\n", + "SELECT \n", + " CASE\n", + " WHEN _errors IS NOT NULL THEN 'Errors'\n", + " WHEN _warnings IS NOT NULL THEN 'Warnings'\n", + " END AS category,\n", + " ROUND((COUNT(*) * 100.0 / (SELECT COUNT(*) FROM $catalog.schema.table)), 2) AS percentage\n", + "FROM $catalog.schema.table\n", + "GROUP BY category" + ] + }, + { + "name": "00_2_dq_error_types", + "displayName": "00_2_dq_error_types", + "queryLines": [ + "/* --title 'Error and Warning Types Breakdown' */\n", + "WITH error_types AS (\n", + " SELECT\n", + " 'Error' AS category,\n", + " error_struct.name AS type,\n", + " COUNT(*) AS count\n", + " FROM $catalog.schema.table\n", + " LATERAL VIEW EXPLODE(_errors) exploded_errors AS error_struct\n", + " WHERE _errors IS NOT NULL\n", + " GROUP BY error_struct.name\n", + "),\n", + "warning_types AS (\n", + " SELECT\n", + " 'Warning' AS category,\n", + " warning_struct.name AS type,\n", + " COUNT(*) AS count\n", + " FROM $catalog.schema.table\n", + " LATERAL VIEW EXPLODE(_warnings) exploded_warnings AS warning_struct\n", + " WHERE _warnings IS NOT NULL\n", + " GROUP BY warning_struct.name\n", + "),\n", + "combined AS (\n", + " SELECT * FROM error_types\n", + " UNION ALL\n", + " SELECT * FROM warning_types\n", + "),\n", + "total AS (\n", + " SELECT SUM(count) AS total_count FROM combined\n", + ")\n", + "SELECT\n", + " category,\n", + " type,\n", + " count,\n", + " ROUND((count * 100.0) / total.total_count, 2) AS percentage\n", + "FROM combined, total\n", + "ORDER BY category, count DESC;" + ] + } + ], + "pages": [ + { + "name": "DQX_Quality_Dashboard", + "displayName": "DQX_Quality_Dashboard", + "layout": [ + { + "widget": { + "name": "00_1_dq_summary_widget", + "queries": [ + { + "name": "main_query", + "query": { + "datasetName": "00_1_dq_summary", + "fields": [ + { + "name": "sum(percentage)", + "expression": "SUM(`percentage`)" + }, + { + "name": "category", + "expression": "`category`" + } + ], + "disaggregated": false + } + } + ], + "spec": { + "version": 3, + "widgetType": "pie", + "encodings": { + "angle": { + "displayName": "Total Row Counts", + "fieldName": "sum(percentage)", + "scale": { + "type": "quantitative" + } + }, + "color": { + "displayName": "category", + "fieldName": "category", + "scale": { + "type": "categorical", + "mappings": [ + { + "color": "#00A972" + }, + { + "color": "#FFAB00" + }, + { + "color": "#FF3621" + } + ] + } + }, + "label": { + "show": true + } + }, + "frame": { + "description": "", + "showDescription": false, + "showTitle": true, + "title": "Data Quality Summary" + } + } + }, + "position": { + "x": 0, + "y": 0, + "width": 3, + "height": 6 + } + }, + { + "widget": { + "name": "00_2_dq_error_types_widget", + "queries": [ + { + "name": "main_query", + "query": { + "datasetName": "00_2_dq_error_types", + "fields": [ + { + "name": "Category", + "expression": "`category`" + }, + { + "name": "Type", + "expression": "`type`" + }, + { + "name": "Count", + "expression": "`count`" + }, + { + "name": "Percentage", + "expression": "`percentage`" + } + ], + "disaggregated": true + } + } + ], + "spec": { + "version": 2, + "widgetType": "table", + "encodings": { + "columns": [ + { + "displayName": "Category", + "fieldName": "Category" + }, + { + "displayName": "Type", + "fieldName": "Type" + }, + { + "displayName": "Count", + "fieldName": "Count" + }, + { + "displayName": "Percentage (%)", + "fieldName": "Percentage" + } + ] + }, + "frame": { + "description": "", + "showDescription": false, + "showTitle": true, + "title": "Error and Warning Types Breakdown" + } + } + }, + "position": { + "x": 3, + "y": 0, + "width": 3, + "height": 6 + } + } + ], + "pageType": "PAGE_TYPE_CANVAS" + } + ] +} \ No newline at end of file diff --git a/src/databricks/labs/dqx/installer/dashboard_installer.py b/src/databricks/labs/dqx/installer/dashboard_installer.py index d9173af64..006a7f526 100644 --- a/src/databricks/labs/dqx/installer/dashboard_installer.py +++ b/src/databricks/labs/dqx/installer/dashboard_installer.py @@ -1,16 +1,15 @@ import functools import logging -import os -import glob +import json +from typing import Any from collections.abc import Callable, Iterable from datetime import timedelta from pathlib import Path from databricks.labs.blueprint.installation import Installation from databricks.labs.blueprint.installer import InstallState from databricks.labs.blueprint.wheels import ProductInfo, find_project_root -from databricks.labs.lsql.dashboards import DashboardMetadata, Dashboards from databricks.sdk import WorkspaceClient -from databricks.sdk.service.dashboards import LifecycleState +from databricks.sdk.service.dashboards import LifecycleState, Dashboard from databricks.sdk.errors import ( InvalidParameterValue, NotFound, @@ -19,15 +18,53 @@ ResourceAlreadyExists, ) from databricks.sdk.retries import retried -from databricks.labs.dqx.config import WorkspaceConfig - +from databricks.labs.dqx.config import WorkspaceConfig, RunConfig logger = logging.getLogger(__name__) +class DashboardMetadata: + """Creates Dashboard Metadata from exported Lakeview dashboard (.lvdash.json)""" + + def __init__( + self, + display_name: str, + dashboard_def: dict[str, Any], + ): + self.display_name = display_name + self.dashboard_def = dashboard_def + + @classmethod + def from_path(cls, file: Path) -> "DashboardMetadata": + """Load dashboard metadata from exported Lakeview dashboard file. + + Expected structure: + dashboard_folder/ + └── dashboard_name.lvdash.json + + Args: + file: The path to the .lvdash.json file. + + Returns: + A DashboardMetadata instance populated with the display name and + dashboard definition from the exported Lakeview dashboard. + """ + logger.info(f"Loading exported Lakeview dashboard: {file.name}") + + with open(file, 'r', encoding="utf-8") as f: + dashboard_data = json.load(f) + + display_name = dashboard_data.get("displayName", file.stem) + + return cls( + display_name=display_name, + dashboard_def=dashboard_data, + ) + + class DashboardInstaller: """ - Creates or updates Lakeview dashboards from bundled SQL queries. + Creates or updates Lakeview dashboards from exported .lvdash.json files. """ def __init__( @@ -46,11 +83,10 @@ def __init__( def get_create_dashboard_tasks(self) -> Iterable[Callable[[], None]]: """ - Returns a generator of tasks to create dashboards from bundled SQL queries. + Returns a generator of tasks to create dashboards from exported Lakeview dashboards. Each task is a callable that, when executed, will create a dashboard in the workspace. - The tasks are created based on the SQL files found in the bundled queries directory. - The tasks will handle the creation of the dashboard, including resolving table names. + The tasks are created based on the .lvdash.json files found in the dashboard directory. """ logger.info("Creating dashboards...") dashboard_folder_remote = f"{self._installation.install_folder()}/dashboards" @@ -59,21 +95,18 @@ def get_create_dashboard_tasks(self) -> Iterable[Callable[[], None]]: except ResourceAlreadyExists: pass - queries_folder = find_project_root(__file__) / "src/databricks/labs/dqx/queries" - logger.debug(f"Dashboard Query Folder is {queries_folder}") - for step_folder in queries_folder.iterdir(): - if not step_folder.is_dir(): + dashboard_folder = find_project_root(__file__) / "src/databricks/labs/dqx/dashboards" + logger.debug(f"Dashboard Query Folder is {dashboard_folder}") + for step_file in dashboard_folder.iterdir(): + if not step_file.is_file(): continue - logger.debug(f"Reading step install folder {step_folder}...") - for dashboard_folder in step_folder.iterdir(): - if not dashboard_folder.is_dir(): - continue - task = functools.partial( - self._create_dashboard, - dashboard_folder, - parent_path=dashboard_folder_remote, - ) - yield task + logger.debug(f"Reading dashboard definition from {step_file}...") + task = functools.partial( + self._create_dashboard, + step_file, + parent_path=dashboard_folder_remote, + ) + yield task def _handle_existing_dashboard(self, dashboard_id: str, display_name: str, parent_path: str) -> str | None: """Handle an existing dashboard @@ -114,41 +147,50 @@ def _handle_existing_dashboard(self, dashboard_id: str, display_name: str, paren return dashboard_id # Update the existing dashboard @staticmethod - def _resolve_table_name_in_queries(src_tbl_name: str, replaced_tbl_name: str, folder: Path) -> bool: - """Replaces table name variable in all .sql files - This method iterates through the dashboard install_folder, and replaces fully qualified tables in *.sql files + def _resolve_table_name_in_dashboard( + src_tbl_name: str, replaced_tbl_name: str, dashboard_def: dict[str, Any] + ) -> dict[str, Any]: + """Replaces table name variable in dashboard definition + + This method replaces the placeholder table name with the actual table name + in all queries within the dashboard definition. Args: - src_tbl_name: The source table name to be replaced + src_tbl_name: The source table name to be replaced (placeholder) replaced_tbl_name: The table name to replace the source table name with - folder: The install_folder containing the SQL files + dashboard_def: The dashboard definition containing datasets with queries Returns: - True if the operation was successful, False otherwise + Updated dashboard definition with replaced table names """ - logger.debug("Preparing .sql files for DQX Dashboard") - dyn_sql_files = glob.glob(os.path.join(folder, "*.sql")) - try: - for sql_file in dyn_sql_files: - sql_file_path = Path(sql_file) - dq_sql_query = sql_file_path.read_text(encoding="utf-8") - dq_sql_query_ref = dq_sql_query.replace(src_tbl_name, replaced_tbl_name) - logger.debug(dq_sql_query_ref) - sql_file_path.write_text(dq_sql_query_ref, encoding="utf-8") - return True - except Exception as e: - err_msg = f"Error during parsing input table name into .sql files: {e}" - logger.error(err_msg) - # Review this - Gracefully handling this internal variable replace operation - return False - - # InternalError and DeadlineExceeded are retried because of Lakeview internal issues - # These issues have been reported to and are resolved by the Lakeview team. - # Keeping the retry for resilience. + logger.debug(f"Replacing '{src_tbl_name}' with '{replaced_tbl_name}' in dashboard queries") + + # Deep copy to avoid modifying original + updated_def = json.loads(json.dumps(dashboard_def)) + + # Replace table names in all datasets + for dataset in updated_def.get("datasets", []): + # Handle queryLines array (exported Lakeview format) + if "queryLines" in dataset and isinstance(dataset["queryLines"], list): + dataset["queryLines"] = [ + line.replace(src_tbl_name, replaced_tbl_name) for line in dataset["queryLines"] + ] + # Handle query string (alternative format) + elif "query" in dataset and isinstance(dataset["query"], str): + dataset["query"] = dataset["query"].replace(src_tbl_name, replaced_tbl_name) + + return updated_def + @retried(on=[InternalError, DeadlineExceeded], timeout=timedelta(minutes=4)) - def _create_dashboard(self, folder: Path, *, parent_path: str) -> None: - """Create a lakeview dashboard from the SQL queries in the install_folder""" - logger.info(f"Reading dashboard assets from {folder}...") + def _create_dashboard(self, file: Path, *, parent_path: str) -> None: + """ + Create a lakeview dashboard from the exported .lvdash.json file. + + Args: + file: Path to the .lvdash.json file + parent_path: Parent path where the dashboard will be created + """ + logger.info(f"Reading dashboard from {file}...") run_config = self._config.get_run_config() if run_config.quarantine_config: @@ -159,28 +201,125 @@ def _create_dashboard(self, folder: Path, *, parent_path: str) -> None: dq_table = run_config.output_config.location.lower() logger.info(f"Using '{dq_table}' output table as the source table for the dashboard...") + try: + self._create_dashboard_from_metadata(file, parent_path, run_config, dq_table) + except Exception as e: + logger.error(f"Failed to create dashboard from {file}: {e}", exc_info=True) + raise + + def _create_dashboard_from_metadata( + self, file: Path, parent_path: str, run_config: RunConfig, dq_table: str + ) -> None: + """ + Create dashboard using exported Lakeview metadata. + + Args: + file: Path to the .lvdash.json file + parent_path: Parent path where the dashboard will be created + run_config: Run configuration containing warehouse settings + dq_table: The actual table name to use in queries + """ + # Load metadata + metadata = self._prepare_dashboard_metadata(file) + reference = file.stem.lower() + dashboard_id = self._install_state.dashboards.get(reference) + + # Handle existing dashboard if needed + if dashboard_id is not None: + dashboard_id = self._handle_existing_dashboard(dashboard_id, metadata.display_name, parent_path) + + # Replace source table name in dashboard definition src_table_name = "$catalog.schema.table" - if self._resolve_table_name_in_queries(src_tbl_name=src_table_name, replaced_tbl_name=dq_table, folder=folder): - metadata = DashboardMetadata.from_path(folder) - logger.debug(f"Dashboard Metadata retrieved is {metadata}") - - metadata.display_name = f"DQX_{folder.parent.stem.title()}_{folder.stem.title()}" - reference = f"{folder.parent.stem}_{folder.stem}".lower() - dashboard_id = self._install_state.dashboards.get(reference) - logger.debug(f"dashboard id retrieved is {dashboard_id}") - - logger.info(f"Installing '{metadata.display_name}' dashboard in '{parent_path}'") - if dashboard_id is not None: - dashboard_id = self._handle_existing_dashboard(dashboard_id, metadata.display_name, parent_path) - dashboard = Dashboards(self._ws).create_dashboard( - metadata, - parent_path=parent_path, - dashboard_id=dashboard_id, - warehouse_id=run_config.warehouse_id, - publish=True, + updated_dashboard_def = self._resolve_table_name_in_dashboard( + src_tbl_name=src_table_name, replaced_tbl_name=dq_table, dashboard_def=metadata.dashboard_def + ) + + # Create or update the dashboard + dashboard = self._create_or_update_dashboard( + dashboard_id=dashboard_id, + metadata=metadata, + parent_path=parent_path, + run_config=run_config, + dashboard_def=updated_dashboard_def, + ) + + # Publish the dashboard + self._publish_dashboard(dashboard, metadata.display_name) + + # Save the dashboard reference + assert dashboard.dashboard_id is not None + self._install_state.dashboards[reference] = dashboard.dashboard_id + logger.info(f"Successfully installed dashboard: {metadata.display_name} ({dashboard.dashboard_id})") + + def _prepare_dashboard_metadata(self, file: Path) -> DashboardMetadata: + """ + Load and prepare dashboard metadata from dashboard file. + + Args: + file: Path to the .lvdash.json file + + Returns: + Prepared dashboard metadata with formatted display name + """ + metadata = DashboardMetadata.from_path(file) + logger.debug(f"Dashboard Metadata retrieved: {metadata.display_name}") + stem = file.stem.title() + if stem.lower().endswith(".lvdash"): + stem = stem[: -len(".lvdash")] + metadata.display_name = f"DQX_{stem}" + return metadata + + def _create_or_update_dashboard( + self, + dashboard_id: str | None, + metadata: DashboardMetadata, + parent_path: str, + run_config: RunConfig, + dashboard_def: dict[str, Any], + ) -> Dashboard: + """ + Create a new dashboard or update an existing one. + + Args: + dashboard_id: Existing dashboard ID if updating, None if creating new + metadata: Dashboard metadata configuration + parent_path: Parent path where the dashboard will be created + run_config: Run configuration containing warehouse settings + dashboard_def: The dashboard definition with resolved table names + + Returns: + Created or updated Dashboard object + """ + if dashboard_id is None: + logger.info(f"Creating new dashboard: {metadata.display_name}") + return self._ws.lakeview.create( + dashboard=Dashboard( + display_name=metadata.display_name, + parent_path=parent_path, + warehouse_id=run_config.warehouse_id, + serialized_dashboard=json.dumps(dashboard_def), + ) ) - assert dashboard.dashboard_id is not None - self._install_state.dashboards[reference] = dashboard.dashboard_id + logger.info(f"Updating existing dashboard: {metadata.display_name} ({dashboard_id})") + return self._ws.lakeview.update( + dashboard_id=dashboard_id, + dashboard=Dashboard( + display_name=metadata.display_name, + serialized_dashboard=json.dumps(dashboard_def), + ), + ) + + def _publish_dashboard(self, dashboard: Dashboard, display_name: str) -> None: + """ + Publish a dashboard and handle publication failures gracefully. - # Revert back SQL queries to placeholder format regardless of success - self._resolve_table_name_in_queries(src_tbl_name=dq_table, replaced_tbl_name=src_table_name, folder=folder) + Args: + dashboard: Dashboard object to publish + display_name: Display name for logging purposes + """ + if dashboard.dashboard_id: + try: + self._ws.lakeview.publish(dashboard.dashboard_id) + logger.info(f"Published dashboard: {display_name}") + except Exception as e: + logger.warning(f"Failed to publish dashboard: {e}") diff --git a/src/databricks/labs/dqx/installer/logs.py b/src/databricks/labs/dqx/installer/logs.py index a77a73831..2822d5027 100644 --- a/src/databricks/labs/dqx/installer/logs.py +++ b/src/databricks/labs/dqx/installer/logs.py @@ -178,7 +178,7 @@ def _exclusive_open(cls, filename: str, **kwargs): lockfile = cls._create_lock(lockfile_name) try: - with open(filename, encoding="utf8", **kwargs) as f: + with open(filename, encoding="utf-8", **kwargs) as f: yield f finally: try: diff --git a/src/databricks/labs/dqx/installer/workflow_task.py b/src/databricks/labs/dqx/installer/workflow_task.py index a76d15119..3515b7b68 100644 --- a/src/databricks/labs/dqx/installer/workflow_task.py +++ b/src/databricks/labs/dqx/installer/workflow_task.py @@ -3,7 +3,6 @@ from dataclasses import dataclass from databricks.labs.blueprint.installation import Installation -from databricks.labs.lsql.backends import SqlBackend from databricks.sdk import WorkspaceClient from databricks.labs.dqx.config import WorkspaceConfig @@ -18,7 +17,7 @@ class Task: workflow: str name: str doc: str - fn: Callable[[WorkspaceConfig, WorkspaceClient, SqlBackend, Installation], None] + fn: Callable[[WorkspaceConfig, WorkspaceClient, Installation], None] depends_on: list[str] | None = None job_cluster: str | None = None # cluster key for job clusters; if None, uses serverless environment override_clusters: dict[str, str] | None = None diff --git a/src/databricks/labs/dqx/queries/quality/dashboard/00_1_dq_summary.sql b/src/databricks/labs/dqx/queries/quality/dashboard/00_1_dq_summary.sql deleted file mode 100644 index ec1906a7c..000000000 --- a/src/databricks/labs/dqx/queries/quality/dashboard/00_1_dq_summary.sql +++ /dev/null @@ -1,9 +0,0 @@ -/* --title 'Data Quality Summary' */ -SELECT - CASE - WHEN _errors IS NOT NULL THEN 'Errors' - WHEN _warnings IS NOT NULL THEN 'Warnings' - END AS category, - ROUND((COUNT(*) * 100.0 / (SELECT COUNT(*) FROM $catalog.schema.table)), 2) AS percentage -FROM $catalog.schema.table -GROUP BY category \ No newline at end of file diff --git a/src/databricks/labs/dqx/queries/quality/dashboard/00_2_dq_error_types.sql b/src/databricks/labs/dqx/queries/quality/dashboard/00_2_dq_error_types.sql deleted file mode 100644 index a7acf4eed..000000000 --- a/src/databricks/labs/dqx/queries/quality/dashboard/00_2_dq_error_types.sql +++ /dev/null @@ -1,36 +0,0 @@ -/* --title 'Error and Warning Types Breakdown' */ -WITH error_types AS ( - SELECT - 'Error' AS category, - error_struct.name AS type, - COUNT(*) AS count - FROM $catalog.schema.table - LATERAL VIEW EXPLODE(_errors) exploded_errors AS error_struct - WHERE _errors IS NOT NULL - GROUP BY error_struct.name -), -warning_types AS ( - SELECT - 'Warning' AS category, - warning_struct.name AS type, - COUNT(*) AS count - FROM $catalog.schema.table - LATERAL VIEW EXPLODE(_warnings) exploded_warnings AS warning_struct - WHERE _warnings IS NOT NULL - GROUP BY warning_struct.name -), -combined AS ( - SELECT * FROM error_types - UNION ALL - SELECT * FROM warning_types -), -total AS ( - SELECT SUM(count) AS total_count FROM combined -) -SELECT - category, - type, - count, - ROUND((count * 100.0) / total.total_count, 2) AS percentage -FROM combined, total -ORDER BY category, count DESC; \ No newline at end of file diff --git a/src/databricks/labs/dqx/queries/quality/dashboard/dashboard.yml b/src/databricks/labs/dqx/queries/quality/dashboard/dashboard.yml deleted file mode 100644 index 0c14771e2..000000000 --- a/src/databricks/labs/dqx/queries/quality/dashboard/dashboard.yml +++ /dev/null @@ -1,73 +0,0 @@ -display_name: Data Quality Summary -tiles: - 00_1_dq_summary: - title: DQ Summary - overrides: - queries: - - name: main_query - query: - datasetName: 00_1_dq_summary - disaggregated: false - fields: - - expression: SUM(`percentage`) - name: sum(percentage) - - expression: '`category`' - name: category - spec: - encodings: - angle: - displayName: Total Row Counts - fieldName: sum(percentage) - scale: - type: quantitative - color: - displayName: category - fieldName: category - scale: - mappings: - - color: '#00A972' - value: Good - - color: '#FFAB00' - value: Warnings - - color: '#FF3621' - value: Errors - type: categorical - label: - show: true - version: 3 - widgetType: pie - 00_2_dq_error_types: - title: DQ Error and Warning Types Breakdown - hidden: false - overrides: - queries: - - name: main_query - query: - datasetName: 00_2_dq_error_types - disaggregated: true - fields: - - expression: '`category`' - name: Category - - expression: '`type`' - name: Type - - expression: '`count`' - name: Count - - expression: '`percentage`' - name: Percentage - spec: - encodings: - columns: - - fieldName: Category - displayName: Category - sortable: true - - fieldName: Type - displayName: Type - sortable: true - - fieldName: Count - displayName: Count - sortable: true - - fieldName: Percentage - displayName: Percentage (%) - sortable: true - version: 2 - widgetType: table diff --git a/src/databricks/labs/dqx/telemetry.py b/src/databricks/labs/dqx/telemetry.py index 43c50e393..4ebe5796f 100644 --- a/src/databricks/labs/dqx/telemetry.py +++ b/src/databricks/labs/dqx/telemetry.py @@ -2,6 +2,7 @@ import sys import functools import logging +import hashlib from io import StringIO from collections.abc import Callable from pyspark.sql import DataFrame, SparkSession @@ -72,10 +73,14 @@ def wrapper(self, *args, **kwargs): def log_dataframe_telemetry(ws: WorkspaceClient, spark: SparkSession, df: DataFrame): """ Log telemetry information about a Spark DataFrame to the Databricks workspace including: - - Number of input tables and non-table inputs + - List of tables used as inputs (hashed) + - List of file paths used as inputs (hashed, excluding paths from tables) - Whether the DataFrame is streaming - Whether running in a Delta Live Tables (DLT) pipeline + This function is designed to never throw exceptions - it will log errors but continue execution + to ensure telemetry failures don't break the main application flow. + Args: ws: WorkspaceClient spark: SparkSession @@ -84,38 +89,114 @@ def log_dataframe_telemetry(ws: WorkspaceClient, spark: SparkSession, df: DataFr Returns: None """ - input_table_count = count_tables_in_spark_plan(df) - log_telemetry(ws, "table_input_count", str(input_table_count)) - # assume 1 input if no tables - log_telemetry(ws, "non_table_input_count", str(0 if input_table_count > 0 else 1)) log_telemetry(ws, "streaming", str(df.isStreaming).lower()) log_telemetry(ws, "dlt", str(is_dlt_pipeline(spark)).lower()) + plan_str = get_spark_plan_as_string(df) + if plan_str: + input_tables = get_tables_from_spark_plan(plan_str) + for table in input_tables: + log_telemetry(ws, "input_table", hashlib.sha256(("id:" + table).encode("utf-8")).hexdigest()) + + input_paths = get_paths_from_spark_plan(plan_str, input_tables) + for path in input_paths: + log_telemetry(ws, "input_path", hashlib.sha256(("id:" + path).encode("utf-8")).hexdigest()) + + +def get_tables_from_spark_plan(plan_str: str) -> set[str]: + """ + Extract table names from the Analyzed Logical Plan section of a Spark execution plan. + + This function parses the Analyzed Logical Plan section and identifies table references + by finding SubqueryAlias nodes, which Spark uses to represent table references in the + logical plan. File-based sources (e.g., Delta files from volumes) and in-memory DataFrames + do not create SubqueryAlias nodes and therefore won't be counted as tables. + + Args: + plan_str: The complete Spark execution plan string (from df.explain(True)) + + Returns: + A set of distinct table names found in the plan. Returns empty set if no + Analyzed Logical Plan section is found or no tables are referenced. + """ + try: + return _extract_tables_from_analyzed_plan(plan_str) + except Exception as e: + logger.debug(f"Failed to extract tables from Spark plan: {e}") + return set() + -def count_tables_in_spark_plan(df: DataFrame) -> int: +def get_paths_from_spark_plan(plan_str: str, table_names: set[str] | None = None) -> set[str]: """ - Count the number of tables referenced in a DataFrame's Spark execution plan. + Extract file paths from the Physical Plan section of a Spark execution plan. - This function analyzes the Analyzed Logical Plan section of the Spark execution plan - to identify table references (via SubqueryAlias nodes). File-based DataFrames and - in-memory DataFrames will return 0. + This function parses the Physical Plan section and identifies file path references + by finding any *FileIndex patterns in the Location field (e.g., PreparedDeltaFileIndex, + ParquetFileIndex, etc.). These paths represent direct file-based data sources + (e.g., files from volumes, DBFS, S3, etc.) that are not registered as tables. Args: - df: The Spark DataFrame to analyze + plan_str: The complete Spark execution plan string (from df.explain(True)) + table_names: Optional set of table names to exclude (paths associated with tables are skipped) Returns: - The number of distinct tables found in the execution plan. Returns 0 if the plan - cannot be retrieved or contains no table references. + A set of distinct file paths found in the plan. Returns empty set if no + Physical Plan section is found or no paths are referenced. """ try: - plan_str = _get_spark_plan_as_string(df) - if not plan_str: - return 0 - tables = _extract_tables_from_spark_plan(plan_str) - return len(tables) + return _extract_paths_from_physical_plan(plan_str, table_names) except Exception as e: - logger.debug(f"Failed to count tables in Spark plan: {e}") - return 0 + logger.debug(f"Failed to extract paths from Spark plan: {e}") + return set() + + +def _extract_tables_from_analyzed_plan(plan_str: str) -> set[str]: + """Helper function to extract tables from the Analyzed Logical Plan section.""" + tables: set[str] = set() + + # Extract Analyzed Logical Plan section (stop at next "==") + match = re.search(r"== Analyzed Logical Plan ==\s*(.*?)\n==", plan_str, re.DOTALL) + if not match: + return tables + + analyzed_text = match.group(1) + + # Extract SubqueryAlias names (only present if table is used) + subquery_aliases = re.findall(r"SubqueryAlias\s+([^\s]+)", analyzed_text) + tables.update(alias.replace("`", "") for alias in subquery_aliases) + + return tables + + +def _extract_paths_from_physical_plan(plan_str: str, table_names: set[str] | None = None) -> set[str]: + """Helper function to extract paths from the Physical Plan section.""" + paths: set[str] = set() + + # Extract Physical Plan section (stop at next "==") + match = re.search(r"== Physical Plan ==\s*(.*?)(?:\n==|$)", plan_str, re.DOTALL) + if not match: + return paths + + physical_text = match.group(1) + if table_names is None: + table_names = set() + + # Process line by line to check for table associations + for line in physical_text.split('\n'): + location_match = re.search(r"Location:\s+\w+FileIndex\([^)]+\)\[([^\]]+)\]", line) + if not location_match: + continue + + # Skip if this line contains any table name (indicating it's a table-based path) + if any(table_name in line for table_name in table_names): + continue + + # Extract and add non-empty paths + paths_str = location_match.group(1) + path_list = [p.strip() for p in paths_str.split(',') if p.strip()] + paths.update(path_list) + + return paths def is_dlt_pipeline(spark: SparkSession) -> bool: @@ -137,7 +218,7 @@ def is_dlt_pipeline(spark: SparkSession) -> bool: return False -def _get_spark_plan_as_string(df: DataFrame) -> str: +def get_spark_plan_as_string(df: DataFrame) -> str: """ Retrieve the Spark execution plan as a string by capturing df.explain() output. @@ -161,35 +242,3 @@ def _get_spark_plan_as_string(df: DataFrame) -> str: finally: sys.stdout = old_stdout return buf.getvalue() - - -def _extract_tables_from_spark_plan(plan_str: str) -> set[str]: - """ - Extract table names from the Analyzed Logical Plan section of a Spark execution plan. - - This function parses the Analyzed Logical Plan section and identifies table references - by finding SubqueryAlias nodes, which Spark uses to represent table references in the - logical plan. File-based sources (e.g., Delta files from volumes) and in-memory DataFrames - do not create SubqueryAlias nodes and therefore won't be counted as tables. - - Args: - plan_str: The complete Spark execution plan string (from df.explain(True)) - - Returns: - A set of distinct table names found in the plan. Returns empty set if no - Analyzed Logical Plan section is found or no tables are referenced. - """ - tables: set[str] = set() - - # Extract Analyzed Logical Plan section (stop at next "==") - match = re.search(r"== Analyzed Logical Plan ==\s*(.*?)\n==", plan_str, re.DOTALL) # non-greedy until next section - if not match: - return tables - - analyzed_text = match.group(1) - - # Extract SubqueryAlias names (only present if table is used) - subquery_aliases = re.findall(r"SubqueryAlias\s+([^\s]+)", analyzed_text) - tables.update(subquery_aliases) - - return tables diff --git a/tests/integration/test_ai_rules_generator.py b/tests/integration/test_ai_rules_generator.py index e5c089b51..6558174c3 100644 --- a/tests/integration/test_ai_rules_generator.py +++ b/tests/integration/test_ai_rules_generator.py @@ -131,7 +131,7 @@ def test_generate_dq_rules_ai_assisted_with_is_not_equal_to_str(ws, spark): def test_generate_dq_rules_ai_assisted_with_sql_expression(ws, spark): - user_input = "Users email must not end with @gmail.com checked using sql expression." + user_input = "Users email must not end with @gmail.com checked using sql expression, skip msg." generator = DQGenerator(ws, spark) actual_checks = generator.generate_dq_rules_ai_assisted(user_input=user_input) diff --git a/tests/integration/test_telemetry.py b/tests/integration/test_telemetry.py index 39689e554..42508790d 100644 --- a/tests/integration/test_telemetry.py +++ b/tests/integration/test_telemetry.py @@ -1,10 +1,14 @@ import pyspark.sql.types as T -from databricks.labs.dqx.telemetry import count_tables_in_spark_plan +from databricks.labs.dqx.telemetry import ( + get_tables_from_spark_plan, + get_paths_from_spark_plan, + get_spark_plan_as_string, +) from tests.conftest import TEST_CATALOG -def test_count_tables_from_file_based_dataframe(ws, spark, make_schema, make_random, make_volume): +def test_get_tables_and_paths_from_file_delta_dataframe(ws, spark, make_schema, make_random, make_volume): catalog_name = TEST_CATALOG schema_name = make_schema(catalog_name=catalog_name).name volume_name = make_volume(catalog_name=catalog_name, schema_name=schema_name).name @@ -15,21 +19,46 @@ def test_count_tables_from_file_based_dataframe(ws, spark, make_schema, make_ran df.write.format("delta").save(volume_path) test_df = spark.read.format("delta").load(volume_path) - count = count_tables_in_spark_plan(test_df) + plan = get_spark_plan_as_string(test_df) + tables = get_tables_from_spark_plan(plan) + paths = get_paths_from_spark_plan(plan, tables) - assert count == 0, f"Expected 0 tables, but found {count}" + assert len(tables) == 0, f"Expected 0 tables, but found {len(tables)}" + assert len(paths) == 1, f"Expected 1 path, but found {len(paths)}" -def test_count_tables_from_code_base_dataframe(ws, spark, make_schema, make_random, make_volume): +def test_get_tables_and_paths_from_file_parquet_dataframe(ws, spark, make_schema, make_random, make_volume): + catalog_name = TEST_CATALOG + schema_name = make_schema(catalog_name=catalog_name).name + volume_name = make_volume(catalog_name=catalog_name, schema_name=schema_name).name + volume_path = f"/Volumes/{catalog_name}/{schema_name}/{volume_name}/" + input_schema = T.StructType([T.StructField("id", T.IntegerType())]) df = spark.createDataFrame([[1]], schema=input_schema) + df.write.mode("overwrite").format("parquet").save(volume_path) + test_df = spark.read.format("parquet").load(volume_path) - count = count_tables_in_spark_plan(df) + plan = get_spark_plan_as_string(test_df) + tables = get_tables_from_spark_plan(plan) + paths = get_paths_from_spark_plan(plan, tables) - assert count == 0, f"Expected 0 tables, but found {count}" + assert len(tables) == 0, f"Expected 0 tables, but found {len(tables)}" + assert len(paths) == 1, f"Expected 1 path, but found {len(paths)}" -def test_count_tables_from_table_based_dataframe(ws, spark, make_schema, make_random): +def test_get_tables_and_paths_from_code_dataframe(ws, spark, make_schema, make_random, make_volume): + input_schema = T.StructType([T.StructField("id", T.IntegerType())]) + test_df = spark.createDataFrame([[1]], schema=input_schema) + + plan = get_spark_plan_as_string(test_df) + tables = get_tables_from_spark_plan(plan) + paths = get_paths_from_spark_plan(plan, tables) + + assert len(tables) == 0, f"Expected 0 tables, but found {len(tables)}" + assert len(paths) == 0, f"Expected 0 paths, but found {len(paths)}" + + +def test_get_tables_and_paths_from_table_dataframe(ws, spark, make_schema, make_random): catalog_name = TEST_CATALOG schema_name = make_schema(catalog_name=catalog_name).name table_name = f"{catalog_name}.{schema_name}.{make_random(10).lower()}" @@ -39,12 +68,16 @@ def test_count_tables_from_table_based_dataframe(ws, spark, make_schema, make_ra df.write.format("delta").saveAsTable(table_name) test_df = spark.table(table_name) - count = count_tables_in_spark_plan(test_df) + plan = get_spark_plan_as_string(test_df) + tables = get_tables_from_spark_plan(plan) + paths = get_paths_from_spark_plan(plan, tables) - assert count == 1, f"Expected 1 table, but found {count}" + assert len(tables) == 1, f"Expected 1 table, but found {len(tables)}" + assert table_name in tables, f"Expected table with name {table_name}" + assert len(paths) == 0, f"Expected 0 paths, but found {len(paths)}" -def test_count_tables_from_aggregated_table_based_dataframe(ws, spark, make_schema, make_random): +def test_get_tables_and_paths_from_aggregated_table_dataframe(ws, spark, make_schema, make_random): catalog_name = TEST_CATALOG schema_name = make_schema(catalog_name=catalog_name).name table_name = f"{catalog_name}.{schema_name}.{make_random(10).lower()}" @@ -54,12 +87,16 @@ def test_count_tables_from_aggregated_table_based_dataframe(ws, spark, make_sche df.write.format("delta").saveAsTable(table_name) test_df = spark.table(table_name) - count = count_tables_in_spark_plan(test_df) + plan = get_spark_plan_as_string(test_df) + tables = get_tables_from_spark_plan(plan) + paths = get_paths_from_spark_plan(plan, tables) - assert count == 1, f"Expected 1 table, but found {count}" + assert len(tables) == 1, f"Expected 1 table, but found {len(tables)}" + assert table_name in tables, f"Expected table with name {table_name}" + assert len(paths) == 0, f"Expected 0 paths, but found {len(paths)}" -def test_count_multiple_joined_tables_from_table_based_dataframe(ws, spark, make_schema, make_random): +def test_get_tables_and_paths_from_joined_tables_dataframe(ws, spark, make_schema, make_random): catalog_name = TEST_CATALOG schema_name = make_schema(catalog_name=catalog_name).name table_name1 = f"{catalog_name}.{schema_name}.{make_random(10).lower()}" @@ -73,12 +110,17 @@ def test_count_multiple_joined_tables_from_table_based_dataframe(ws, spark, make test_df2 = spark.table(table_name2) test_df = test_df1.join(test_df2, on="id", how="inner") - count = count_tables_in_spark_plan(test_df) + plan = get_spark_plan_as_string(test_df) + tables = get_tables_from_spark_plan(plan) + paths = get_paths_from_spark_plan(plan, tables) - assert count == 2, f"Expected 2 tables, but found {count}" + assert len(tables) == 2, f"Expected 2 tables, but found {len(tables)}" + assert table_name1 in tables, f"Expected table with name {table_name1}" + assert table_name2 in tables, f"Expected table with name {table_name2}" + assert len(paths) == 0, f"Expected 0 paths, but found {len(paths)}" -def test_count_multiple_unioned_tables_from_table_based_dataframe(ws, spark, make_schema, make_random): +def test_get_tables_and_paths_from_unioned_tables_dataframe(ws, spark, make_schema, make_random): catalog_name = TEST_CATALOG schema_name = make_schema(catalog_name=catalog_name).name table_name1 = f"{catalog_name}.{schema_name}.{make_random(10).lower()}" @@ -92,17 +134,22 @@ def test_count_multiple_unioned_tables_from_table_based_dataframe(ws, spark, mak test_df2 = spark.table(table_name2) test_df = test_df1.union(test_df2) - count = count_tables_in_spark_plan(test_df) + plan = get_spark_plan_as_string(test_df) + tables = get_tables_from_spark_plan(plan) + paths = get_paths_from_spark_plan(plan, tables) - assert count == 2, f"Expected 2 tables, but found {count}" + assert len(tables) == 2, f"Expected 2 tables, but found {len(tables)}" + assert table_name1 in tables, f"Expected table with name {table_name1}" + assert table_name2 in tables, f"Expected table with name {table_name2}" + assert len(paths) == 0, f"Expected 0 paths, but found {len(paths)}" -def test_count_path_and_table_based_dataframe(ws, spark, make_schema, make_random, make_volume): +def test_get_tables_and_paths_from_mixed_dataframe(ws, spark, make_schema, make_random, make_volume): catalog_name = TEST_CATALOG schema_name = make_schema(catalog_name=catalog_name).name table_name = f"{catalog_name}.{schema_name}.{make_random(10).lower()}" volume_name = make_volume(catalog_name=catalog_name, schema_name=schema_name).name - volume_path = f"/Volumes/{catalog_name}/{schema_name}/{volume_name}/" + volume_path = f"/Volumes/{catalog_name}/{schema_name}/{volume_name}" input_schema = T.StructType([T.StructField("id", T.IntegerType())]) df = spark.createDataFrame([[1]], schema=input_schema) @@ -113,12 +160,18 @@ def test_count_path_and_table_based_dataframe(ws, spark, make_schema, make_rando test_df2 = spark.read.format("delta").load(volume_path) test_df = test_df1.join(test_df2, on="id", how="inner") - count = count_tables_in_spark_plan(test_df) + plan = get_spark_plan_as_string(test_df) + tables = get_tables_from_spark_plan(plan) + paths = get_paths_from_spark_plan(plan, tables) - assert count == 1, f"Expected 1 table, but found {count}" + assert len(tables) == 1, f"Expected 1 table, but found {len(tables)}" + assert table_name in tables, f"Expected table with name {table_name}" + assert len(paths) == 1, f"Expected 1 path, but found {len(paths)}" + # Volume paths in Spark plans are shown with dbfs: prefix + assert any(volume_path in path for path in paths), f"Expected path containing {volume_path}, got {paths}" -def test_count_tables_from_streaming_table_based_dataframe(ws, spark, make_schema, make_random): +def test_get_tables_and_paths_from_streaming_table_based_dataframe(ws, spark, make_schema, make_random): catalog_name = TEST_CATALOG schema_name = make_schema(catalog_name=catalog_name).name table_name = f"{catalog_name}.{schema_name}.{make_random(10).lower()}" @@ -128,6 +181,10 @@ def test_count_tables_from_streaming_table_based_dataframe(ws, spark, make_schem df.write.format("delta").saveAsTable(table_name) test_df = spark.readStream.table(table_name) - count = count_tables_in_spark_plan(test_df) + plan = get_spark_plan_as_string(test_df) + tables = get_tables_from_spark_plan(plan) + paths = get_paths_from_spark_plan(plan, tables) - assert count == 1, f"Expected 1 table, but found {count}" + assert len(tables) == 1, f"Expected 1 table, but found {len(tables)}" + assert table_name in tables, f"Expected table with name {table_name}" + assert len(paths) == 0, f"Expected 0 paths, but found {len(paths)}" diff --git a/tests/integration/test_upload_dependencies.py b/tests/integration/test_upload_dependencies.py index fe5999708..83c842a8a 100644 --- a/tests/integration/test_upload_dependencies.py +++ b/tests/integration/test_upload_dependencies.py @@ -134,7 +134,6 @@ def test_workflow_deployment_uploads_dependencies(ws, installation_with_upload_d # Verify core dependencies are included assert any("databricks_sdk" in dep for dep in call_args) - assert any("databricks_labs_lsql" in dep for dep in call_args) # Verify job was created with the uploaded wheels mock_create_job.assert_called_once() diff --git a/tests/unit/test_telemetry.py b/tests/unit/test_telemetry.py index 07a539170..a554edfbd 100644 --- a/tests/unit/test_telemetry.py +++ b/tests/unit/test_telemetry.py @@ -1,4 +1,9 @@ -from databricks.labs.dqx.telemetry import is_dlt_pipeline, count_tables_in_spark_plan +from databricks.labs.dqx.telemetry import ( + is_dlt_pipeline, + get_tables_from_spark_plan, + get_paths_from_spark_plan, + get_spark_plan_as_string, +) class DummySparkConf: @@ -54,34 +59,229 @@ def test_is_dlt_pipeline_false_exception(): assert is_dlt_pipeline(dummy_spark) is False -def test_count_tables_with_missing_analyzed_plan_section(): - """Test that count_tables_in_spark_plan returns 0 when plan has no Analyzed Logical Plan section.""" +def test_get_tables_with_missing_analyzed_plan_section(): + """Test that get_tables_from_spark_plan returns 0 when plan has no Analyzed Logical Plan section.""" # Create a plan string without the "== Analyzed Logical Plan ==" section plan_without_analyzed_section = """== Physical Plan == Some physical plan details here == Optimized Logical Plan == Some optimized plan details here """ + tables = get_tables_from_spark_plan(plan_without_analyzed_section) - mock_df = MockDataFrameWithPlan(plan_without_analyzed_section) - count = count_tables_in_spark_plan(mock_df) + assert len(tables) == 0, f"Expected 0 tables when Analyzed Logical Plan section is missing, but found {len(tables)}" - assert count == 0, f"Expected 0 tables when Analyzed Logical Plan section is missing, but found {count}" - -def test_count_tables_with_explain_exception(): - """Test that count_tables_in_spark_plan returns 0 when df.explain() raises an exception.""" +def test_get_spark_plan_exception(): + """Test that get_tables_from_spark_plan returns 0 when df.explain() raises an exception.""" mock_df = MockDataFrameWithException(RuntimeError("Failed to explain plan")) - count = count_tables_in_spark_plan(mock_df) + tables = get_spark_plan_as_string(mock_df) + + assert len(tables) == 0, f"Expected 0 tables when explain() raises exception, but found {len(tables)}" + + +def test_get_paths_with_single_path(): + """Test that get_paths_from_spark_plan extracts a single path from Physical Plan.""" + plan_with_single_path = """== Physical Plan == +*(1) ColumnarToRow ++- PhotonResultStage + +- PhotonScan parquet [vendor_id#1080600] DataFilters: [], DictionaryFilters: [], Format: parquet, Location: PreparedDeltaFileIndex(1 paths)[/Volumes/catalog/schema/volume/nyctaxi_2019], OptionalDataFilters: [], PartitionFilters: [], ReadSchema: struct + +== Photon Explanation == +""" + + paths = get_paths_from_spark_plan(plan_with_single_path) + + assert len(paths) == 1, f"Expected 1 path, but found {len(paths)}" + assert "/Volumes/catalog/schema/volume/nyctaxi_2019" in paths + + +def test_get_paths_with_multiple_paths(): + """Test that get_paths_from_spark_plan extracts multiple paths from Physical Plan.""" + plan_with_multiple_paths = """== Physical Plan == +AdaptiveSparkPlan isFinalPlan=false ++- == Initial Plan == + ColumnarToRow + +- PhotonResultStage + +- PhotonProject [vendor_id#1081424] + +- PhotonShuffledHashJoin [vendor_id#1081424], [vendor_id#1081478], Inner, BuildLeft + :- PhotonShuffleExchangeSource + : +- PhotonShuffleMapStage + : +- PhotonShuffleExchangeSink hashpartitioning(vendor_id#1081424, 200) + : +- PhotonScan parquet [vendor_id#1081424] DataFilters: [isnotnull(vendor_id#1081424)], DictionaryFilters: [], Format: parquet, Location: PreparedDeltaFileIndex(1 paths)[/Volumes/catalog/schema/volume1/nyctaxi_2019], OptionalDataFilters: [], PartitionFilters: [], ReadSchema: struct + +- PhotonShuffleExchangeSource + +- PhotonShuffleMapStage + +- PhotonShuffleExchangeSink hashpartitioning(vendor_id#1081478, 200) + +- PhotonScan parquet users.marcin_wojtyczka.nyctaxi[vendor_id#1081478] DataFilters: [isnotnull(vendor_id#1081478)], DictionaryFilters: [], Format: parquet, Location: PreparedDeltaFileIndex(1 paths)[s3://databricks-e2demofieldengwest/b169b504-4c54-49f2-bc3a-adf4b1], OptionalDataFilters: [], PartitionFilters: [], ReadSchema: struct + +== Photon Explanation == +""" + + paths = get_paths_from_spark_plan(plan_with_multiple_paths) + + assert len(paths) == 2, f"Expected 2 paths, but found {len(paths)}" + assert "/Volumes/catalog/schema/volume1/nyctaxi_2019" in paths + assert "s3://databricks-e2demofieldengwest/b169b504-4c54-49f2-bc3a-adf4b1" in paths + + +def test_get_paths_with_missing_physical_plan_section(): + """Test that get_paths_from_spark_plan returns 0 when plan has no Physical Plan section.""" + plan_without_physical_section = """== Analyzed Logical Plan == +Some analyzed plan details here +== Optimized Logical Plan == +Some optimized plan details here +""" + + paths = get_paths_from_spark_plan(plan_without_physical_section) + + assert len(paths) == 0, f"Expected 0 paths when Physical Plan section is missing, but found {len(paths)}" + + +def test_get_paths_with_no_location_patterns(): + """Test that get_paths_from_spark_plan returns 0 when Physical Plan has no Location patterns.""" + plan_without_locations = """== Physical Plan == +*(1) ColumnarToRow ++- Some plan without Location patterns +== Photon Explanation == +""" + + paths = get_paths_from_spark_plan(plan_without_locations) + + assert len(paths) == 0, f"Expected 0 paths when no Location patterns found, but found {len(paths)}" + + +def test_get_paths_with_different_fileindex_types(): + """Test that get_paths_from_spark_plan extracts paths from different FileIndex types.""" + plan_with_various_fileindex = """== Physical Plan == +*(1) ColumnarToRow ++- PhotonResultStage + +- PhotonScan parquet [id#100] DataFilters: [], Format: parquet, Location: InMemoryFileIndex(1 paths)[file:/tmp/data/dataset1], ReadSchema: struct + +- PhotonScan parquet [id#200] DataFilters: [], Format: parquet, Location: ParquetFileIndex(1 paths)[s3://bucket/path/dataset2], ReadSchema: struct + +- PhotonScan parquet [id#300] DataFilters: [], Format: parquet, Location: PreparedDeltaFileIndex(1 paths)[/Volumes/catalog/schema/volume/dataset3], ReadSchema: struct + +== Photon Explanation == +""" + + paths = get_paths_from_spark_plan(plan_with_various_fileindex) + + assert len(paths) == 3, f"Expected 3 paths from different FileIndex types, but found {len(paths)}" + assert "file:/tmp/data/dataset1" in paths, "Expected to find InMemoryFileIndex path" + assert "s3://bucket/path/dataset2" in paths, "Expected to find ParquetFileIndex path" + assert "/Volumes/catalog/schema/volume/dataset3" in paths, "Expected to find PreparedDeltaFileIndex path" + + +def test_get_paths_excludes_table_paths(): + """Test that get_paths_from_spark_plan excludes paths associated with tables when requested.""" + plan_with_table_and_file = """== Analyzed Logical Plan == +vendor_id: string, pickup_datetime: timestamp +Project [vendor_id#1081424, pickup_datetime#1081425] ++- Join Inner, (vendor_id#1081424 = vendor_id#1081478) + :- Relation [vendor_id#1081424,pickup_datetime#1081425] parquet + +- SubqueryAlias users.marcin_wojtyczka.nyctaxi + +- Relation users.marcin_wojtyczka.nyctaxi[vendor_id#1081478,pickup_datetime#1081479] parquet + +== Physical Plan == +AdaptiveSparkPlan isFinalPlan=false ++- PhotonProject [vendor_id#1081424, pickup_datetime#1081425] + +- PhotonShuffledHashJoin [vendor_id#1081424], [vendor_id#1081478], Inner, BuildLeft + :- PhotonScan parquet [vendor_id#1081424] DataFilters: [], Format: parquet, Location: PreparedDeltaFileIndex(1 paths)[/Volumes/catalog/schema/volume/nyctaxi_2019], ReadSchema: struct + +- PhotonScan parquet users.marcin_wojtyczka.nyctaxi[vendor_id#1081478] DataFilters: [], Format: parquet, Location: PreparedDeltaFileIndex(1 paths)[s3://databricks-e2demofieldengwest/b169b504-4c54-49f2-bc3a-adf4b1], ReadSchema: struct + +== Photon Explanation == +""" + + table_names = {"users.marcin_wojtyczka.nyctaxi"} + + # Without excluding table paths + paths_all = get_paths_from_spark_plan(plan_with_table_and_file) + assert len(paths_all) == 2, f"Expected 2 paths without exclusion, but found {len(paths_all)}" + + # With excluding table paths + paths_filtered = get_paths_from_spark_plan(plan_with_table_and_file, table_names) + assert len(paths_filtered) == 1, f"Expected 1 path with table exclusion, but found {len(paths_filtered)}" + assert "/Volumes/catalog/schema/volume/nyctaxi_2019" in paths_filtered, "Expected to find the file-based path" + assert ( + "s3://databricks-e2demofieldengwest/b169b504-4c54-49f2-bc3a-adf4b1" not in paths_filtered + ), "Table path should be excluded" + + +def test_get_tables_handles_none_input(): + """Test that get_tables_from_spark_plan handles None input gracefully.""" + tables = get_tables_from_spark_plan(None) + assert len(tables) == 0, "Expected empty set for None input" + + +def test_get_tables_handles_empty_string(): + """Test that get_tables_from_spark_plan handles empty string gracefully.""" + tables = get_tables_from_spark_plan("") + assert len(tables) == 0, "Expected empty set for empty string" - assert count == 0, f"Expected 0 tables when explain() raises exception, but found {count}" +def test_get_paths_handles_none_input(): + """Test that get_paths_from_spark_plan handles None input gracefully.""" + paths = get_paths_from_spark_plan(None) + assert len(paths) == 0, "Expected empty set for None input" -def test_count_tables_with_general_exception(): - """Test that count_tables_in_spark_plan returns 0 when an unexpected exception occurs.""" - mock_df = MockDataFrameWithException(ValueError("Unexpected error")) - count = count_tables_in_spark_plan(mock_df) +def test_get_paths_handles_empty_string(): + """Test that get_paths_from_spark_plan handles empty string gracefully.""" + paths = get_paths_from_spark_plan("") + assert len(paths) == 0, "Expected empty set for empty string" + + +def test_get_tables_handles_regex_exception(): + """Test that get_tables_from_spark_plan handles exceptions during regex processing.""" + # Create a pathological string that could cause regex issues + # Use a very large string to potentially trigger catastrophic backtracking + large_plan = "== Analyzed Logical Plan ==\n" + "SubqueryAlias " + "a" * 100000 + "\n==" + + # Should not raise exception, just return empty or partial results + tables = get_tables_from_spark_plan(large_plan) + assert isinstance(tables, set), "Expected set return type even with problematic input" + + +def test_get_paths_handles_malformed_location_pattern(): + """Test that get_paths_from_spark_plan handles malformed Location patterns gracefully.""" + # Malformed pattern with unbalanced brackets and special characters + malformed_plan = """== Physical Plan == ++- PhotonScan Location: PreparedDeltaFileIndex(1 paths)[[[unclosed ++- PhotonScan Location: FileIndex(broken)[path1, path2 ++- Normal line without location +""" + + # Should not raise exception + paths = get_paths_from_spark_plan(malformed_plan) + assert isinstance(paths, set), "Expected set return type even with malformed input" + + +def test_get_tables_with_unicode_and_special_chars(): + """Test that get_tables_from_spark_plan handles unicode and special characters gracefully.""" + plan_with_unicode = """== Analyzed Logical Plan == +SubqueryAlias 表名_with_中文 +SubqueryAlias table.with.dots +SubqueryAlias table-with-dashes +SubqueryAlias table$with$dollars +SubqueryAlias table@with@symbols ++- Other content 🚀 with emojis +== Next Section == +""" + + # Should handle gracefully and extract what it can + tables = get_tables_from_spark_plan(plan_with_unicode) + assert isinstance(tables, set), "Expected set return type with unicode input" + # At least some tables should be extracted + assert len(tables) >= 1, "Should extract at least some table names" + + +def test_get_paths_with_unicode_paths(): + """Test that get_paths_from_spark_plan handles unicode in paths gracefully.""" + plan_with_unicode_paths = """== Physical Plan == ++- PhotonScan Location: PreparedDeltaFileIndex(1 paths)[/Volumes/目录/schema/卷/数据] ++- PhotonScan Location: ParquetFileIndex(1 paths)[s3://bucket/path/🚀data] +== Photon Explanation == +""" - assert count == 0, f"Expected 0 tables when exception occurs, but found {count}" + # Should handle gracefully + paths = get_paths_from_spark_plan(plan_with_unicode_paths) + assert isinstance(paths, set), "Expected set return type with unicode paths"