add code to handle unspecified excludeName better

acryldata · Dec 16, 2023 · 9fd8fb6 · 9fd8fb6
1 parent 6762f0f
commit 9fd8fb6
Show file tree

Hide file tree

Showing 4 changed files with 52 additions and 4 deletions.
diff --git a/datahub-classify/src/datahub_classify/infotype_predictor.py b/datahub-classify/src/datahub_classify/infotype_predictor.py
@@ -58,8 +58,10 @@ def predict_infotypes(
             config_dict = global_config[infotype]
 
             # convert exclude_name list into a set for o(1) checking
-            if EXCLUDE_NAME in config_dict:
+            if EXCLUDE_NAME in config_dict and config_dict[EXCLUDE_NAME] is not None:
                 config_dict[EXCLUDE_NAME] = set(config_dict[EXCLUDE_NAME])
+            else:
+                config_dict[EXCLUDE_NAME] = set()
 
             # call the infotype prediction function
             column_info.values = [

diff --git a/datahub-classify/src/datahub_classify/infotype_utils.py b/datahub-classify/src/datahub_classify/infotype_utils.py
@@ -100,7 +100,9 @@ def perform_basic_checks(
             f"does not meet minimum threshold for {infotype}"
         )
         basic_checks_status = False
-    elif metadata.name in config_dict.get(EXCLUDE_NAME, set()):
+    elif config_dict[EXCLUDE_NAME] is not None and metadata.name in config_dict.get(
+        EXCLUDE_NAME, set()
+    ):
         logger.warning(f"Excluding match for {infotype} on column {metadata.name}")
         basic_checks_status = False
     # TODO: Add more basic checks

diff --git a/datahub-classify/tests/exclude_name_test_config.py b/datahub-classify/tests/exclude_name_test_config.py
@@ -28,3 +28,32 @@
         },
     },
 }
+
+none_exclude_name_test_config = {
+    "Email_Address": {
+        "Prediction_Factors_and_Weights": {
+            "Name": 1,
+            "Description": 0,
+            "Datatype": 0,
+            "Values": 0,
+        },
+        "ExcludeName": None,
+        "Name": {
+            "regex": [
+                "^.*mail.*id.*$",
+                "^.*id.*mail.*$",
+                "^.*mail.*add.*$",
+                "^.*add.*mail.*$",
+                "email",
+                "mail",
+            ]
+        },
+        "Description": {"regex": []},
+        "Datatype": {"type": ["str"]},
+        "Values": {
+            "prediction_type": "regex",
+            "regex": [],
+            "library": [],
+        },
+    },
+}
diff --git a/datahub-classify/tests/test_info_type_utils.py b/datahub-classify/tests/test_info_type_utils.py
@@ -1,4 +1,7 @@
-from exclude_name_test_config import exclude_name_test_config
+from exclude_name_test_config import (
+    exclude_name_test_config,
+    none_exclude_name_test_config,
+)
 from datahub_classify.helper_classes import ColumnInfo, Metadata
 from datahub_classify.infotype_utils import perform_basic_checks
 
@@ -85,7 +88,7 @@ def column_infos():
     ]
 
 
-def test_perform_basic_checks():
+def test_perform_basic_checks_with_exclude_name():
     for col_data in column_infos():
         result = perform_basic_checks(
             Metadata(col_data.metadata.meta_info),
@@ -98,3 +101,15 @@ def test_perform_basic_checks():
             assert not result
         else:
             assert result
+
+
+def test_perform_basic_checks_with_none_exclude_name():
+    for col_data in column_infos():
+        result = perform_basic_checks(
+            Metadata(col_data.metadata.meta_info),
+            col_data.values,
+            none_exclude_name_test_config["Email_Address"],
+            "Email_Address",
+            1,
+        )
+        assert result