openmetadata-ingestion 1.7.0.1__py3-none-any.whl → 1.7.1.0__py3-none-any.whl

metadata/ingestion/source/pipeline/airbyte/metadata.py

@@ -29,7 +29,6 @@ from metadata.generated.schema.entity.data.table import Table
 from metadata.generated.schema.entity.services.connections.pipeline.airbyteConnection import (
     AirbyteConnection,
 )
-from metadata.generated.schema.entity.services.databaseService import DatabaseService
 from metadata.generated.schema.metadataIngestion.workflow import (
     Source as WorkflowSource,
 )
@@ -46,12 +45,16 @@ from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import InvalidSourceException
 from metadata.ingestion.models.pipeline_status import OMetaPipelineStatus
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.ingestion.source.pipeline.openlineage.models import TableDetails
+from metadata.ingestion.source.pipeline.openlineage.utils import FQNNotFoundException
 from metadata.ingestion.source.pipeline.pipeline_service import PipelineServiceSource
 from metadata.utils import fqn
 from metadata.utils.helpers import clean_uri
 from metadata.utils.logger import ingestion_logger
 from metadata.utils.time_utils import convert_timestamp_to_milliseconds
 
+from .utils import get_destination_table_details, get_source_table_details
+
 logger = ingestion_logger()
 
 
@@ -189,6 +192,26 @@ class AirbyteSource(PipelineServiceSource):
                     )
                 )
 
+    def _get_table_fqn(self, table_details: TableDetails) -> Optional[str]:
+        """
+        Get the FQN of the table
+        """
+        try:
+            if self.get_db_service_names():
+                return self._get_table_fqn_from_om(table_details)
+
+            return fqn.build(
+                metadata=self.metadata,
+                entity_type=Table,
+                service_name="*",
+                database_name=table_details.database,
+                schema_name=table_details.schema,
+                table_name=table_details.name,
+            )
+        except FQNNotFoundException:
+            return None
+
+    # pylint: disable=too-many-locals
     def yield_pipeline_lineage_details(
         self, pipeline_details: AirbytePipelineDetails
     ) -> Iterable[Either[AddLineageRequest]]:
