openmetadata-ingestion 1.6.0.0rc1__py3-none-any.whl → 1.6.0.0rc2__py3-none-any.whl

metadata/sampler/processor.py

@@ -0,0 +1,158 @@
+#  Copyright 2021 Collate
+#  Licensed under the Apache License, Version 2.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
+#  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.
+"""
+Data Sampler for the PII Workflow
+"""
+import traceback
+from copy import deepcopy
+from typing import Optional, cast
+
+from metadata.generated.schema.entity.data.database import Database
+from metadata.generated.schema.entity.data.table import Table
+from metadata.generated.schema.entity.services.databaseService import DatabaseConnection
+from metadata.generated.schema.entity.services.ingestionPipelines.status import (
+    StackTraceError,
+)
+from metadata.generated.schema.entity.services.serviceType import ServiceType
+from metadata.generated.schema.metadataIngestion.databaseServiceAutoClassificationPipeline import (
+    DatabaseServiceAutoClassificationPipeline,
+)
+from metadata.generated.schema.metadataIngestion.workflow import (
+    OpenMetadataWorkflowConfig,
+)
+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.ometa.ometa_api import OpenMetadata
+from metadata.profiler.api.models import ProfilerProcessorConfig
+from metadata.profiler.source.metadata import ProfilerSourceAndEntity
+from metadata.sampler.config import get_config_for_table
+from metadata.sampler.models import SampleConfig, SampleData, SamplerResponse
+from metadata.sampler.sampler_interface import SamplerInterface
+from metadata.utils.profiler_utils import get_context_entities
+from metadata.utils.service_spec.service_spec import import_sampler_class
+
+
+class SamplerProcessor(Processor):
+    """Use the profiler interface to fetch the sample data"""
+
+    def __init__(self, config: OpenMetadataWorkflowConfig, metadata: OpenMetadata):
+        super().__init__()
+
+        self.config = config
+        self.metadata = metadata
+
+        self.source_config: DatabaseServiceAutoClassificationPipeline = cast(
+            DatabaseServiceAutoClassificationPipeline,
+            self.config.source.sourceConfig.config,
+        )  # Used to satisfy type checked
+        # We still rely on the orm-processor. We should decouple this in the future
+        self.profiler_config = ProfilerProcessorConfig.model_validate(
+            self.config.processor.model_dump().get("config")
+        )
+
+        self._interface_type: str = config.source.type.lower()
+        self.sampler_class = import_sampler_class(
+            ServiceType.Database, source_type=self._interface_type
+        )
+
+    @property
+    def name(self) -> str:
+        return "Sampler"
+
+    def _run(self, record: ProfilerSourceAndEntity) -> Either[SamplerResponse]:
+        """Fetch the sample data and pass it down the pipeline"""
+
+        try:
+            entity = cast(Table, record.entity)
+            schema_entity, database_entity, _ = get_context_entities(
+                entity=entity, metadata=self.metadata
+            )
+            service_conn_config = self._copy_service_config(
+                self.config, database_entity
+            )
+
+            sampler_interface: SamplerInterface = self.sampler_class.create(
+                service_connection_config=service_conn_config,
+                ometa_client=self.metadata,
+                entity=entity,
+                schema_entity=schema_entity,
+                database_entity=database_entity,
+                table_config=get_config_for_table(entity, self.profiler_config),
+                default_sample_config=SampleConfig(
+                    profile_sample=self.source_config.profileSample,
+                    profile_sample_type=self.source_config.profileSampleType,
+                    sampling_method_type=self.source_config.samplingMethodType,
+                ),
+                default_sample_data_count=self.source_config.sampleDataCount,
+            )
+            sample_data = SampleData(
+                data=sampler_interface.generate_sample_data(),
+                store=self.source_config.storeSampleData,
+            )
+
+            return Either(
+                right=SamplerResponse(
+                    table=entity,
+                    sample_data=sample_data,
+                )
+            )
+
+        except Exception as exc:
+            return Either(
+                left=StackTraceError(
+                    name=record.entity.fullyQualifiedName.root,
+                    error=f"Unexpected exception processing entity {record.entity.fullyQualifiedName.root}: {exc}",
+                    stackTrace=traceback.format_exc(),
+                )
+            )
+
+    @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 _copy_service_config(
+        self, config: OpenMetadataWorkflowConfig, database: Database
+    ) -> DatabaseConnection:
+        """Make a copy of the service config and update the database name
+
+        Args:
+            database (_type_): a database entity
+
+        Returns:
+            DatabaseService.__config__
+        """
+        config_copy = deepcopy(
+            config.source.serviceConnection.root.config  # type: ignore
+        )
+        if hasattr(
+            config_copy,  # type: ignore
+            "supportsDatabase",
+        ):
+            if hasattr(config_copy, "database"):
+                config_copy.database = database.name.root  # type: ignore
+            if hasattr(config_copy, "catalog"):
+                config_copy.catalog = database.name.root  # type: ignore
+
+        # we know we'll only be working with DatabaseConnection, we cast the type to satisfy type checker
+        config_copy = cast(DatabaseConnection, config_copy)
+
+        return config_copy
+
+    def close(self) -> None:
+        """Nothing to close"""

