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

metadata/pii/processor.py

@@ -1,8 +1,8 @@
-#  Copyright 2021 Collate
-#  Licensed under the Apache License, Version 2.0 (the "License");
+#  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
-#  http://www.apache.org/licenses/LICENSE-2.0
+#  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.
@@ -12,16 +12,10 @@
 """
 Processor util to fetch pii sensitive columns
 """
-import traceback
-from typing import List, Optional, cast
+from typing import Any, Sequence
 
-from metadata.generated.schema.entity.data.table import Column, TableData
-from metadata.generated.schema.entity.services.ingestionPipelines.status import (
-    StackTraceError,
-)
-from metadata.generated.schema.metadataIngestion.databaseServiceAutoClassificationPipeline import (
-    DatabaseServiceAutoClassificationPipeline,
-)
+from metadata.generated.schema.entity.classification.tag import Tag
+from metadata.generated.schema.entity.data.table import Column
 from metadata.generated.schema.metadataIngestion.workflow import (
     OpenMetadataWorkflowConfig,
 )
@@ -31,160 +25,79 @@ from metadata.generated.schema.type.tagLabel import (
     TagLabel,
     TagSource,
 )
-from metadata.ingestion.api.models import Either
-from metadata.ingestion.api.parser import parse_workflow_config_gracefully
-from metadata.ingestion.api.step import Step
-from metadata.ingestion.api.steps import Processor
-from metadata.ingestion.models.table_metadata import ColumnTag
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.pii.algorithms.tags import PIISensitivityTag
+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.pii.scanners.column_name_scanner import ColumnNameScanner
-from metadata.pii.scanners.ner_scanner import NERScanner
-from metadata.sampler.models import SamplerResponse
+from metadata.utils import fqn
 from metadata.utils.logger import profiler_logger
 
 logger = profiler_logger()
 
 
-class PIIProcessor(Processor):
+class PIIProcessor(AutoClassificationProcessor):
     """
-    A scanner that uses Spacy NER for entity recognition
+    An AutoClassificationProcessor that uses a PIISensitive classifier to tag columns.
     """
 
     def __init__(
         self,
         config: OpenMetadataWorkflowConfig,
         metadata: OpenMetadata,
+        tolerance: float = 0.01,
     ):
-        super().__init__()
-        self.config = config
-        self.metadata = metadata
-
-        # Init and type the source config
-        self.source_config: DatabaseServiceAutoClassificationPipeline = cast(
-            DatabaseServiceAutoClassificationPipeline,
-            self.config.source.sourceConfig.config,
-        )  # Used to satisfy type checked
-
-        self._ner_scanner = None
-        self.name_scanner = ColumnNameScanner()
-        self.confidence_threshold = self.source_config.confidence
-
-    @property
-    def name(self) -> str:
-        return "Auto Classification Processor"
-
-    @property
-    def ner_scanner(self) -> NERScanner:
-        """Load the NER Scanner only if called"""
-        if self._ner_scanner is None:
-            self._ner_scanner = NERScanner()
-
-        return self._ner_scanner
-
-    @classmethod
-    def create(
-        cls,
-        config_dict: dict,
-        metadata: OpenMetadata,
-        pipeline_name: Optional[str] = None,
-    ) -> "Step":
-        config = parse_workflow_config_gracefully(config_dict)
-        return cls(config=config, metadata=metadata)
+        super().__init__(config, metadata)
+
+        from metadata.pii.algorithms.classifiers import (  # pylint: disable=import-outside-toplevel
+            ColumnClassifier,
+            PIISensitiveClassifier,
+        )
 
-    def close(self) -> None:
-        """Nothing to close"""
+        self._classifier: ColumnClassifier[PIISensitivityTag] = PIISensitiveClassifier()
+
+        self.confidence_threshold = self.source_config.confidence / 100
+        self._tolerance = tolerance
 
     @staticmethod
-    def build_column_tag(tag_fqn: str, column_fqn: str) -> ColumnTag:
-        """
-        Build the tag and run the PATCH
-        """
+    def build_tag_label(tag: PIISensitivityTag, score: float) -> TagLabel:
+        tag_fqn = fqn.build(
+            metadata=None,
+            entity_type=Tag,
+            classification_name=PII,
+            tag_name=tag.value,
+        )
+
         tag_label = TagLabel(
             tagFQN=tag_fqn,
             source=TagSource.Classification,
             state=State.Suggested,
-            labelType=LabelType.Automated,
+            labelType=LabelType.Generated,
+            reason=build_reason(tag_fqn, score),
         )
 
-        return ColumnTag(column_fqn=column_fqn, tag_label=tag_label)
+        return tag_label
 
-    def process_column(
-        self,
-        idx: int,
-        column: Column,
-        table_data: Optional[TableData],
-        confidence_threshold: float,
-    ) -> Optional[List[ColumnTag]]:
+    def create_column_tag_labels(
+        self, column: Column, sample_data: Sequence[Any]
+    ) -> Sequence[TagLabel]:
         """
-        Tag a column with PII if we find it using our scanners
+        Create tags for the column based on the sample data.
         """