@@ -197,47 +220,65 @@ class AirbyteSource(PipelineServiceSource):
         :param pipeline_details: pipeline_details object from airbyte
         :return: Lineage from inlets and outlets
         """
+        pipeline_name = pipeline_details.connection.get("name")
         source_connection = self.client.get_source(
             pipeline_details.connection.get("sourceId")
         )
         destination_connection = self.client.get_destination(
             pipeline_details.connection.get("destinationId")
         )
-        source_service = self.metadata.get_by_name(
-            entity=DatabaseService, fqn=source_connection.get("name")
-        )
-        destination_service = self.metadata.get_by_name(
-            entity=DatabaseService, fqn=destination_connection.get("name")
-        )
-        if not source_service or not destination_service:
-            return
+        source_name = source_connection.get("sourceName")
+        destination_name = destination_connection.get("destinationName")
 
         for task in (
             pipeline_details.connection.get("syncCatalog", {}).get("streams") or []
         ):
             stream = task.get("stream")
-            from_fqn = fqn.build(
-                self.metadata,
-                Table,
-                table_name=stream.get("name"),
-                database_name=None,
-                schema_name=stream.get("namespace"),
-                service_name=source_connection.get("name"),
-            )
 
-            to_fqn = fqn.build(
-                self.metadata,
-                Table,
-                table_name=stream.get("name"),
-                database_name=None,
-                schema_name=stream.get("namespace"),
-                service_name=destination_connection.get("name"),
+            source_table_details = get_source_table_details(stream, source_connection)
+            destination_table_details = get_destination_table_details(
+                stream, destination_connection
             )
 
+            if not source_table_details or not destination_table_details:
+                continue
+
+            from_fqn = self._get_table_fqn(source_table_details)
+            to_fqn = self._get_table_fqn(destination_table_details)
+
+            if not from_fqn:
+                logger.warning(
+                    f"While extracting lineage: [{pipeline_name}],"
+                    f" source table: [{source_table_details.database or '*'}]"
+                    f".[{source_table_details.schema}].[{source_table_details.name}]"
+                    f" (type: {source_name}) not found in openmetadata"
+                )
+                continue
+            if not to_fqn:
+                logger.warning(
+                    f"While extracting lineage: [{pipeline_name}],"
+                    f" destination table: [{destination_table_details.database or '*'}]"
+                    f".[{destination_table_details.schema}].[{destination_table_details.name}]"
+                    f" (type: {destination_name}) not found in openmetadata"
+                )
+                continue
+
             from_entity = self.metadata.get_by_name(entity=Table, fqn=from_fqn)
             to_entity = self.metadata.get_by_name(entity=Table, fqn=to_fqn)
 
-            if not from_entity and not to_entity:
+            if not from_entity:
+                logger.warning(
+                    f"While extracting lineage: [{pipeline_name}],"
+                    f" source table (fqn: [{from_fqn}], type: {source_name}) not found"
+                    " in openmetadata"
+                )
+                continue
+            if not to_entity:
+                logger.warning(
+                    f"While extracting lineage: [{pipeline_name}],"
+                    f" destination table (fqn: [{to_fqn}], type: {destination_name}) not found"
+                    " in openmetadata"
+                )
                 continue
 
             pipeline_fqn = fqn.build(
@@ -279,4 +320,4 @@ class AirbyteSource(PipelineServiceSource):
         """
         Get Pipeline Name
         """
-        return pipeline_details.connection.get("connectionId")
+        return pipeline_details.connection.get("name")

metadata/ingestion/source/pipeline/airbyte/utils.py

@@ -0,0 +1,99 @@
+#  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.
+
+"""
+Utils for Airbyte
+"""
+
+from metadata.ingestion.source.pipeline.openlineage.models import TableDetails
+from metadata.utils.logger import ingestion_logger
+
+from .constants import AirbyteDestination, AirbyteSource
+
+logger = ingestion_logger()
+
+
+def get_source_table_details(stream: dict, source_connection: dict) -> TableDetails:
+    """
+    Get the source table details
+    """
+    source_name = source_connection.get("sourceName")
+    source_database = source_connection.get("connectionConfiguration", {}).get(
+        "database"
+    )
+    source_schema = stream.get("namespace")
+
+    # Check if source is supported
+    if source_name not in [
+        AirbyteSource.POSTGRES.value,
+        AirbyteSource.MSSQL.value,
+        AirbyteSource.MYSQL.value,
+        AirbyteSource.MONGODB.value,
+    ]:
+        logger.warning(
+            f"Lineage of airbyte pipeline with source [{source_name}] is not supported yet"
+        )
+        return None
+
+    # Handle specific database configurations
+    if source_name == AirbyteSource.MYSQL.value:
+        source_schema = source_database
+        source_database = None
+    elif source_name == AirbyteSource.MONGODB.value:
+        source_schema = (
+            source_connection.get("connectionConfiguration", {})
+            .get("database_config", {})
+            .get("database")
+        )
+        source_database = None
+
+    return TableDetails(
+        name=stream["name"],
+        schema=source_schema,
+        database=source_database,
+    )
+
+
+def get_destination_table_details(
+    stream: dict, destination_connection: dict
+) -> TableDetails:
+    """
+    Get the destination table details
+    """
+    destination_name = destination_connection.get("destinationName")
+    destination_database = destination_connection.get(
+        "connectionConfiguration", {}
+    ).get("database")
+    destination_schema = destination_connection.get("connectionConfiguration", {}).get(
+        "schema"
+    )
+
+    # Check if destination is supported
+    if destination_name not in [
+        AirbyteDestination.POSTGRES.value,
+        AirbyteDestination.MSSQL.value,
+        AirbyteDestination.MYSQL.value,
+    ]:
+        logger.warning(
+            f"Lineage of airbyte pipeline with destination [{destination_name}] is not supported yet"
+        )
+        return None
+
+    # Handle specific database configurations
+    if destination_name == AirbyteDestination.MYSQL.value:
+        destination_schema = destination_database
+        destination_database = None
+
+    return TableDetails(
+        name=stream["name"],
+        schema=destination_schema,
+        database=destination_database,
+    )

metadata/ingestion/source/pipeline/openlineage/models.py

@@ -14,7 +14,7 @@ Openlineage Source Model module
 
 from dataclasses import dataclass
 from enum import Enum
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Optional
 
 
 @dataclass
@@ -76,8 +76,9 @@ class TableDetails:
     Minimal table information.
     """
 
