openmetadata-ingestion 1.9.0.0__py3-none-any.whl → 1.10.14.4__py3-none-any.whl

metadata/pii/algorithms/presidio_utils.py

@@ -13,24 +13,48 @@ Utilities for working with the Presidio Library.
 """
 import inspect
 import logging
-from typing import Iterable, Type, Union
+from typing import Any, Callable, Iterable, List, Optional, Set, Type, Union, cast
 
 import spacy
+from dateutil import parser
 from presidio_analyzer import (
     AnalyzerEngine,
     EntityRecognizer,
     PatternRecognizer,
+    RecognizerRegistry,
+    RecognizerResult,
     predefined_recognizers,
 )
-from presidio_analyzer.nlp_engine import SpacyNlpEngine
+from presidio_analyzer.nlp_engine import NlpArtifacts, SpacyNlpEngine
+from presidio_analyzer.predefined_recognizers import (
+    AuTfnRecognizer,
+    CreditCardRecognizer,
+    DateRecognizer,
+    NhsRecognizer,
+    UsLicenseRecognizer,
+)
 from spacy.cli.download import download  # pyright: ignore[reportUnknownVariableType]
 
+from metadata.pii.algorithms import patterns
 from metadata.pii.constants import PRESIDIO_LOGGER, SPACY_EN_MODEL, SUPPORTED_LANG
+from metadata.utils.dispatch import class_register
 from metadata.utils.logger import pii_logger
 
 logger = pii_logger()
 
 
+def load_nlp_engine(
+    model_name: str = SPACY_EN_MODEL,
+    supported_language: str = SUPPORTED_LANG,
+) -> SpacyNlpEngine:
+    _load_spacy_model(model_name)
+    model = {
+        "lang_code": supported_language,
+        "model_name": model_name,
+    }
+    return SpacyNlpEngine(models=[model])
+
+
 def build_analyzer_engine(
     model_name: str = SPACY_EN_MODEL,
 ) -> AnalyzerEngine:
@@ -39,27 +63,18 @@ def build_analyzer_engine(
 
     If the model is not found locally, it will be downloaded.
     """
-    _load_spacy_model(model_name)
-
-    model = {
-        "lang_code": SUPPORTED_LANG,
-        "model_name": model_name,
-    }
-
-    nlp_engine = SpacyNlpEngine(models=[model])
+    nlp_engine = load_nlp_engine(
+        model_name=model_name, supported_language=SUPPORTED_LANG
+    )
+    recognizer_registry = RecognizerRegistry(
+        recognizers=list(_get_all_pattern_recognizers()),
+        supported_languages=[SUPPORTED_LANG],
+    )
     analyzer_engine = AnalyzerEngine(
-        nlp_engine=nlp_engine, supported_languages=[SUPPORTED_LANG]
+        nlp_engine=nlp_engine,
+        supported_languages=[SUPPORTED_LANG],
+        registry=recognizer_registry,
     )
-    for recognizer in _get_all_pattern_recognizers():
-        # Register the recognizer by setting the appropriate language.
-        # Presidio recognizers are language-dependent: when analyzing text,
-        # Presidio filters recognizers based on the specified language, assuming
-        # language-specific patterns (e.g., for country-specific formats).
-        # However, our use case involves analyzing structured table data rather than free text,
-        # so this language-based approach doesn't always make sense.
-        # To fix this, we manually set the recognizer supported language to the one we want.
-        recognizer.supported_language = SUPPORTED_LANG
-        analyzer_engine.registry.add_recognizer(recognizer)
 
     return analyzer_engine
 
@@ -96,16 +111,182 @@ def _get_all_entity_recognizer_classes() -> Iterable[Type[EntityRecognizer]]:
             yield obj
 
 
