diff --git a/src/databricks/labs/dqx/checks_resolver.py b/src/databricks/labs/dqx/checks_resolver.py index 0c410776b..02997fdbc 100644 --- a/src/databricks/labs/dqx/checks_resolver.py +++ b/src/databricks/labs/dqx/checks_resolver.py @@ -7,6 +7,14 @@ from databricks.labs.dqx import check_funcs from databricks.labs.dqx.geo import check_funcs as geo_check_funcs + +try: + from databricks.labs.dqx.pii import pii_detection_funcs as pii_check_funcs + + PII_ENABLED = True +except ImportError: + PII_ENABLED = False + from databricks.labs.dqx.errors import InvalidCheckError logger = logging.getLogger(__name__) @@ -32,8 +40,11 @@ def resolve_check_function( logger.debug(f"Resolving function: {function_name}") func = getattr(check_funcs, function_name, None) # resolve using predefined checks first if not func: - # resolve using predefined geo checks, requires Databricks serverless or DBR >= 17.1 + # try to resolve using predefined geo checks, requires Databricks serverless or DBR >= 17.1 func = getattr(geo_check_funcs, function_name, None) + if not func and PII_ENABLED: + # try to resolve using predefined pii detection checks + func = getattr(pii_check_funcs, function_name, None) if not func and custom_check_functions: func = custom_check_functions.get(function_name) # returns None if not found if fail_on_missing and not func: diff --git a/tests/e2e/notebooks/pii_detection_notebook.py b/tests/e2e/notebooks/pii_detection_notebook.py index 67b4688f5..df0b640d0 100644 --- a/tests/e2e/notebooks/pii_detection_notebook.py +++ b/tests/e2e/notebooks/pii_detection_notebook.py @@ -33,8 +33,13 @@ def test_import_pii_module_fails_without_installation(): # COMMAND ---------- +import yaml +from databricks.sdk import WorkspaceClient +from databricks.labs.dqx.engine import DQEngine +from databricks.labs.dqx.rule import DQRowRule import pyspark.sql.functions as F from databricks.labs.dqx.pii.nlp_engine_config import NLPEngineConfig +from databricks.labs.dqx.check_funcs import is_not_null from databricks.labs.dqx.pii.pii_detection_funcs import does_not_contain_pii from chispa import assert_df_equality @@ -222,3 +227,75 @@ def test_does_not_contain_pii_with_custom_nlp_config_dict(): assert_df_equality(actual, expected, transforms=transforms) test_does_not_contain_pii_with_custom_nlp_config_dict() + +# COMMAND ---------- +# DBTITLE 1,test_does_not_contain_pii_apply_checks_by_metadata_and_split + +def test_does_not_contain_pii_apply_checks_by_metadata_and_split(): + checks = yaml.safe_load(""" + - criticality: error + check: + function: is_not_null + arguments: + column: val + - criticality: error + check: + function: does_not_contain_pii + arguments: + column: val + """) + + # Initialize the DQX engine: + dq_engine = DQEngine(WorkspaceClient()) + + # Create some sample data: + data = [ + ["My name is John Smith"], + ["The sky is blue, road runner"], + ["Jane Smith sent an email to sara@info.com"], + [None], + ] + df = spark.createDataFrame(data, "val string") + + # Run the checks and display the output: + valid_df, invalid_df = dq_engine.apply_checks_by_metadata_and_split(df, checks) + assert valid_df.count() == 1 + assert invalid_df.count() == 3 + +test_does_not_contain_pii_apply_checks_by_metadata_and_split() + +# COMMAND ---------- +# DBTITLE 1,test_does_not_contain_pii_apply_checks_and_split + +def test_does_not_contain_pii_apply_checks_and_split(): + checks = [ + DQRowRule( + criticality="error", + check_func=is_not_null, + column="val", + ), + DQRowRule( + criticality="error", + check_func=does_not_contain_pii, + column="val", + ) + ] + + # Initialize the DQX engine: + dq_engine = DQEngine(WorkspaceClient()) + + # Create some sample data: + data = [ + ["My name is John Smith"], + ["The sky is blue, road runner"], + ["Jane Smith sent an email to sara@info.com"], + [None], + ] + df = spark.createDataFrame(data, "val string") + + # Run the checks and display the output: + valid_df, invalid_df = dq_engine.apply_checks_and_split(df, checks) + assert valid_df.count() == 1 + assert invalid_df.count() == 3 + +test_does_not_contain_pii_apply_checks_and_split() diff --git a/tests/unit/test_resolve_check_function.py b/tests/unit/test_resolve_check_function.py index 257b346a4..53f81e4aa 100644 --- a/tests/unit/test_resolve_check_function.py +++ b/tests/unit/test_resolve_check_function.py @@ -1,6 +1,9 @@ import textwrap +import sys +import importlib import pytest +from databricks.labs import dqx from databricks.labs.dqx.checks_resolver import resolve_check_function, resolve_custom_check_functions_from_path from databricks.labs.dqx.errors import InvalidCheckError @@ -101,3 +104,31 @@ def test_resolve_custom_check_functions_from_path_with_dependency(tmp_path): func = resolve_custom_check_functions_from_path({"main_func": str(main_path)})["main_func"] assert func() == "dependency ok" + + +def test_pii_module_import_failure(monkeypatch): + """Test that the code handles gracefully when PII module is not available.""" + # Save the original module state + original_pii_module = sys.modules.get('databricks.labs.dqx.pii') + + # Remove the PII module from sys.modules to simulate ImportError + monkeypatch.setitem(sys.modules, 'databricks.labs.dqx.pii', None) + + # Force module reload to trigger the import logic + importlib.reload(dqx.checks_resolver) + + # Verify that PII_ENABLED is False when import fails + assert dqx.checks_resolver.PII_ENABLED is False + + # Verify that standard check functions still work + func = dqx.checks_resolver.resolve_check_function('is_not_null') + assert func is not None + + # Verify that a non-existent function (like PII) cannot be resolved when PII is disabled + func = dqx.checks_resolver.resolve_check_function('some_missing_func', fail_on_missing=False) + assert func is None + + # Restore the module to its original state + if original_pii_module is not None: + sys.modules['databricks.labs.dqx.pii'] = original_pii_module + importlib.reload(dqx.checks_resolver)