metadata/sampler/sampler_interface.py

@@ -0,0 +1,251 @@
+#  Copyright 2021 Collate
+#  Licensed under the Apache License, Version 2.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
+#  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.
+"""
+Interface for sampler
+"""
+import traceback
+from abc import ABC, abstractmethod
+from typing import Dict, List, Optional, Set, Union
+
+from metadata.generated.schema.entity.data.database import Database
+from metadata.generated.schema.entity.data.databaseSchema import DatabaseSchema
+from metadata.generated.schema.entity.data.table import (
+    ColumnProfilerConfig,
+    Table,
+    TableData,
+)
+from metadata.generated.schema.entity.services.connections.connectionBasicType import (
+    DataStorageConfig,
+)
+from metadata.generated.schema.entity.services.connections.database.datalakeConnection import (
+    DatalakeConnection,
+)
+from metadata.generated.schema.entity.services.databaseService import DatabaseConnection
+from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.profiler.api.models import TableConfig
+from metadata.profiler.processor.sample_data_handler import upload_sample_data
+from metadata.sampler.config import (
+    get_exclude_columns,
+    get_include_columns,
+    get_profile_sample_config,
+    get_sample_data_count_config,
+    get_sample_query,
+)
+from metadata.sampler.models import SampleConfig
+from metadata.sampler.partition import get_partition_details
+from metadata.utils.constants import SAMPLE_DATA_DEFAULT_COUNT
+from metadata.utils.execution_time_tracker import calculate_execution_time
+from metadata.utils.logger import sampler_logger
+from metadata.utils.sqa_like_column import SQALikeColumn
+from metadata.utils.ssl_manager import get_ssl_connection
+
+logger = sampler_logger()
+
+
+class SamplerInterface(ABC):
+    """Sampler interface
+    This should be the entrypoint for computing any metrics that are required downstream for
+    data quality, profiling, etc.
+    """
+
+    # pylint: disable=too-many-instance-attributes, too-many-arguments
+    def __init__(
+        self,
+        service_connection_config: Union[DatabaseConnection, DatalakeConnection],
+        ometa_client: OpenMetadata,
+        entity: Table,
+        include_columns: Optional[List[ColumnProfilerConfig]] = None,
+        exclude_columns: Optional[List[str]] = None,
+        sample_config: SampleConfig = SampleConfig(),
+        partition_details: Optional[Dict] = None,
+        sample_query: Optional[str] = None,
+        storage_config: DataStorageConfig = None,
+        sample_data_count: Optional[int] = SAMPLE_DATA_DEFAULT_COUNT,
+        **__,
+    ):
+        self.ometa_client = ometa_client
+        self._sample = None
+        self._columns: Optional[List[SQALikeColumn]] = None
+        self.sample_config = sample_config
+
+        if not self.sample_config.profile_sample:
+            self.sample_config.profile_sample = 100
+
+        self.entity = entity
+        self.include_columns = include_columns
+        self.exclude_columns = exclude_columns
+        self.sample_query = sample_query
+        self.sample_limit = sample_data_count
+        self.partition_details = partition_details
+        self.storage_config = storage_config
+
+        self.service_connection_config = service_connection_config
+        self.connection = get_ssl_connection(self.service_connection_config)
+        self.client = self.get_client()
+
+    # pylint: disable=too-many-arguments, too-many-locals
+    @classmethod
+    def create(
+        cls,
+        service_connection_config: Union[DatabaseConnection, DatalakeConnection],
+        ometa_client: OpenMetadata,
+        entity: Table,
+        schema_entity: DatabaseSchema,
+        database_entity: Database,
+        table_config: Optional[TableConfig] = None,
+        storage_config: Optional[DataStorageConfig] = None,
+        default_sample_config: Optional[SampleConfig] = None,
+        default_sample_data_count: Optional[int] = SAMPLE_DATA_DEFAULT_COUNT,
+        **kwargs,
+    ) -> "SamplerInterface":
+        """Create sampler"""
+
+        sample_data_count = get_sample_data_count_config(
+            entity=entity,
+            schema_entity=schema_entity,
+            database_entity=database_entity,
+            entity_config=table_config,
+            default_sample_data_count=default_sample_data_count,
+        )
+        sample_config = get_profile_sample_config(
+            entity=entity,
+            schema_entity=schema_entity,
+            database_entity=database_entity,
+            entity_config=table_config,
+            default_sample_config=default_sample_config,
+        )
+        sample_query = get_sample_query(entity=entity, entity_config=table_config)
+        partition_details = get_partition_details(
+            entity=entity, entity_config=table_config
+        )
+        include_columns = get_include_columns(entity, entity_config=table_config)
+        exclude_columns = get_exclude_columns(entity, entity_config=table_config)
+
+        return cls(
+            service_connection_config=service_connection_config,
+            ometa_client=ometa_client,
+            entity=entity,
+            include_columns=include_columns,
+            exclude_columns=exclude_columns,
+            sample_config=sample_config,
+            partition_details=partition_details,
+            sample_query=sample_query,
+            storage_config=storage_config,
+            sample_data_count=sample_data_count,
+            **kwargs,
+        )
+
+    @property
+    def columns(self) -> List[SQALikeColumn]:
+        """
+        Return the list of columns to profile
+        by skipping the columns to ignore.
+        """
+
+        if self._columns:
+            return self._columns
+
+        if self._get_included_columns():
+            self._columns = [
+                column
+                for column in self.get_columns()
+                if column.name in self._get_included_columns()
+            ]
+
+        if not self._get_included_columns():
+            self._columns = [
+                column
+                for column in self._columns or self.get_columns()
+                if column.name not in self._get_excluded_columns()
+            ]
+
+        return self._columns
+
+    def _get_excluded_columns(self) -> Optional[Set[str]]:
+        """Get excluded  columns for table being profiled"""
+        if self.exclude_columns:
+            return set(self.exclude_columns)
+        return set()
+
+    def _get_included_columns(self) -> Optional[Set[str]]:
+        """Get include columns for table being profiled"""
+        if self.include_columns:
+            return {include_col.columnName for include_col in self.include_columns}
+        return set()
+
+    @property
+    @abstractmethod
+    def raw_dataset(self):
+        """Table object to run the sampling"""
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_client(self):
+        """Get client"""
+        raise NotImplementedError
+
+    @abstractmethod
+    def _rdn_sample_from_user_query(self):
+        """Get random sample from user query"""
+        raise NotImplementedError
+
+    @abstractmethod
+    def _fetch_sample_data_from_user_query(self) -> TableData:
+        """Fetch sample data from user query"""
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_dataset(self, **kwargs):
+        """Get random sample"""
+        raise NotImplementedError
+
+    @abstractmethod
+    def fetch_sample_data(self, columns: Optional[List[SQALikeColumn]]) -> TableData:
+        """Fetch sample data
+
+        Args:
+            columns (Optional[List]): List of columns to fetch
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_columns(self) -> List[SQALikeColumn]:
+        """get columns"""
+        raise NotImplementedError
+
+    @calculate_execution_time(store=False)
+    def generate_sample_data(self) -> Optional[TableData]:
+        """Fetch and ingest sample data
+
+        Returns:
+            TableData: sample data
+        """
+        try:
+            logger.debug(
+                f"Fetching sample data for {self.entity.fullyQualifiedName.root}..."
+            )
+            table_data = self.fetch_sample_data(self.columns)
+            # Only store the data if configured to do so
+            if self.storage_config:
+                upload_sample_data(
+                    data=table_data,
+                    entity=self.entity,
+                    sample_storage_config=self.storage_config,
+                )
+            table_data.rows = table_data.rows[
+                : min(SAMPLE_DATA_DEFAULT_COUNT, self.sample_limit)
+            ]
+            return table_data
+
+        except Exception as err:
+            logger.debug(traceback.format_exc())
+            logger.warning(f"Error fetching sample data: {err}")
+            raise err

metadata/{profiler/processor/sampler → sampler}/sqlalchemy/azuresql/sampler.py

@@ -17,7 +17,7 @@ from typing import List, Optional
 from sqlalchemy import Column
 
 from metadata.generated.schema.entity.data.table import TableData
-from metadata.profiler.processor.sampler.sqlalchemy.sampler import SQASampler
+from metadata.sampler.sqlalchemy.sampler import SQASampler
 
 
 class AzureSQLSampler(SQASampler):

metadata/{profiler/processor/sampler → sampler}/sqlalchemy/bigquery/sampler.py

@@ -12,15 +12,28 @@
 Helper module to handle data sampling
 for the profiler
 """
