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

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

metadata/profiler/interface/sqlalchemy/profiler_interface.py

@@ -18,6 +18,7 @@ supporting sqlalchemy abstraction layer
 import concurrent.futures
 import math
 import threading
+import time
 import traceback
 from collections import defaultdict
 from datetime import datetime
@@ -397,43 +398,72 @@ class SQAProfilerInterface(ProfilerInterface, SQAInterfaceMixin):
             f"Running profiler for {metric_func.table.__tablename__} on thread {threading.current_thread()}"
         )
         Session = self.session_factory  # pylint: disable=invalid-name
-        with Session() as session:
-            self.set_session_tag(session)
-            self.set_catalog(session)
-            runner = self._create_thread_safe_runner(session, metric_func.column)
-            row = None
-            try:
-                row = self._get_metric_fn[metric_func.metric_type.value](
-                    metric_func.metrics,
-                    runner=runner,
-                    session=session,
-                    column=metric_func.column,
-                    sample=runner.dataset,
-                )
-                if isinstance(row, dict):
-                    row = self._validate_nulls(row)
-                if isinstance(row, list):
-                    row = [
-                        self._validate_nulls(r) if isinstance(r, dict) else r
-                        for r in row
-                    ]
-
-            except Exception as exc:
-                error = (
-                    f"{metric_func.column if metric_func.column is not None else metric_func.table.__tablename__} "
-                    f"metric_type.value: {exc}"
-                )
-                logger.error(error)
-                self.status.failed_profiler(error, traceback.format_exc())
-
-            if metric_func.column is not None:
-                column = metric_func.column.name
-                self.status.scanned(f"{metric_func.table.__tablename__}.{column}")
-            else:
-                self.status.scanned(metric_func.table.__tablename__)
-                column = None
+        max_retries = 3
+        retry_count = 0
+        initial_backoff = 5
+        max_backoff = 30
+        row = None
+
+        while retry_count < max_retries:
+            with Session() as session:
+                self.set_session_tag(session)
+                self.set_catalog(session)
+                runner = self._create_thread_safe_runner(session, metric_func.column)
+                try:
+                    row = self._get_metric_fn[metric_func.metric_type.value](
+                        metric_func.metrics,
+                        runner=runner,
+                        session=session,
+                        column=metric_func.column,
+                        sample=runner.dataset,
+                    )
+                    if isinstance(row, dict):
+                        row = self._validate_nulls(row)
+                    if isinstance(row, list):
+                        row = [
+                            self._validate_nulls(r) if isinstance(r, dict) else r
+                            for r in row
+                        ]
+
+                    # On success, log the scan and break out of the retry loop
+                    if metric_func.column is not None:
+                        column = metric_func.column.name
+                        self.status.scanned(
+                            f"{metric_func.table.__tablename__}.{column}"
+                        )
+                    else:
+                        self.status.scanned(metric_func.table.__tablename__)
+                        column = None
+
+                    return row, column, metric_func.metric_type.value
+
+                except Exception as exc:
+                    dialect = session.get_bind().dialect
+                    if dialect.is_disconnect(exc, session.get_bind(), None):
+                        retry_count += 1
+                        if retry_count < max_retries:
+                            backoff = min(
+                                initial_backoff * (2 ** (retry_count - 1)), max_backoff
+                            )
+                            logger.debug(
+                                f"Connection error detected, retrying ({retry_count}/{max_retries}) "
+                                f"after {backoff:.2f} seconds..."
+                            )
+                            session.rollback()
+                            time.sleep(backoff)
+                            continue
+                        logger.error(
+                            f"Max retries ({max_retries}) exceeded for disconnection"
+                        )
+                    error = (
+                        f"{metric_func.column if metric_func.column is not None else metric_func.table.__tablename__} "
+                        f"metric_type.value: {exc}"
+                    )
+                    logger.error(error)
+                    self.status.failed_profiler(error, traceback.format_exc())
 