+recognizer_factories = class_register()
+
+
+class SanitizedCreditCardRecognizer(CreditCardRecognizer):
+    def analyze(
+        self,
+        text: str,
+        entities: List[str],
+        nlp_artifacts: Optional[NlpArtifacts] = None,
+        regex_flags: Optional[int] = None,
+    ) -> List[RecognizerResult]:
+        return super().analyze(
+            self.sanitize_value(text, self.replacement_pairs),
+            entities,
+            nlp_artifacts,
+            regex_flags,
+        )
+
+
+@recognizer_factories.add(  # pyright: ignore[reportUnknownMemberType, reportUntypedFunctionDecorator]
+    CreditCardRecognizer
+)
+def credit_card_factory(**kwargs: Any) -> SanitizedCreditCardRecognizer:
+    return SanitizedCreditCardRecognizer(
+        patterns=patterns.credit_cards,
+        **kwargs,
+    )
+
+
+@recognizer_factories.add(  # pyright: ignore[reportUnknownMemberType, reportUntypedFunctionDecorator]
+    UsLicenseRecognizer
+)
+def us_license_factory(**kwargs: Any) -> UsLicenseRecognizer:
+    return UsLicenseRecognizer(patterns=patterns.us_driving_license, **kwargs)
+
+
+@recognizer_factories.add(  # pyright: ignore[reportUnknownMemberType, reportUntypedFunctionDecorator]
+    AuTfnRecognizer
+)
+def au_tfn_factory(**kwargs: Any) -> AuTfnRecognizer:
+    return AuTfnRecognizer(
+        patterns=patterns.au_tfn_number,
+        **kwargs,
+    )
+
+
+class ContextAwareNhsRecognizer(NhsRecognizer):
+    TIMESTAMP_KEYWORDS: Set[str] = {
+        "time",
+        "timestamp",
+        "date",
+        "created",
+        "updated",
+        "modified",
+        "deleted",
+        "at",
+        "on",
+        "when",
+        "epoch",
+        "unix",
+        "millis",
+        "seconds",
+        "utc",
+        "gmt",
+        "datetime",
+        "expired",
+        "expires",
+        "transaction",
+        "logged",
+        "recorded",
+        "started",
+        "ended",
+        "finished",
+    }
+
+    def enhance_using_context(
+        self,
+        text: str,
+        raw_recognizer_results: List[RecognizerResult],
+        other_raw_recognizer_results: List[RecognizerResult],
+        nlp_artifacts: NlpArtifacts,
+        context: Optional[List[str]] = None,
+    ) -> List[RecognizerResult]:
+        """Enhance confidence score using context of the entity.
+
+        Filter out NHS number false positives when context suggests
+        the column contains timestamp data.
+
+        :param text: The actual text that was analyzed
+        :param raw_recognizer_results: This recognizer's results, to be updated
+        based on recognizer specific context.
+        :param other_raw_recognizer_results: Other recognizer results matched in
+        the given text to allow related entity context enhancement
+        :param nlp_artifacts: The nlp artifacts contains elements
+                              such as lemmatized tokens for better
+                              accuracy of the context enhancement process
+        :param context: list of context words
+        """
+        if context is None:
+            return raw_recognizer_results
+
+        if self._is_timestamp_context(context):
+            return []
+
+        return raw_recognizer_results
+
+    def _is_timestamp_context(self, context: List[str]) -> bool:
+        """Check if the context contains timestamp-related keywords."""
+        context_lower = {word.lower() for word in context}
+        return bool(context_lower & self.TIMESTAMP_KEYWORDS)
+
+
+@recognizer_factories.add(  # pyright: ignore[reportUnknownMemberType, reportUntypedFunctionDecorator]
+    NhsRecognizer
+)
+def nhs_recognizer(**kwargs: Any) -> NhsRecognizer:
+    return ContextAwareNhsRecognizer(**kwargs)
+
+
+class ValidatedDateRecognizer(DateRecognizer):
+    def validate_result(self, pattern_text: str) -> Optional[bool]:
+        try:
+            _ = parser.parse(pattern_text)
+        except Exception as e:
+            logger.debug(f"Failed to parse {pattern_text}: {e}")
+            # Return None so score isn't modified, relying on Regex score
+            return None
+
+        # Return True so score is boosted to 1.0
+        return True
+
+
+@recognizer_factories.add(  # pyright: ignore[reportUnknownMemberType, reportUntypedFunctionDecorator]
+    DateRecognizer
+)
+def date_recognizer(**kwargs: Any) -> ValidatedDateRecognizer:
+    return ValidatedDateRecognizer(**kwargs)
+
+
 def _get_all_pattern_recognizers() -> Iterable[EntityRecognizer]:
     for cls in _get_all_entity_recognizer_classes():
         if issubclass(cls, PatternRecognizer):
             try:
                 # Try to instantiate the recognizer
