openmetadata-ingestion 1.7.0.0rc2__py3-none-any.whl → 1.7.1.0__py3-none-any.whl

metadata/pii/algorithms/presidio_utils.py

@@ -0,0 +1,119 @@
+#  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.
+"""
+Utilities for working with the Presidio Library.
+"""
+import inspect
+import logging
+from typing import Iterable, Optional, Type
+
+import spacy
+from presidio_analyzer import (
+    AnalyzerEngine,
+    EntityRecognizer,
+    PatternRecognizer,
+    predefined_recognizers,
+)
+from presidio_analyzer.nlp_engine import SpacyNlpEngine
+from spacy.cli.download import download  # pyright: ignore[reportUnknownVariableType]
+
+from metadata.pii.constants import PRESIDIO_LOGGER, SPACY_EN_MODEL, SUPPORTED_LANG
+from metadata.utils.logger import METADATA_LOGGER, pii_logger
+
+logger = pii_logger()
+
+
+def build_analyzer_engine(
+    model_name: str = SPACY_EN_MODEL,
+) -> AnalyzerEngine:
+    """
+    Build a Presidio analyzer engine for the model_name and tailored to our use case.
+
+    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])
+    analyzer_engine = AnalyzerEngine(
+        nlp_engine=nlp_engine, supported_languages=[SUPPORTED_LANG]
+    )
+    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
+
+
+def set_presidio_logger_level(log_level: Optional[int] = None) -> None:
+    """
+    Set the presidio logger to talk less about internal entities unless we are debugging.
+    """
+    if log_level is None:
+        log_level = (
+            logging.INFO
+            if logging.getLogger(METADATA_LOGGER).level == logging.DEBUG
+            else logging.ERROR
+        )
+
+    logging.getLogger(PRESIDIO_LOGGER).setLevel(log_level)
+
+
+def _load_spacy_model(model_name: str) -> None:
+    """
+    Load the spaCy model for the given language.
+    If the model is not found locally, it will be downloaded.
+    """
+
+    try:
+        _ = spacy.load(model_name)
+    except OSError:
+
+        logger.warning(f"Downloading {model_name} language model for the spaCy")
+        download(model_name)
+        _ = spacy.load(model_name)
+
+
+def _get_all_entity_recognizer_classes() -> Iterable[Type[EntityRecognizer]]:
+    """
+    Iterate over all subclasses of the `EntityRecognizer` exposed
+    in the predefined_recognizers module.
+    """
+    for name in getattr(predefined_recognizers, "__all__", []):
+        obj = getattr(predefined_recognizers, name, None)
+        if inspect.isclass(obj) and issubclass(obj, EntityRecognizer):
+            yield obj
+
+
+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]
+            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()

metadata/pii/algorithms/tags.py

@@ -0,0 +1,111 @@
+#  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.
+"""
+Definition of tags for the PII algorithms.
+These tags currently belong to the layer logic of the algorithms.
+"""
+import enum
+from typing import List
+
+
+class PIISensitivityTag(enum.Enum):
+    SENSITIVE = "Sensitive"
+    NONSENSITIVE = "NonSensitive"
+
+
+@enum.unique
+class PIITag(enum.Enum):
+    """
+    PII Tags (borrowed from Presidio https://microsoft.github.io/presidio/supported_entities/).
+    """
+
+    # Global
+    CREDIT_CARD = "CREDIT_CARD"
+    CRYPTO = "CRYPTO"  # Crypto Wallet Address
+    DATE_TIME = "DATE_TIME"
+    EMAIL_ADDRESS = "EMAIL_ADDRESS"
+    IBAN_CODE = "IBAN_CODE"
+    IP_ADDRESS = "IP_ADDRESS"
+    NRP = "NRP"
+    LOCATION = "LOCATION"
+    PERSON = "PERSON"
+    PHONE_NUMBER = "PHONE_NUMBER"
+    MEDICAL_LICENSE = "MEDICAL_LICENSE"
+    URL = "URL"
+
+    # USA
+    US_BANK_NUMBER = "US_BANK_NUMBER"
+    US_DRIVER_LICENSE = "US_DRIVER_LICENSE"
+    US_ITIN = "US_ITIN"
+    US_PASSPORT = "US_PASSPORT"
+    US_SSN = "US_SSN"
+
+    # UK
+    UK_NHS = "UK_NHS"
+
+    # Spain
+    ES_NIF = "ES_NIF"
+    ES_NIE = "ES_NIE"
+
+    # Italy
+    IT_FISCAL_CODE = "IT_FISCAL_CODE"
+    IT_DRIVER_LICENSE = "IT_DRIVER_LICENSE"
+    IT_VAT_CODE = "IT_VAT_CODE"
+    IT_PASSPORT = "IT_PASSPORT"
+    IT_IDENTITY_CARD = "IT_IDENTITY_CARD"
+
+    # Poland
+    PL_PESEL = "PL_PESEL"
+
+    # Singapore
+    SG_NRIC_FIN = "SG_NRIC_FIN"
+    SG_UEN = "SG_UEN"
+
+    # Australia
+    AU_ABN = "AU_ABN"
+    AU_ACN = "AU_ACN"
+    AU_TFN = "AU_TFN"
+    AU_MEDICARE = "AU_MEDICARE"
+
+    # India
+    IN_PAN = "IN_PAN"
+    IN_AADHAAR = "IN_AADHAAR"
+    IN_VEHICLE_REGISTRATION = "IN_VEHICLE_REGISTRATION"
+    IN_VOTER = "IN_VOTER"
+    IN_PASSPORT = "IN_PASSPORT"
+
+    # Finland
+    FI_PERSONAL_IDENTITY_CODE = "FI_PERSONAL_IDENTITY_CODE"
+
+    @classmethod
+    def values(cls) -> List[str]:
+        """
+        Get all the values of the enum as a set of strings.
+        """
+        return [tag.value for tag in cls]
+
+    def sensitivity(self) -> PIISensitivityTag:
+        """
+        Get the sensitivity level of the PII tag.
+        This map is opinionated and can be changed in the future according to users' needs.
+        """
+        if self in DEFAULT_NON_PII_SENSITIVE:
+            return PIISensitivityTag.NONSENSITIVE
+        return PIISensitivityTag.SENSITIVE
+
+
+DEFAULT_NON_PII_SENSITIVE = (
+    PIITag.DATE_TIME,
+    PIITag.NRP,
+    PIITag.LOCATION,
+    PIITag.PHONE_NUMBER,
+    PIITag.URL,
+)

metadata/pii/algorithms/utils.py

@@ -0,0 +1,38 @@
+#  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.
+"""
+Utility functions for PII algorithms
+"""
+from typing import Mapping, Sequence, TypeVar
+
+T = TypeVar("T")
+
+
+def normalize_scores(scores: Mapping[T, float], tol: float = 0.01) -> Mapping[T, float]:
+    """
+    Normalize the scores to sum to 1, while ignoring scores below the tolerance.
+    Scores must be positive.
+    """
+    scores = {key: score for key, score in scores.items() if score > tol}
+    total = sum(scores.values())
+    if total == 0:
+        return scores
+    return {key: score / total for key, score in scores.items()}
+
+
+def get_top_classes(scores: Mapping[T, float], n: int, threshold: float) -> Sequence[T]:
+    """
+    Get the top n scores from the scores mapping that are above the threshold.
+    The classes are sorted in descending order of their scores.
+    """
+    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]

metadata/pii/base_processor.py

@@ -0,0 +1,125 @@
+#  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.
+"""
+Base class for the Auto Classification Processor.
+"""
+import traceback
+from abc import ABC, abstractmethod
+from typing import Any, Optional, Sequence, Type, TypeVar, cast, final
+
+from metadata.generated.schema.entity.data.table import Column
+from metadata.generated.schema.entity.services.ingestionPipelines.status import (
+    StackTraceError,
+)
+from metadata.generated.schema.metadataIngestion.databaseServiceAutoClassificationPipeline import (
+    DatabaseServiceAutoClassificationPipeline,
+)
+from metadata.generated.schema.metadataIngestion.workflow import (
+    OpenMetadataWorkflowConfig,
+)
+from metadata.generated.schema.type.tagLabel import TagLabel
+from metadata.ingestion.api.models import Either
+from metadata.ingestion.api.parser import parse_workflow_config_gracefully
+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.sampler.models import SamplerResponse
+
+C = TypeVar("C", bound="AutoClassificationProcessor")
+
+
+class AutoClassificationProcessor(Processor, ABC):
+    """
+    Abstract class for the Auto Classification Processor.
+
+    Implementations should only provide the logic for creating tags based on sample data,
+    and rely on the running part to be handled by the base class.
+    """
+
+    # Some methods are marked as final to prevent overriding in subclasses thus
+    # ensuring that the workflow is always run in the same way keeping implementer
+    # with the responsibility of *only* implementing the logic for creating tags.
+    def __init__(
+        self,
+        config: OpenMetadataWorkflowConfig,
+        metadata: OpenMetadata,
+    ):
+        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
+
+    @abstractmethod
+    def create_column_tag_labels(
+        self, column: Column, sample_data: Sequence[Any]
+    ) -> Sequence[TagLabel]:
+        """
+        Create tags for the column based on the sample data.
+        """
+
+    @property
+    def name(self) -> str:
+        return "Auto Classification Processor"
+
+    def close(self) -> None:
+        """Nothing to close"""
+
+    @classmethod
+    @final
+    def create(
+        cls: Type[C],
+        config_dict: dict,
+        metadata: OpenMetadata,
+        pipeline_name: Optional[str] = None,
+    ) -> C:
+        config = parse_workflow_config_gracefully(config_dict)
+        return cls(config=config, metadata=metadata)
+
+    @final
+    def _run(self, record: SamplerResponse) -> Either[SamplerResponse]:
+        """
+        Main entrypoint for the processor.
+        """
+
+        # We don't always need to process
+        if not self.source_config.enableAutoClassification:
+            return Either(right=record, left=None)
+
+        column_tags = []
+
+        for idx, column in enumerate(record.table.columns):
+            try:
+                tags = self.create_column_tag_labels(
+                    column=column,
+                    sample_data=[row[idx] for row in record.sample_data.data.rows],
+                )
+                for tag in tags:
+                    column_tag = ColumnTag(
+                        column_fqn=column.fullyQualifiedName.root, tag_label=tag
+                    )
+                    column_tags.append(column_tag)
+            except Exception as err:
+                # TODO: Shouldn't we return a Left here?
+                self.status.failed(
+                    StackTraceError(
+                        name=record.table.fullyQualifiedName.root,
+                        error=f"Error in Processor {self.name} computing tags for [{column}] - [{err}]",
+                        stackTrace=traceback.format_exc(),
+                    )
+                )
+
+        record.column_tags = column_tags
+        return Either(right=record, left=None)

metadata/pii/constants.py

@@ -13,4 +13,12 @@ PII constants
 """
 
 PII = "PII"
+
+# Constants for Presidio
+PRESIDIO_LOGGER = "presidio-analyzer"
 SPACY_EN_MODEL = "en_core_web_md"
+
+# Supported language for Presidio.
+# Don't change this unless you know what you are doing.
+# We are doing some tricks to make Presidio work for our use case.
+SUPPORTED_LANG = "en"

metadata/pii/processor.py

@@ -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,24 +25,21 @@ 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.classifiers import ColumnClassifier, PIISensitiveClassifier
+from metadata.pii.algorithms.tags import PIISensitivityTag
+from metadata.pii.algorithms.utils import get_top_classes, normalize_scores
+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__(
@@ -56,136 +47,49 @@ class PIIProcessor(Processor):
         config: OpenMetadataWorkflowConfig,
         metadata: OpenMetadata,
     ):
-        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
+        super().__init__(config, metadata)
+        self._classifier: ColumnClassifier[PIISensitivityTag] = PIISensitiveClassifier()
 
-    @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)
-
-    def close(self) -> None:
-        """Nothing to close"""
+        self.confidence_threshold = self.source_config.confidence / 100
+        self._tolerance = 0.01
 
     @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) -> 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,
         )
 
-        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.
         """
+        # 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
+        )
 
-        # 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
-
-        # We'll scan first by sample data to prioritize the NER scanner
-        # If we find nothing, we'll check the column name
-        tag_and_confidence = (
-            self.ner_scanner.scan([row[idx] for row in table_data.rows])
-            if table_data
-            else None
-        ) or self.name_scanner.scan(column.name.root)
-
-        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
-        """
-
-        # 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(),
-                    )
-                )
+        scores = normalize_scores(scores, tol=self._tolerance)
 
-        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) for tag in winner]
+        return tag_labels