-            return row, column, metric_func.metric_type.value
+        # If we've exhausted all retries without success, return a tuple of None values
+        return None, None, None
 
     @staticmethod
     def _validate_nulls(row: Dict[str, Any]) -> Dict[str, Any]:

metadata/profiler/processor/runner.py

@@ -23,7 +23,10 @@ from sqlalchemy.orm import DeclarativeMeta, Query, Session
 from sqlalchemy.orm.util import AliasedClass
 
 from metadata.utils.logger import query_runner_logger
-from metadata.utils.sqa_utils import get_query_filter_for_runner
+from metadata.utils.sqa_utils import (
+    get_query_filter_for_runner,
+    get_query_group_by_for_runner,
+)
 
 logger = query_runner_logger()
 
@@ -88,6 +91,11 @@ class QueryRunner:
         """Table name attribute access"""
         return self._session
 
+    @property
+    def dialect(self) -> str:
+        """Dialect attribute access"""
+        return self._session.get_bind().dialect.name
+
     def _build_query(self, *entities, **kwargs) -> Query:
         """Build query object
 
@@ -106,11 +114,15 @@ class QueryRunner:
             **kwargs: kwargs to pass to the query
         """
         filter_ = get_query_filter_for_runner(kwargs)
+        group_by_ = get_query_group_by_for_runner(kwargs)
 
         query = self._build_query(*entities, **kwargs).select_from(self._dataset)
 
         if filter_ is not None:
-            return query.filter(filter_)
+            query = query.filter(filter_)
+
+        if group_by_ is not None:
+            query = query.group_by(*group_by_)
 
         return query
 
@@ -122,7 +134,7 @@ class QueryRunner:
             **kwargs: kwargs to pass to the query
         """
         filter_ = get_query_filter_for_runner(kwargs)
-
+        group_by_ = get_query_group_by_for_runner(kwargs)
         user_query = self._session.query(self._dataset).from_statement(
             text(f"{self.profile_sample_query}")
         )
@@ -130,7 +142,10 @@ class QueryRunner:
         query = self._build_query(*entities, **kwargs).select_from(user_query)
 
         if filter_ is not None:
-            return query.filter(filter_)
+            query = query.filter(filter_)
+
+        if group_by_ is not None:
+            query = query.group_by(*group_by_)
 
         return query
 
@@ -143,13 +158,17 @@ class QueryRunner:
             **kwargs: kwargs to pass to the query
         """
         filter_ = get_query_filter_for_runner(kwargs)
+        group_by_ = get_query_group_by_for_runner(kwargs)
 
         if self.profile_sample_query:
             return self._select_from_user_query(*entities, **kwargs).first()
         query = self._build_query(*entities, **kwargs).select_from(self.table)
 
         if filter_ is not None:
-            return query.filter(filter_).first()
+            query = query.filter(filter_)
+
+        if group_by_ is not None:
+            query = query.group_by(*group_by_)
 
         return query.first()
 