-                yield cls(
-                    supported_language=SUPPORTED_LANG
-                )  # pyright: ignore[reportCallIssue]
+                factory = cast(
+                    Callable[..., PatternRecognizer],
+                    recognizer_factories.get(  # pyright: ignore[reportUnknownMemberType]
+                        cls, cls
+                    ),
+                )
+                yield factory(supported_language=SUPPORTED_LANG)
             except Exception as e:
                 logger.warning(e)
         elif cls == predefined_recognizers.PhoneRecognizer:
             # Not a pattern recognizer, but pretty much the same
             yield predefined_recognizers.PhoneRecognizer()
+        elif issubclass(cls, predefined_recognizers.SpacyRecognizer):
+            yield cls(supported_language=SUPPORTED_LANG)
+
+
+def apply_confidence_threshold(
+    threshold: float,
+) -> Callable[[EntityRecognizer], EntityRecognizer]:
+    def decorate_entity_recognizer(recognizer: EntityRecognizer) -> EntityRecognizer:
+        original_analyze = recognizer.analyze
+
+        def analyze(
+            instance: EntityRecognizer,  # pyright: ignore[reportUnusedParameter]
+            text: str,
+            entities: List[str],
+            nlp_artifacts: NlpArtifacts,
+        ) -> List[RecognizerResult]:
+            results = original_analyze(text, entities, nlp_artifacts)
+            return [result for result in results if result.score >= threshold]
+
+        recognizer.analyze = analyze.__get__(recognizer, type(recognizer))
+        return recognizer
+
+    return decorate_entity_recognizer

metadata/pii/algorithms/utils.py

@@ -36,3 +36,7 @@ def get_top_classes(scores: Mapping[T, float], n: int, threshold: float) -> Sequ
     sorted_scores = sorted(scores.items(), key=lambda x: x[1], reverse=True)
     top_classes = [key for key, score in sorted_scores if score >= threshold]
     return top_classes[:n]
+
+
+def build_reason(tag_name: str, score: float) -> str:
+    return f"Chose {tag_name} with a classification score of {score:.2f}"

metadata/pii/base_processor.py

@@ -100,7 +100,8 @@ class AutoClassificationProcessor(Processor, ABC):
 
         column_tags = []
 
-        for idx, column in enumerate(record.table.columns):
+        for idx, column_name in enumerate(record.sample_data.data.columns):
+            column = next(c for c in record.table.columns if c.name == column_name)
             try:
                 tags = self.create_column_tag_labels(
                     column=column,

metadata/pii/processor.py

@@ -27,7 +27,7 @@ from metadata.generated.schema.type.tagLabel import (
 )
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.pii.algorithms.tags import PIISensitivityTag
-from metadata.pii.algorithms.utils import get_top_classes, normalize_scores
+from metadata.pii.algorithms.utils import build_reason, get_top_classes
 from metadata.pii.base_processor import AutoClassificationProcessor
 from metadata.pii.constants import PII
 from metadata.utils import fqn
@@ -45,6 +45,7 @@ class PIIProcessor(AutoClassificationProcessor):
         self,
         config: OpenMetadataWorkflowConfig,
         metadata: OpenMetadata,
+        tolerance: float = 0.01,
     ):
         super().__init__(config, metadata)
 
@@ -56,10 +57,10 @@ class PIIProcessor(AutoClassificationProcessor):
         self._classifier: ColumnClassifier[PIISensitivityTag] = PIISensitiveClassifier()
 
         self.confidence_threshold = self.source_config.confidence / 100