-    schema: str
     name: str
+    schema: str
+    database: Optional[str] = None
 
 
 class EventType(str, Enum):

metadata/ingestion/source/pipeline/pipeline_service.py

@@ -319,8 +319,7 @@ class PipelineServiceSource(TopologyRunnerMixin, Source, ABC):
     def _get_table_fqn_from_om(self, table_details: TableDetails) -> Optional[str]:
         """
         Based on partial schema and table names look for matching table object in open metadata.
-        :param schema: schema name
-        :param table: table name
+        :param table_details: TableDetails object containing table name, schema, database information
         :return: fully qualified name of a Table in Open Metadata
         """
         result = None
@@ -330,7 +329,7 @@ class PipelineServiceSource(TopologyRunnerMixin, Source, ABC):
                 metadata=self.metadata,
                 entity_type=Table,
                 service_name=db_service,
-                database_name=None,
+                database_name=table_details.database,
                 schema_name=table_details.schema,
                 table_name=table_details.name,
             )

metadata/ingestion/source/storage/s3/metadata.py

@@ -64,6 +64,7 @@ from metadata.readers.file.config_source_factory import get_reader
 from metadata.utils import fqn
 from metadata.utils.filters import filter_by_container
 from metadata.utils.logger import ingestion_logger
+from metadata.utils.s3_utils import list_s3_objects
 from metadata.utils.tag_utils import get_ometa_tag_and_classification, get_tag_label
 
 logger = ingestion_logger()
@@ -345,14 +346,13 @@ class S3Source(StorageServiceSource):
         try:
             prefix = self._get_sample_file_prefix(metadata_entry=metadata_entry)
             if prefix:
-                response = self.s3_client.list_objects_v2(
-                    Bucket=bucket_response.name, Prefix=prefix
-                )
+                kwargs = {"Bucket": bucket_response.name, "Prefix": prefix}
+                response = list_s3_objects(self.s3_client, **kwargs)
                 # total depth is depth of prefix + depth of the metadata entry
                 total_depth = metadata_entry.depth + len(prefix[:-1].split("/"))
                 candidate_keys = {
                     "/".join(entry.get("Key").split("/")[:total_depth]) + "/"
-                    for entry in response[S3_CLIENT_ROOT_RESPONSE]
+                    for entry in response
                     if entry
                     and entry.get("Key")
                     and len(entry.get("Key").split("/")) > total_depth
@@ -464,12 +464,11 @@ class S3Source(StorageServiceSource):
         parent: Optional[EntityReference] = None,
     ):
         bucket_name = bucket_response.name
-        response = self.s3_client.list_objects_v2(
-            Bucket=bucket_name, Prefix=metadata_entry.dataPath
-        )
+        kwargs = {"Bucket": bucket_name, "Prefix": metadata_entry.dataPath}
+        response = list_s3_objects(self.s3_client, **kwargs)
         candidate_keys = [
             entry["Key"]
-            for entry in response[S3_CLIENT_ROOT_RESPONSE]
+            for entry in response
             if entry and entry.get("Key") and not entry.get("Key").endswith("/")
         ]
         for key in candidate_keys:

metadata/pii/algorithms/classifiers.py