-from typing import Dict, Optional
+from typing import Dict, Optional, Union
 
 from sqlalchemy import Column
+from sqlalchemy import Table as SqaTable
+from sqlalchemy import text
 from sqlalchemy.orm import Query
 
-from metadata.generated.schema.entity.data.table import ProfileSampleType, TableType
-from metadata.profiler.api.models import ProfileSampleConfig
-from metadata.profiler.processor.handle_partition import partition_filter_handler
-from metadata.profiler.processor.sampler.sqlalchemy.sampler import SQASampler
+from metadata.generated.schema.entity.data.table import (
+    ProfileSampleType,
+    Table,
+    TableType,
+)
+from metadata.generated.schema.entity.services.connections.connectionBasicType import (
+    DataStorageConfig,
+)
+from metadata.generated.schema.entity.services.connections.database.datalakeConnection import (
+    DatalakeConnection,
+)
+from metadata.generated.schema.entity.services.databaseService import DatabaseConnection
+from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.sampler.models import SampleConfig
+from metadata.sampler.sqlalchemy.sampler import SQASampler
 from metadata.utils.constants import SAMPLE_DATA_DEFAULT_COUNT
 
 
@@ -33,23 +46,44 @@ class BigQuerySampler(SQASampler):
     # pylint: disable=too-many-arguments
     def __init__(
         self,
-        client,
-        table,
-        profile_sample_config: Optional[ProfileSampleConfig] = None,
+        service_connection_config: Union[DatabaseConnection, DatalakeConnection],
+        ometa_client: OpenMetadata,
+        entity: Table,
+        sample_config: Optional[SampleConfig] = None,
         partition_details: Optional[Dict] = None,
-        profile_sample_query: Optional[str] = None,
+        sample_query: Optional[str] = None,
+        storage_config: DataStorageConfig = None,
         sample_data_count: Optional[int] = SAMPLE_DATA_DEFAULT_COUNT,
         table_type: TableType = None,
+        **kwargs,
     ):
         super().__init__(
-            client,
-            table,
-            profile_sample_config,
-            partition_details,
-            profile_sample_query,
-            sample_data_count,
+            service_connection_config=service_connection_config,
+            ometa_client=ometa_client,
+            entity=entity,
+            sample_config=sample_config,
+            partition_details=partition_details,
+            sample_query=sample_query,
+            storage_config=storage_config,
+            sample_data_count=sample_data_count,
+            **kwargs,
         )