-        self._tolerance = 0.01
+        self._tolerance = tolerance
 
     @staticmethod
-    def build_tag_label(tag: PIISensitivityTag) -> TagLabel:
+    def build_tag_label(tag: PIISensitivityTag, score: float) -> TagLabel:
         tag_fqn = fqn.build(
             metadata=None,
             entity_type=Tag,
@@ -72,6 +73,7 @@ class PIIProcessor(AutoClassificationProcessor):
             source=TagSource.Classification,
             state=State.Suggested,
             labelType=LabelType.Generated,
+            reason=build_reason(tag_fqn, score),
         )
 
         return tag_label
@@ -92,9 +94,10 @@ class PIIProcessor(AutoClassificationProcessor):
             sample_data, column_name=column.name.root, column_data_type=column.dataType
         )
 
-        scores = normalize_scores(scores, tol=self._tolerance)
+        # Filter noise and cap at 1.0 (don't normalize to sum=1)
+        scores = {k: min(v, 1.0) for k, v in scores.items() if v > self._tolerance}
 
         # winner is at most 1 tag
         winner = get_top_classes(scores, 1, self.confidence_threshold)
-        tag_labels = [self.build_tag_label(tag) for tag in winner]
+        tag_labels = [self.build_tag_label(tag, scores[tag]) for tag in winner]
         return tag_labels

metadata/pii/processor_factory.py

@@ -0,0 +1,28 @@
+# pyright: reportUnknownMemberType=false, reportUnknownVariableType=false
+from typing import Any
+
+from metadata.generated.schema.metadataIngestion.workflow import (
+    OpenMetadataWorkflowConfig,
+)
+from metadata.ingestion.api.parser import parse_workflow_config_gracefully
+from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.pii.base_processor import AutoClassificationProcessor
+from metadata.pii.processor import PIIProcessor
+from metadata.pii.tag_processor import TagAnalyzerGenerator, TagProcessor
+
+
+def create_pii_processor(
+    metadata: OpenMetadata[Any, Any], openmetadata_config: OpenMetadataWorkflowConfig
+) -> AutoClassificationProcessor:
+    if getattr(openmetadata_config.processor, "type") == "tag-pii-processor":
+        return TagProcessor(
+            config=parse_workflow_config_gracefully(openmetadata_config.model_dump()),
+            metadata=metadata,
+            generate_tag_analyzers=TagAnalyzerGenerator(
+                metadata=metadata,
+            ),
+        )
+    return PIIProcessor.create(
+        openmetadata_config.model_dump(),
+        metadata,
+    )

metadata/pii/scanners/custom_ner_scanner.py