@@ -0,0 +1,180 @@
+#  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.
+"""
+Classifier for PII detection and sensitivity tagging.
+"""
+from abc import ABC, abstractmethod
+from collections import defaultdict
+from typing import (
+    Any,
+    DefaultDict,
+    Dict,
+    Generic,
+    Hashable,
+    Mapping,
+    Optional,
+    Sequence,
+    Set,
+    TypeVar,
+    final,
+)
+
+from presidio_analyzer import AnalyzerEngine
+
+from metadata.generated.schema.entity.data.table import DataType
+from metadata.pii.algorithms.column_patterns import get_pii_column_name_patterns
+from metadata.pii.algorithms.feature_extraction import (
+    extract_pii_from_column_names,
+    extract_pii_tags,
+    is_non_pii_datatype,
+    split_column_name,
+)
+from metadata.pii.algorithms.preprocessing import preprocess_values
+from metadata.pii.algorithms.presidio_patches import url_patcher
+from metadata.pii.algorithms.presidio_utils import (
+    build_analyzer_engine,
+    set_presidio_logger_level,
+)
+from metadata.pii.algorithms.tags import PIISensitivityTag, PIITag
+
+T = TypeVar("T", bound=Hashable)
+
+
+class ColumnClassifier(ABC, Generic[T]):
+    """
+    Base class for column classifiers.
+    This class defines the interface for classifiers that predict the class
+    of a column based on its data and metadata.
+    """
+
+    @abstractmethod
+    def predict_scores(
+        self,
+        sample_data: Sequence[Any],
+        column_name: Optional[str] = None,
+        column_data_type: Optional[DataType] = None,
+    ) -> Mapping[T, float]:
+        """
+        Predict the scores for the given column and sample data of the column.
+        The scores are a mapping of class labels to their respective scores:
+        higher scores indicate a higher likelihood of the class for the given inputs.
+        """
+
+
+# Implementations
+
+
+@final
+class HeuristicPIIClassifier(ColumnClassifier[PIITag]):
+    """
+    Heuristic PII Column Classifier
+    """
+
+    def __init__(
+        self,
+        *,
+        column_name_contribution: float = 0.5,
+        score_cutoff: float = 0.1,
+        relative_cardinality_cutoff: float = 0.01,
+    ):
+        set_presidio_logger_level()
+        self._presidio_analyzer: AnalyzerEngine = build_analyzer_engine()
+        self._column_name_patterns = get_pii_column_name_patterns()
+
+        self._column_name_contribution = column_name_contribution
+        self._score_cutoff = score_cutoff
+        self._relative_cardinality_cutoff = relative_cardinality_cutoff
+
+    def predict_scores(
+        self,
+        sample_data: Sequence[Any],
+        column_name: Optional[str] = None,
+        column_data_type: Optional[DataType] = None,
+    ) -> Mapping[PIITag, float]:
+
+        if column_data_type is not None and is_non_pii_datatype(column_data_type):
+            return {}
+
+        str_values = preprocess_values(sample_data)
+
+        if not str_values:
+            return {}
+
+        # Relative cardinality test
+        unique_values = set(str_values)
+
+        if len(unique_values) / len(str_values) < self._relative_cardinality_cutoff:
+            return {}
+        context = split_column_name(column_name) if column_name else None
+
+        content_results = extract_pii_tags(
+            self._presidio_analyzer,
+            str_values,
+            context=context,
+            recognizer_result_patcher=url_patcher,
+        )
+
+        column_name_matches: Set[PIITag] = set()
+
+        if column_name is not None:
+            column_name_matches = extract_pii_from_column_names(
+                column_name, patterns=self._column_name_patterns
+            )
+
+        final_results: Dict[PIITag, float] = {}
+
+        for tag, score in content_results.items():
+            final_score = score
+            if tag in column_name_matches:
+                final_score += self._column_name_contribution
+            # Apply the score cutoff
+            if final_score >= self._score_cutoff:
+                final_results[tag] = final_score
+
+        return final_results
+
+
+class PIISensitiveClassifier(ColumnClassifier[PIISensitivityTag]):
+    """
+    Implements a classifier for PII sensitivity tags based on a given
+    PII column classifier. If no classifier is provided, it defaults to
+    using the HeuristicPIIColumnClassifier.
+    """
+
+    def __init__(self, classifier: Optional[ColumnClassifier[PIITag]] = None):
+        self.classifier: ColumnClassifier[PIITag] = (
+            classifier or HeuristicPIIClassifier()
+        )
+
+    def predict_scores(
+        self,
+        sample_data: Sequence[Any],
+        column_name: Optional[str] = None,
+        column_data_type: Optional[DataType] = None,
+    ) -> Mapping[PIISensitivityTag, float]:
+        pii_tags = self.classifier.predict_scores(
+            sample_data, column_name, column_data_type
+        )
+        results: DefaultDict[PIISensitivityTag, float] = defaultdict(float)
+        counts: DefaultDict[PIISensitivityTag, int] = defaultdict(int)
+
+        for tag, score in pii_tags.items():
+            # Convert PIITag to PIISensitivityTag
+            pii_sensitivity = tag.sensitivity()
+            results[pii_sensitivity] += score
+            counts[pii_sensitivity] += 1
+
+        # Normalize the scores
+        for tag in results:
+            if counts[tag] > 0:
+                results[tag] /= counts[tag]
+
+        return results