-
-        # First, check if the column we are about to process
-        # already has PII tags or not
-        column_has_pii_tag = any((PII in tag.tagFQN.root for tag in column.tags or []))
-
-        # If it has PII tags, we skip the processing
-        # for the column
-        if column_has_pii_tag is True:
-            return None
-
-        # Scan by column name. If no results there, check the sample data, if any
-        tag_and_confidence = self.name_scanner.scan(column.name.root) or (
-            self.ner_scanner.scan([row[idx] for row in table_data.rows])
-            if table_data
-            else None
+        # If the column we are about to process already has PII tags return empty
+        for tag in column.tags or []:
+            if PII in tag.tagFQN.root:
+                return []
+
+        # Get the tags and confidence
+        scores = self._classifier.predict_scores(
+            sample_data, column_name=column.name.root, column_data_type=column.dataType
         )
 
-        if (
-            tag_and_confidence
-            and tag_and_confidence.tag_fqn
-            and tag_and_confidence.confidence >= confidence_threshold / 100
-        ):
-            # We support returning +1 tags for a single column in _run
-            return [
-                self.build_column_tag(
-                    tag_fqn=tag_and_confidence.tag_fqn,
-                    column_fqn=column.fullyQualifiedName.root,
-                )
-            ]
-
-        return None
-
-    def _run(
-        self,
-        record: SamplerResponse,
-    ) -> Either[SamplerResponse]:
-        """
-        Main entrypoint for the scanner.
-
-        Adds PII tagging based on the column names
-        and TableData
-        """
+        # 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}
 
-        # We don't always need to process
-        if not self.source_config.enableAutoClassification:
-            return Either(right=record)
-
-        column_tags = []
-        for idx, column in enumerate(record.table.columns):
-            try:
-                col_tags = self.process_column(
-                    idx=idx,
-                    column=column,
-                    table_data=record.sample_data.data,
-                    confidence_threshold=self.confidence_threshold,
-                )
-                if col_tags:
-                    column_tags.extend(col_tags)
-            except Exception as err:
-                self.status.failed(
-                    StackTraceError(
-                        name=record.table.fullyQualifiedName.root,
-                        error=f"Error computing PII tags for [{column}] - [{err}]",
-                        stackTrace=traceback.format_exc(),
-                    )
-                )
-
-        record.column_tags = column_tags
-        return Either(right=record)
+        # winner is at most 1 tag
+        winner = get_top_classes(scores, 1, self.confidence_threshold)
+        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/base.py

@@ -1,8 +1,8 @@
-#  Copyright 2021 Collate
-#  Licensed under the Apache License, Version 2.0 (the "License");
+#  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
-#  http://www.apache.org/licenses/LICENSE-2.0
+#  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.

metadata/pii/scanners/column_name_scanner.py

@@ -1,8 +1,8 @@
-#  Copyright 2021 Collate
-#  Licensed under the Apache License, Version 2.0 (the "License");
+#  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
-#  http://www.apache.org/licenses/LICENSE-2.0
+#  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.
@@ -22,16 +22,14 @@ from metadata.utils import fqn
 
 
 class ColumnNameScanner(BaseScanner):
-    """
-    Column Name Scanner to scan column name
-    """
+    """Column Name Scanner to scan column name"""
 
     sensitive_regex = {
         "PASSWORD": re.compile("^.*password.*$", re.IGNORECASE),
         "US_SSN": re.compile("^.*(ssn|social).*$", re.IGNORECASE),
         "CREDIT_CARD": re.compile("^.*(credit).*(card).*$", re.IGNORECASE),
         "BANK_ACCOUNT": re.compile("^.*bank.*(acc|num).*$", re.IGNORECASE),
-        "EMAIL_ADDRESS": re.compile("^.*(email|e-mail|mail).*$", re.IGNORECASE),
+        "EMAIL_ADDRESS": re.compile("^(email|e-mail|mail)(.*address)?$", re.IGNORECASE),
         "USER_NAME": re.compile("^.*(user|client|person).*(name).*$", re.IGNORECASE),
         "PERSON": re.compile(
             "^.*(firstname|lastname|fullname|maidenname|nickname|name_suffix).*$",

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

metadata/pii/scanners/ner_scanner.py

@@ -1,8 +1,8 @@
-#  Copyright 2021 Collate
-#  Licensed under the Apache License, Version 2.0 (the "License");
+#  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
-#  http://www.apache.org/licenses/LICENSE-2.0
+#  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.
@@ -22,6 +22,7 @@ from typing import Any, Dict, List, Optional, Tuple, Union
 from pydantic import BaseModel, ConfigDict
 
 from metadata.generated.schema.entity.classification.tag import Tag
+from metadata.pii.algorithms.presidio_utils import _load_spacy_model
 from metadata.pii.constants import PII, SPACY_EN_MODEL
 from metadata.pii.models import TagAndConfidence
 from metadata.pii.ner import NEREntity
@@ -56,18 +57,10 @@ class NERScanner(BaseScanner):
     """Based on https://microsoft.github.io/presidio/"""
 
     def __init__(self):
-        import spacy
         from presidio_analyzer import AnalyzerEngine
         from presidio_analyzer.nlp_engine.spacy_nlp_engine import SpacyNlpEngine
 
-        try:
-            spacy.load(SPACY_EN_MODEL)
-        except OSError:
-            logger.warning("Downloading en_core_web_md language model for the spaCy")
-            from spacy.cli import download
-
-            download(SPACY_EN_MODEL)
-            spacy.load(SPACY_EN_MODEL)
+        _load_spacy_model(SPACY_EN_MODEL)
 
         nlp_engine_model = NLPEngineModel(
             lang_code=SUPPORTED_LANG, model_name=SPACY_EN_MODEL