@@ -0,0 +1,290 @@
+#  Copyright 2025 Collate
+#  Licensed under the Collate Community License, Version 1.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  https://github.com/open-metadata/OpenMetadata/blob/main/ingestion/LICENSE
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+"""
+Enhanced NER Scanner that uses custom recognizers from OpenMetadata classifications.
+"""
+from typing import Dict, List, Optional, Set, Tuple
+
+from presidio_analyzer import RecognizerResult
+
+from metadata.generated.schema.entity.classification.classification import (
+    AutoClassificationConfig,
+    Classification,
+)
+from metadata.generated.schema.entity.classification.tag import Tag
+from metadata.generated.schema.type.tagLabel import TagLabel
+from metadata.pii.algorithms.presidio_recognizer_factory import RecognizerRegistry
+from metadata.pii.algorithms.presidio_utils import build_analyzer_engine
+from metadata.pii.constants import SUPPORTED_LANG
+from metadata.utils.logger import pii_logger
+
+logger = pii_logger()
+
+
+class CustomNERScanner:
+    """NER Scanner that uses custom recognizers from OpenMetadata configurations."""
+
+    def __init__(
+        self,
+        classifications: List[Classification],
+        tags: List[Tag],
+        model_name: Optional[str] = None,
+    ):
+        """
+        Initialize the scanner with classifications and tags.
+
+        Args:
+            classifications: List of classifications with auto-classification config
+            tags: List of tags with recognizer configurations
+            model_name: Optional spaCy model name
+        """
+        self.classifications = {c.fullyQualifiedName: c for c in classifications}
+        self.tags = {t.fullyQualifiedName: t for t in tags}
+        self.tag_by_classification: Dict[
+            str, List[Tag]
+        ] = self._group_tags_by_classification()
+
+        # Build base analyzer engine
+        self.analyzer_engine = (
+            build_analyzer_engine(model_name) if model_name else build_analyzer_engine()
+        )
+
+        # Register custom recognizers
+        self.recognizer_registry = RecognizerRegistry()
+        self._register_custom_recognizers()
+
+    def _group_tags_by_classification(self) -> Dict[str, List[Tag]]:
+        """Group tags by their classification."""
+        grouped = {}
+        for tag in self.tags.values():
+            if tag.classification:
+                classification_fqn = tag.classification.fullyQualifiedName
+                if classification_fqn not in grouped:
+                    grouped[classification_fqn] = []
+                grouped[classification_fqn].append(tag)
+        return grouped
+
+    def _register_custom_recognizers(self) -> None:
+        """Register all custom recognizers from tags."""
+        for tag in self.tags.values():
+            if tag.autoClassificationEnabled and tag.recognizers:
+                self.recognizer_registry.register_tag_recognizers(tag)
+
+                # Add recognizers to the analyzer engine
+                for recognizer in self.recognizer_registry.get_recognizers_for_tag(
+                    tag.fullyQualifiedName
+                ):
+                    recognizer.supported_language = SUPPORTED_LANG
+                    self.analyzer_engine.registry.add_recognizer(recognizer)
+                    logger.info(
+                        f"Registered recognizer {recognizer.name} for tag {tag.fullyQualifiedName}"
+                    )
+
+    def scan_text(
+        self, text: str, classification_fqn: Optional[str] = None
+    ) -> List[TagLabel]:
+        """
+        Scan text for PII using custom recognizers.
+
+        Args:
+            text: Text to scan
+            classification_fqn: Optional classification to limit scanning to
+
+        Returns:
+            List of TagLabel objects for detected entities
+        """
+        if not text:
+            return []
+
+        # Analyze text with all recognizers
+        results = self.analyzer_engine.analyze(
+            text=text,
+            language=SUPPORTED_LANG,
+            return_decision_process=True,
+        )
+
+        # Map results to tags
+        tag_labels = self._map_results_to_tags(results, classification_fqn)
+
+        # Apply conflict resolution if needed
+        if classification_fqn:
+            classification = self.classifications.get(classification_fqn)
+            if classification and classification.mutuallyExclusive:
+                tag_labels = self._resolve_conflicts(
+                    tag_labels, classification.autoClassificationConfig
+                )
+
+        return tag_labels
+
+    def scan_column_name(
+        self, column_name: str, classification_fqn: Optional[str] = None
+    ) -> List[TagLabel]:
+        """
+        Scan a column name for patterns indicating PII.
+
+        Args:
+            column_name: The column name to scan
+            classification_fqn: Optional classification to limit scanning to
+
+        Returns:
+            List of TagLabel objects for detected entities
+        """
+        # Column names often have patterns like "user_ssn", "email_address"
+        # We can use the same scanning but with adjusted confidence
+        results = self.scan_text(column_name, classification_fqn)
+
+        # Boost confidence for column name matches
+        for label in results:
+            if label.confidence:
+                label.confidence = min(1.0, label.confidence * 1.2)
+
+        return results
+
+    def _map_results_to_tags(
+        self, results: List[RecognizerResult], classification_fqn: Optional[str] = None
+    ) -> List[TagLabel]:
+        """
+        Map Presidio recognizer results to OpenMetadata tag labels.
+
+        Args:
+            results: List of recognizer results from Presidio
+            classification_fqn: Optional classification to filter tags
+
+        Returns:
+            List of TagLabel objects
+        """
+        tag_labels = []
+        detected_entities: Set[Tuple[str, float]] = set()
+
+        for result in results:
+            # Find tags that have recognizers for this entity type
+            matching_tags = self._find_matching_tags(
+                result.entity_type, classification_fqn
+            )
+
+            for tag in matching_tags:
+                # Check if confidence meets threshold
+                confidence_threshold = (
+                    self.recognizer_registry.get_tag_confidence_threshold(
+                        tag.fullyQualifiedName
+                    )
+                )
+
+                if result.score >= confidence_threshold:
+                    # Create tag label
+                    tag_label = TagLabel(
+                        tagFQN=tag.fullyQualifiedName,
+                        confidence=result.score,
+                        source="auto-classification",
+                        detectedEntity=result.entity_type,
+                    )
+                    tag_labels.append(tag_label)
+                    detected_entities.add((tag.fullyQualifiedName, result.score))
+
+        return tag_labels
+
+    def _find_matching_tags(
+        self, entity_type: str, classification_fqn: Optional[str] = None
+    ) -> List[Tag]:
+        """
+        Find tags that have recognizers for the given entity type.
+
+        Args:
+            entity_type: The entity type detected by Presidio
+            classification_fqn: Optional classification to filter tags
+
+        Returns:
+            List of matching tags
+        """
+        matching_tags = []
+
+        tags_to_check = []
+        if classification_fqn and classification_fqn in self.tag_by_classification:
+            tags_to_check = self.tag_by_classification[classification_fqn]
+        else:
+            tags_to_check = list(self.tags.values())
+
+        for tag in tags_to_check:
+            if not tag.autoClassificationEnabled or not tag.recognizers:
+                continue
+
+            for recognizer in tag.recognizers:
+                if recognizer.recognizerConfig and hasattr(
+                    recognizer.recognizerConfig, "supportedEntity"
+                ):
+                    if recognizer.recognizerConfig.supportedEntity == entity_type:
+                        matching_tags.append(tag)
+                        break
+
+        return matching_tags
+
+    def _resolve_conflicts(
+        self,
+        tag_labels: List[TagLabel],
+        config: Optional[AutoClassificationConfig],
+    ) -> List[TagLabel]:
+        """
+        Resolve conflicts when multiple tags match for a mutually exclusive classification.
+
+        Args:
+            tag_labels: List of detected tag labels
+            config: Auto-classification configuration
+
+        Returns:
+            List with at most one tag label (the winner of conflict resolution)
+        """
+        if not tag_labels or not config:
+            return tag_labels
+
+        if len(tag_labels) <= 1:
+            return tag_labels
+
+        # Apply conflict resolution strategy
+        if config.conflictResolution == "highest_confidence":
+            # Sort by confidence and return the highest
+            sorted_labels = sorted(
+                tag_labels, key=lambda x: x.confidence or 0, reverse=True
+            )
+            return [sorted_labels[0]]
+
+        elif config.conflictResolution == "highest_priority":
+            # Sort by tag priority
+            sorted_labels = sorted(
+                tag_labels,
+                key=lambda x: self.recognizer_registry.get_tag_priority(x.tagFQN),
+                reverse=True,
+            )
+            return [sorted_labels[0]]
+
+        elif config.conflictResolution == "most_specific":
+            # For now, use confidence as a proxy for specificity
+            # In the future, we could use pattern complexity or other metrics
+            sorted_labels = sorted(
+                tag_labels, key=lambda x: x.confidence or 0, reverse=True
+            )
+            return [sorted_labels[0]]
+
+        # Default: return highest confidence
+        sorted_labels = sorted(
+            tag_labels, key=lambda x: x.confidence or 0, reverse=True
+        )
+        return [sorted_labels[0]]
+
+    def get_supported_classifications(self) -> List[str]:
+        """Get list of classifications that have auto-classification enabled."""
+        enabled_classifications = []
+        for classification in self.classifications.values():
+            if (
+                classification.autoClassificationConfig
+                and classification.autoClassificationConfig.enabled
+            ):
+                enabled_classifications.append(classification.fullyQualifiedName)
+        return enabled_classifications