metadata/pii/algorithms/column_patterns.py

@@ -0,0 +1,61 @@
+#  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 custom patterns for the PII detection.
+Only patterns for column names are implemented here; for content,
+we rely on the Presidio library.
+"""
+import re
+from collections import defaultdict
+from functools import lru_cache
+from typing import DefaultDict, List, Mapping, Union
+
+from metadata.pii.algorithms.tags import PIITag
+
+# Regex patterns for PII detection in column names, not for content
+_pii_column_name_regexes: Mapping[PIITag, Union[str, List[str]]] = {
+    PIITag.US_SSN: "^.*(ssn|social).*$",
+    PIITag.CREDIT_CARD: "^.*(credit).*(card).*$",
+    PIITag.US_BANK_NUMBER: [
+        r"\b(account|acct|acc)[_-]?(number|num|no)\b",  # account_number, account_num
+        r"\bbank[_-]?(account|number|num|no)?\b",  # bank_account, bank_number
+    ],
+    PIITag.IBAN_CODE: [
+        r"\b(account|acct|acc)[_-]?(number|num|no)\b",  # account_number, account_num
+        r"\bbank[_-]?(account|number|num|no)?\b",  # bank_account, bank_number
+        r"\biban(?:[_]?(number|code))?\b",  # iban, iban_number, iban_code
+        r"\bbank[_]?iban\b",  # bank_iban
+        r"\binternational[_]?(account|bank[_]?number)\b",  # international_account, international_bank_number
+    ],
+    PIITag.EMAIL_ADDRESS: "^(email|e-mail|mail)(.*address)?$",
+    PIITag.PERSON: "^.*(user|client|person|first|last|maiden|nick).*(name).*$",
+    PIITag.DATE_TIME: "^.*(date|time|dob|birthday|dod).*$",
+    PIITag.NRP: "^.*(gender|nationality).*$",
+    PIITag.LOCATION: "^.*(address|city|state|county|country|zipcode|zip|postal|zone|borough).*$",
+    PIITag.PHONE_NUMBER: "^.*(phone).*$",
+}
+
+
+@lru_cache
+def get_pii_column_name_patterns() -> Mapping[PIITag, List[re.Pattern[str]]]:
+    """
+    Returns the regex patterns for PII detection in column names.
+    The patterns are cached for performance.
+    """
+    patterns: DefaultDict[PIITag, List[re.Pattern[str]]] = defaultdict(list)
+
+    for pii_type, regexes in _pii_column_name_regexes.items():
+        if isinstance(regexes, str):
+            regexes = [regexes]
+        for regex in regexes:
+            patterns[pii_type].append(re.compile(regex, re.IGNORECASE))
+
+    return patterns