-        self.table_type: TableType = table_type
+        self.raw_dataset_type: TableType = table_type
+
+    def set_tablesample(self, selectable: SqaTable):
+        """Set the TABLESAMPLE clause for BigQuery
+        Args:
+            selectable (Table): Table object
+        """
+        if (
+            self.sample_config.profile_sample_type == ProfileSampleType.PERCENTAGE
+            and self.raw_dataset_type != TableType.View
+        ):
+            return selectable.tablesample(
+                text(f"{self.sample_config.profile_sample or 100} PERCENT")
+            )
+
+        return selectable
 
     def _base_sample_query(self, column: Optional[Column], label=None):
         """Base query for sampling
@@ -73,25 +107,20 @@ class BigQuerySampler(SQASampler):
                 # FROM sample TABLESAMPLE SYSTEM (n PERCENT)
                 column = Column(column_parts[0], STRUCT)
                 # pylint: disable=protected-access
-                column._set_parent(self.table.__table__)
+                column._set_parent(self.raw_dataset.__table__)
                 # pylint: enable=protected-access
 
         return super()._base_sample_query(column, label=label)
 
-    @partition_filter_handler(build_sample=True)
     def get_sample_query(self, *, column=None) -> Query:
         """get query for sample data"""
         # TABLESAMPLE SYSTEM is not supported for views
         if (
-            self.profile_sample_type == ProfileSampleType.PERCENTAGE
-            and self.table_type != TableType.View
+            self.sample_config.profile_sample_type == ProfileSampleType.PERCENTAGE
+            and self.raw_dataset_type != TableType.View
         ):
-            return (
-                self._base_sample_query(column)
-                .suffix_with(
-                    f"TABLESAMPLE SYSTEM ({self.profile_sample or 100} PERCENT)",
-                )
-                .cte(f"{self.table.__tablename__}_sample")
+            return self._base_sample_query(column).cte(
+                f"{self.raw_dataset.__tablename__}_sample"
             )
 
         return super().get_sample_query(column=column)

metadata/sampler/sqlalchemy/postgres/sampler.py

@@ -0,0 +1,91 @@
+#  Copyright 2021 Collate
+#  Licensed under the Apache License, Version 2.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
+#  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.
+"""
+Helper module to handle data sampling for the profiler
+"""
+from typing import Dict, Optional, Union
+
+from sqlalchemy import Table as SqaTable
+from sqlalchemy import func
+from sqlalchemy.orm import Query
+
+from metadata.generated.schema.entity.data.table import ProfileSampleType, Table
+from metadata.generated.schema.entity.services.connections.connectionBasicType import (
+    DataStorageConfig,
+)
+from metadata.generated.schema.entity.services.connections.database.datalakeConnection import (
+    DatalakeConnection,
+)
+from metadata.generated.schema.entity.services.databaseService import DatabaseConnection
+from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.sampler.models import SampleConfig
+from metadata.sampler.sqlalchemy.sampler import SQASampler
+from metadata.sampler.sqlalchemy.snowflake.sampler import SamplingMethodType
+from metadata.utils.constants import SAMPLE_DATA_DEFAULT_COUNT
+
+
+class PostgresSampler(SQASampler):
+    """
+    Generates a sample of the data to not
+    run the query in the whole table.
+    """
+
+    # pylint: disable=too-many-arguments
+    def __init__(
+        self,
+        service_connection_config: Union[DatabaseConnection, DatalakeConnection],
+        ometa_client: OpenMetadata,
+        entity: Table,
+        sample_config: Optional[SampleConfig] = None,
+        partition_details: Optional[Dict] = None,
+        sample_query: Optional[str] = None,
+        storage_config: DataStorageConfig = None,
+        sample_data_count: Optional[int] = SAMPLE_DATA_DEFAULT_COUNT,
+        **kwargs,
+    ):
+        super().__init__(
+            service_connection_config=service_connection_config,
+            ometa_client=ometa_client,
+            entity=entity,
+            sample_config=sample_config,
+            partition_details=partition_details,
+            sample_query=sample_query,
+            storage_config=storage_config,
+            sample_data_count=sample_data_count,
+            **kwargs,
+        )
+        self.sampling_fn = func.bernoulli
+        self.sampling_method_type = SamplingMethodType.BERNOULLI
+        if (
+            sample_config
+            and sample_config.sampling_method_type == SamplingMethodType.SYSTEM
+        ):
+            self.sampling_fn = func.system
+
+    def set_tablesample(self, selectable: SqaTable):
+        """Set the TABLESAMPLE clause for postgres
+        Args:
+            selectable (Table): _description_
+        """
+        if self.sample_config.profile_sample_type == ProfileSampleType.PERCENTAGE:
+            return selectable.tablesample(
+                self.sampling_fn(self.sample_config.profile_sample or 100)
+            )
+
+        return selectable
+
+    def get_sample_query(self, *, column=None) -> Query:
+        if self.sample_config.profile_sample_type == ProfileSampleType.PERCENTAGE:
+            return self._base_sample_query(column).cte(
+                f"{self.raw_dataset.__tablename__}_rnd"
+            )
+
+        return super().get_sample_query(column=column)