@@ -162,6 +181,7 @@ class QueryRunner:
             **kwargs: kwargs to pass to the query
         """
         filter_ = get_query_filter_for_runner(kwargs)
+        group_by_ = get_query_group_by_for_runner(kwargs)
 
         if self.profile_sample_query:
             return self._select_from_user_query(*entities, **kwargs).all()
@@ -169,7 +189,10 @@ class QueryRunner:
         query = self._build_query(*entities, **kwargs).select_from(self.table)
 
         if filter_ is not None:
-            return query.filter(filter_).all()
+            query = query.filter(filter_)
+
+        if group_by_ is not None:
+            query = query.group_by(*group_by_)
 
         return query.all()
 

metadata/profiler/source/database/mssql/profiler_source.py

@@ -0,0 +1,86 @@
+"""Extend the ProfilerSource class to add support for MSSQL is_disconnect SQA method"""
+from metadata.generated.schema.configuration.profilerConfiguration import (
+    ProfilerConfiguration,
+)
+from metadata.generated.schema.entity.data.database import Database
+from metadata.generated.schema.entity.services.connections.database.mssqlConnection import (
+    MssqlScheme,
+)
+from metadata.generated.schema.metadataIngestion.workflow import (
+    OpenMetadataWorkflowConfig,
+)
+from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.profiler.source.database.base.profiler_source import ProfilerSource
+
+
+def is_disconnect(is_disconnect_original):
+    """Wrapper to add custom is_disconnect method for the MSSQL dialects"""
+
+    def inner_is_disconnect(self, e, connection, cursor):
+        """is_disconnect method for the MSSQL dialects"""
+        error_str = str(e)
+
+        mssql_disconnect_codes = [
+            "08S01",  # Communication link failure
+            "08001",  # Cannot connect
+            "HY000",  # General error often used for connection issues
+        ]
+
+        mssql_disconnect_messages = [
+            "Server closed connection",
+            "ClosedConnectionError",
+            "Connection is closed",
+            "Connection reset by peer",
+            "Timeout expired",
+            "Socket closed",
+        ]
+
+        if any(code in error_str for code in mssql_disconnect_codes) or any(
+            message in error_str for message in mssql_disconnect_messages
+        ):
+            return True
+
+        # If none of our custom checks match, fall back to SQLAlchemy's built-in detection
+        return is_disconnect_original(self, e, connection, cursor)
+
+    return inner_is_disconnect
+
+
+class MssqlProfilerSource(ProfilerSource):
+    """MSSQL Profiler source"""
+
+    def __init__(
+        self,
+        config: OpenMetadataWorkflowConfig,
+        database: Database,
+        ometa_client: OpenMetadata,
+        global_profiler_config: ProfilerConfiguration,
+    ):
+        super().__init__(config, database, ometa_client, global_profiler_config)
+        self.set_is_disconnect(config)
+
+    def set_is_disconnect(self, config: OpenMetadataWorkflowConfig):
+        """Set the is_disconnect method based on the configured connection scheme"""
+        # pylint: disable=import-outside-toplevel
+
+        # Get the configured scheme from the source connection
+        scheme = config.source.serviceConnection.root.config.scheme
+
+        # Set the appropriate is_disconnect method based on the scheme
+        if scheme == MssqlScheme.mssql_pytds:
+            from sqlalchemy_pytds.dialect import MSDialect_pytds
+
+            original_is_disconnect = MSDialect_pytds.is_disconnect
+            MSDialect_pytds.is_disconnect = is_disconnect(original_is_disconnect)
+        elif scheme == MssqlScheme.mssql_pyodbc:
+            from sqlalchemy.dialects.mssql.pyodbc import MSDialect_pyodbc
+
+            original_is_disconnect = MSDialect_pyodbc.is_disconnect
+            MSDialect_pyodbc.is_disconnect = is_disconnect(original_is_disconnect)
+        elif scheme == MssqlScheme.mssql_pymssql:
+            from sqlalchemy.dialects.mssql.pymssql import MSDialect_pymssql
+
+            original_is_disconnect = MSDialect_pymssql.is_disconnect
+            MSDialect_pymssql.is_disconnect = is_disconnect(original_is_disconnect)
+        else:
+            raise ValueError(f"Unsupported MSSQL scheme: {scheme}")

metadata/profiler/source/fetcher/profiler_source_factory.py

@@ -21,6 +21,9 @@ from metadata.generated.schema.entity.services.connections.database.bigQueryConn
 from metadata.generated.schema.entity.services.connections.database.databricksConnection import (
     DatabricksType,
 )
+from metadata.generated.schema.entity.services.connections.database.mssqlConnection import (
+    MssqlType,
+)
 from metadata.profiler.source.profiler_source_interface import ProfilerSourceInterface
 
 
@@ -79,10 +82,20 @@ class ProfilerSourceFactory:
 
         return DataBricksProfilerSource
 
+    @staticmethod
+    def mssql() -> Type[ProfilerSourceInterface]:
+        """Lazy loading of the MSSQL source"""
+        from metadata.profiler.source.database.mssql.profiler_source import (
+            MssqlProfilerSource,
+        )
+
+        return MssqlProfilerSource
+
 
 source = {
     BigqueryType.BigQuery.value.lower(): ProfilerSourceFactory.bigquery,
     DatabricksType.Databricks.value.lower(): ProfilerSourceFactory.databricks,
+    MssqlType.Mssql.value.lower(): ProfilerSourceFactory.mssql,
 }
 
 profiler_source_factory = ProfilerSourceFactory()

metadata/readers/dataframe/json.py

@@ -35,7 +35,11 @@ def _get_json_text(key: str, text: bytes, decode: bool) -> Union[str, bytes]:
         with zipfile.ZipFile(io.BytesIO(text)) as zip_file:
             processed_text = zip_file.read(zip_file.infolist()[0])
     if decode:
-        return processed_text.decode(UTF_8) if isinstance(text, bytes) else text
+        return (
+            processed_text.decode(UTF_8, errors="ignore")
+            if isinstance(text, bytes)
+            else text
+        )
     return processed_text
 
 

metadata/readers/dataframe/parquet.py

@@ -68,8 +68,16 @@ class ParquetDataFrameReader(DataFrameReader):
         from pyarrow.parquet import ParquetDataset
 
         client_kwargs = {
-            "endpoint_override": str(self.config_source.securityConfig.endPointURL),
-            "region": self.config_source.securityConfig.awsRegion,
+            "endpoint_override": (
+                str(self.config_source.securityConfig.endPointURL)
+                if self.config_source.securityConfig.endPointURL
+                else None
+            ),
+            "region": (
+                self.config_source.securityConfig.awsRegion
+                if self.config_source.securityConfig.awsRegion
+                else None
+            ),
             "access_key": self.config_source.securityConfig.awsAccessKeyId,
             "session_token": self.config_source.securityConfig.awsSessionToken,
             "role_arn": self.config_source.securityConfig.assumeRoleArn,

metadata/readers/dataframe/reader_factory.py

@@ -39,6 +39,10 @@ class SupportedTypes(Enum):
     TSV = "tsv"
     AVRO = "avro"
     PARQUET = "parquet"
+    PARQUET_PQ = "pq"
+    PARQUET_PQT = "pqt"
+    PARQUET_PARQ = "parq"
+    PARQUET_SNAPPY = "parquet.snappy"
     JSON = "json"
     JSONGZ = "json.gz"
     JSONZIP = "json.zip"
@@ -52,6 +56,10 @@ DF_READER_MAP = {
     SupportedTypes.TSV.value: TSVDataFrameReader,
     SupportedTypes.AVRO.value: AvroDataFrameReader,
     SupportedTypes.PARQUET.value: ParquetDataFrameReader,
+    SupportedTypes.PARQUET_PQ.value: ParquetDataFrameReader,
+    SupportedTypes.PARQUET_PQT.value: ParquetDataFrameReader,
+    SupportedTypes.PARQUET_PARQ.value: ParquetDataFrameReader,
+    SupportedTypes.PARQUET_SNAPPY.value: ParquetDataFrameReader,
     SupportedTypes.JSON.value: JSONDataFrameReader,
     SupportedTypes.JSONGZ.value: JSONDataFrameReader,
     SupportedTypes.JSONZIP.value: JSONDataFrameReader,

metadata/sampler/processor.py

@@ -99,7 +99,7 @@ class SamplerProcessor(Processor):
                 data=sampler_interface.generate_sample_data(),
                 store=self.source_config.storeSampleData,
             )
-
+            sampler_interface.close()
             return Either(
                 right=SamplerResponse(
                     table=entity,

metadata/sampler/sampler_interface.py

@@ -245,3 +245,6 @@ class SamplerInterface(ABC):
             logger.debug(traceback.format_exc())
             logger.warning(f"Error fetching sample data: {err}")
             raise err
+
+    def close(self):
+        """Default noop"""