openmetadata-ingestion 1.5.11.0__py3-none-any.whl → 1.6.0.0rc1__py3-none-any.whl

metadata/profiler/processor/core.py

@@ -273,9 +273,6 @@ class Profiler(Generic[TMetric]):
 
         Data should be saved under self.results
         """
-
-        logger.debug("Running post Profiler...")
-
         current_col_results: Dict[str, Any] = self._column_results.get(col.name)
         if not current_col_results:
             logger.debug(
@@ -326,6 +323,10 @@ class Profiler(Generic[TMetric]):
     def _prepare_table_metrics(self) -> List:
         """prepare table metrics"""
         metrics = []
+
+        if self.source_config and not self.source_config.computeTableMetrics:
+            return metrics
+
         table_metrics = [
             metric
             for metric in self.metric_filter.static_metrics
@@ -380,6 +381,9 @@ class Profiler(Generic[TMetric]):
     def _prepare_column_metrics(self) -> List:
         """prepare column metrics"""
         column_metrics_for_thread_pool = []
+        if self.source_config and not self.source_config.computeColumnMetrics:
+            return column_metrics_for_thread_pool
+
         columns = [
             column
             for column in self.columns
@@ -613,7 +617,7 @@ class Profiler(Generic[TMetric]):
 
             if self._system_results:
                 system_profile = [
-                    SystemProfile(**system_result)
+                    SystemProfile.model_validate(system_result)
                     for system_result in self._system_results
                 ]
             else:

metadata/profiler/processor/processor.py

@@ -30,7 +30,7 @@ from metadata.ingestion.api.steps import Processor
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.profiler.api.models import ProfilerProcessorConfig, ProfilerResponse
 from metadata.profiler.processor.core import Profiler
-from metadata.profiler.source.metadata import ProfilerSourceAndEntity
+from metadata.profiler.source.model import ProfilerSourceAndEntity
 
 
 class ProfilerProcessor(Processor):

metadata/profiler/processor/sampler/sqlalchemy/snowflake/sampler.py

@@ -13,14 +13,19 @@ Helper module to handle data sampling
 for the profiler
 """
 
-from typing import cast
+from typing import Dict, Optional, cast
 
 from sqlalchemy import Table
 from sqlalchemy.sql.selectable import CTE
 
-from metadata.generated.schema.entity.data.table import ProfileSampleType
+from metadata.generated.schema.entity.data.table import (
+    ProfileSampleType,
+    SamplingMethodType,
+)
+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.utils.constants import SAMPLE_DATA_DEFAULT_COUNT
 
 
 class SnowflakeSampler(SQASampler):
@@ -29,8 +34,26 @@ class SnowflakeSampler(SQASampler):
     run the query in the whole table.
     """
 
-    def __init__(self, *args, **kwargs) -> None:
-        super().__init__(*args, **kwargs)
+    def __init__(
+        self,
+        client,
+        table,
+        profile_sample_config: Optional[ProfileSampleConfig] = None,
+        partition_details: Optional[Dict] = None,
+        profile_sample_query: Optional[str] = None,
+        sample_data_count: Optional[int] = SAMPLE_DATA_DEFAULT_COUNT,
+    ):
+        super().__init__(
+            client,
+            table,
+            profile_sample_config,
+            partition_details,
+            profile_sample_query,
+            sample_data_count,
+        )
+        self.sampling_method_type = SamplingMethodType.BERNOULLI
+        if profile_sample_config and profile_sample_config.sampling_method_type:
+            self.sampling_method_type = profile_sample_config.sampling_method_type
 
     @partition_filter_handler(build_sample=True)
     def get_sample_query(self, *, column=None) -> CTE:
@@ -44,7 +67,7 @@ class SnowflakeSampler(SQASampler):
                     column,
                 )
                 .suffix_with(
-                    f"SAMPLE BERNOULLI ({self.profile_sample or 100})",
+                    f"SAMPLE {self.sampling_method_type.value} ({self.profile_sample or 100})",
                 )
                 .cte(f"{self.table.__tablename__}_rnd")
             )

metadata/profiler/source/{base → database/base}/profiler_source.py

@@ -14,7 +14,7 @@ Base source for the profiler used to instantiate a profiler runner with
 its interface
 """
 from copy import deepcopy
-from typing import List, Optional, Tuple, cast
+from typing import List, Optional, Tuple, Type, cast
 
 from sqlalchemy import MetaData
 
@@ -31,6 +31,7 @@ from metadata.generated.schema.entity.services.databaseService import (
     DatabaseConnection,
     DatabaseService,
 )
+from metadata.generated.schema.entity.services.serviceType import ServiceType
 from metadata.generated.schema.metadataIngestion.databaseServiceProfilerPipeline import (
     DatabaseServiceProfilerPipeline,
 )
@@ -44,9 +45,14 @@ from metadata.profiler.metrics.registry import Metrics
 from metadata.profiler.processor.core import Profiler
 from metadata.profiler.processor.default import DefaultProfiler, get_default_metrics
 from metadata.profiler.source.profiler_source_interface import ProfilerSourceInterface
+from metadata.utils.importer import import_from_module
+from metadata.utils.logger import profiler_logger
+from metadata.utils.service_spec.service_spec import BaseSpec
 
 NON_SQA_DATABASE_CONNECTIONS = (DatalakeConnection,)
 
+logger = profiler_logger()
+
 
 class ProfilerSource(ProfilerSourceInterface):
     """
@@ -61,15 +67,14 @@ class ProfilerSource(ProfilerSourceInterface):
         global_profiler_configuration: ProfilerConfiguration,
     ):
         self.service_conn_config = self._copy_service_config(config, database)
-        self.source_config = config.source.sourceConfig.config
-        self.source_config = cast(
-            DatabaseServiceProfilerPipeline, self.source_config
-        )  # satisfy type checker
+        self.source_config = DatabaseServiceProfilerPipeline.model_validate(
+            config.source.sourceConfig.config
+        )
         self.profiler_config = ProfilerProcessorConfig.model_validate(
             config.processor.model_dump().get("config")
         )
         self.ometa_client = ometa_client
-        self.profiler_interface_type: str = self._get_profiler_interface_type(config)
+        self.profiler_interface_type: str = config.source.type.lower()
         self.sqa_metadata = self._set_sqa_metadata()
         self._interface = None
         self.global_profiler_configuration = global_profiler_configuration
@@ -92,18 +97,6 @@ class ProfilerSource(ProfilerSourceInterface):
             return MetaData()
         return None
 
-    def _get_profiler_interface_type(self, config) -> str:
-        """_summary_
-
-        Args:
-            config (_type_): profiler config
-        Returns:
-            str:
-        """
-        if isinstance(self.service_conn_config, NON_SQA_DATABASE_CONNECTIONS):
-            return self.service_conn_config.__class__.__name__
-        return config.source.serviceConnection.root.config.__class__.__name__
-
     @staticmethod
     def get_config_for_table(entity: Table, profiler_config) -> Optional[TableConfig]:
         """Get config for a specific entity
@@ -196,12 +189,10 @@ class ProfilerSource(ProfilerSourceInterface):
         db_service: Optional[DatabaseService],
     ) -> ProfilerInterface:
         """Create sqlalchemy profiler interface"""
-        from metadata.profiler.interface.profiler_interface_factory import (  # pylint: disable=import-outside-toplevel
-            profiler_interface_factory,
+        profiler_class = self.import_profiler_class(
+            ServiceType.Database, source_type=self.profiler_interface_type
         )
-
-        profiler_interface: ProfilerInterface = profiler_interface_factory.create(
-            self.profiler_interface_type,
+        profiler_interface: ProfilerInterface = profiler_class.create(
             entity,
             schema_entity,
             database_entity,
@@ -217,6 +208,12 @@ class ProfilerSource(ProfilerSourceInterface):
         self.interface = profiler_interface
         return self.interface
 
+    def import_profiler_class(
+        self, service_type: ServiceType, source_type: str
+    ) -> Type[ProfilerInterface]:
+        class_path = BaseSpec.get_for_source(service_type, source_type).profiler_class
+        return cast(Type[ProfilerInterface], import_from_module(class_path))
+
     def _get_context_entities(
         self, entity: Table
     ) -> Tuple[DatabaseSchema, Database, DatabaseService]:

metadata/profiler/source/{bigquery → database/bigquery}/profiler_source.py

@@ -27,7 +27,7 @@ from metadata.generated.schema.security.credentials.gcpValues import (
     MultipleProjectId,
     SingleProjectId,
 )
-from metadata.profiler.source.base.profiler_source import ProfilerSource
+from metadata.profiler.source.database.base.profiler_source import ProfilerSource
 
 
 class BigQueryProfilerSource(ProfilerSource):

metadata/profiler/source/{databricks → database/databricks}/profiler_source.py

@@ -8,9 +8,10 @@ from metadata.generated.schema.metadataIngestion.workflow import (
     OpenMetadataWorkflowConfig,
 )
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
-from metadata.profiler.source.base.profiler_source import ProfilerSource
+from metadata.profiler.source.database.base.profiler_source import ProfilerSource
 
 
+# pylint: disable=unused-argument
 def is_disconnect(self, e, connection, cursor):
     """is_disconnect method for the Databricks dialect"""
     if "Invalid SessionHandle: SessionHandle" in str(e):
@@ -33,8 +34,7 @@ class DataBricksProfilerSource(ProfilerSource):
 
     def set_is_disconnect(self):
         """Set the is_disconnect method for the Databricks dialect"""
-        from databricks.sqlalchemy import (
-            DatabricksDialect,  # pylint: disable=import-outside-toplevel
-        )
+        # pylint: disable=import-outside-toplevel
+        from databricks.sqlalchemy import DatabricksDialect
 
         DatabricksDialect.is_disconnect = is_disconnect

metadata/profiler/source/{mariadb → database/mariadb}/metrics/window/first_quartile.py

@@ -1,7 +1,7 @@
 """Override first quartile metric definition for MariaDB"""
 
 from metadata.profiler.metrics.window.first_quartile import FirstQuartile
-from metadata.profiler.source.mariadb.functions.median import MariaDBMedianFn
+from metadata.profiler.source.database.mariadb.functions.median import MariaDBMedianFn
 
 
 class MariaDBFirstQuartile(FirstQuartile):

metadata/profiler/source/{mariadb → database/mariadb}/metrics/window/median.py

@@ -1,7 +1,7 @@
 """Override first quartile metric definition for MariaDB"""
 
 from metadata.profiler.metrics.window.median import Median
-from metadata.profiler.source.mariadb.functions.median import MariaDBMedianFn
+from metadata.profiler.source.database.mariadb.functions.median import MariaDBMedianFn
 
 
 class MariaDBMedian(Median):

metadata/profiler/source/{mariadb → database/mariadb}/metrics/window/third_quartile.py

@@ -1,7 +1,7 @@
 """Override first quartile metric definition for MariaDB"""
 
 from metadata.profiler.metrics.window.third_quartile import ThirdQuartile
-from metadata.profiler.source.mariadb.functions.median import MariaDBMedianFn
+from metadata.profiler.source.database.mariadb.functions.median import MariaDBMedianFn
 
 
 class MariaDBThirdQuartile(ThirdQuartile):

metadata/profiler/source/{single_store → database/single_store}/metrics/window/first_quartile.py

@@ -1,7 +1,9 @@
 """Override first quartile metric definition for SingleStore"""
 
 from metadata.profiler.metrics.window.first_quartile import FirstQuartile
-from metadata.profiler.source.single_store.functions.median import SingleStoreMedianFn
+from metadata.profiler.source.database.single_store.functions.median import (
+    SingleStoreMedianFn,
+)
 
 
 class SingleStoreFirstQuartile(FirstQuartile):

metadata/profiler/source/{single_store → database/single_store}/metrics/window/median.py

@@ -1,7 +1,9 @@
 """Override first quartile metric definition for SingleStore"""
 
 from metadata.profiler.metrics.window.median import Median
-from metadata.profiler.source.single_store.functions.median import SingleStoreMedianFn
+from metadata.profiler.source.database.single_store.functions.median import (
+    SingleStoreMedianFn,
+)
 
 
 class SingleStoreMedian(Median):

metadata/profiler/source/{single_store → database/single_store}/metrics/window/third_quartile.py

@@ -1,7 +1,9 @@
 """Override first quartile metric definition for SingleStore"""
 
 from metadata.profiler.metrics.window.third_quartile import ThirdQuartile
-from metadata.profiler.source.single_store.functions.median import SingleStoreMedianFn
+from metadata.profiler.source.database.single_store.functions.median import (
+    SingleStoreMedianFn,
+)
 
 
 class SingleStoreThirdQuartile(ThirdQuartile):

metadata/profiler/source/fetcher/entity_fetcher.py

@@ -0,0 +1,62 @@
+#  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.
+"""
+Entity Fetcher
+"""
+
+from typing import Iterator, Optional
+
+from metadata.generated.schema.entity.services.databaseService import DatabaseService
+from metadata.generated.schema.metadataIngestion.workflow import (
+    OpenMetadataWorkflowConfig,
+)
+from metadata.generated.schema.settings.settings import Settings
+from metadata.ingestion.api.models import Either
+from metadata.ingestion.api.status import Status
+from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.profiler.source.fetcher.fetcher_strategy import (
+    DatabaseFetcherStrategy,
+    FetcherStrategy,
+)
+from metadata.profiler.source.model import ProfilerSourceAndEntity
+from metadata.utils.entity_utils import service_class
+
+
+class EntityFetcher:
+    """Entity fetcher context class"""
+
+    def __init__(
+        self,
+        config: OpenMetadataWorkflowConfig,
+        metadata: OpenMetadata,
+        global_profiler_config: Optional[Settings],
+        status: Status,
+    ):
+        self.config = config
+        self.metadata = metadata
+        self.global_profiler_config = global_profiler_config
+        self.status = status
+        self.strategy = self._get_strategy()
+
+    def _get_strategy(self) -> FetcherStrategy:
+        """Get strategy for entity fetcher"""
+        if service_class(self.config.source.type) is DatabaseService:
+            return DatabaseFetcherStrategy(
+                self.config, self.metadata, self.global_profiler_config, self.status
+            )
+
+        raise NotImplementedError(
+            "Fetcher strategy not implemented for this connection type"
+        )
+
+    def fetch(self) -> Iterator[Either[ProfilerSourceAndEntity]]:
+        """Fetch entities"""
+        yield from self.strategy.fetch()

metadata/profiler/source/fetcher/fetcher_strategy.py

@@ -0,0 +1,301 @@
+#  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.
+"""
+Entity Fetcher Strategy
+"""
+
+import traceback
+from abc import ABC, abstractmethod
+from typing import Iterable, Iterator, Optional, cast
+
+from metadata.generated.schema.entity.data.database import Database
+from metadata.generated.schema.entity.services.ingestionPipelines.status import (
+    StackTraceError,
+)
+from metadata.generated.schema.metadataIngestion.databaseServiceProfilerPipeline import (
+    DatabaseServiceProfilerPipeline,
+)
+from metadata.generated.schema.metadataIngestion.workflow import (
+    OpenMetadataWorkflowConfig,
+)
+from metadata.generated.schema.settings.settings import Settings
+from metadata.ingestion.api.models import Either
+from metadata.ingestion.api.status import Status
+from metadata.ingestion.models.entity_interface import EntityInterfaceWithTags
+from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.profiler.source.fetcher.profiler_source_factory import (
+    profiler_source_factory,
+)
+from metadata.profiler.source.model import ProfilerSourceAndEntity
+from metadata.utils.db_utils import Table
+from metadata.utils.filters import (
+    filter_by_classification,
+    filter_by_database,
+    filter_by_schema,
+    filter_by_table,
+)
+
+FIELDS = ["tableProfilerConfig", "columns", "customMetrics", "tags"]
+
+
+class FetcherStrategy(ABC):
+    """Fetcher strategy interface"""
+
+    def __init__(
+        self,
+        config: OpenMetadataWorkflowConfig,
+        metadata: OpenMetadata,
+        global_profiler_config: Optional[Settings],
+        status: Status,
+    ) -> None:
+        self.config = config
+        self.source_config = config.source.sourceConfig.config
+        self.metadata = metadata
+        self.global_profiler_config = global_profiler_config
+        self.status = status
+
+    def filter_classifications(self, entity: EntityInterfaceWithTags) -> bool:
+        """Given a list of entities, filter out entities that do not match the classification filter pattern
+
+        Args:
+            entities (List[EntityInterfaceWithTags]): List of entities to filter implemnenting the `tags` attribute
+
+        Raises:
+            NotImplementedError: Must be implemented by subclass
+        """
+        classification_filter_pattern = getattr(
+            self.source_config, "classificationFilterPattern", None
+        )
+        if not classification_filter_pattern:
+            return False
+
+        use_fqn_for_filtering = getattr(self.source_config, "useFqnForFiltering", False)
+
+        if not entity.tags:
+            # if we are not explicitly including entities with tags we'll add the ones without tags
+            if not classification_filter_pattern.includes:
+                return False
+            return True
+
+        for tag in entity.tags:
+            tag_name = tag.tagFQN.root if use_fqn_for_filtering else tag.name
+            if not tag_name:
+                continue
+            if filter_by_classification(classification_filter_pattern, tag_name):
+                self.status.filter(
+                    tag_name,
+                    f"Classification pattern not allowed for entity {entity.fullyQualifiedName.root}",
+                )  # type: ignore
+                return True
+
+        return False
+
+    @abstractmethod
+    def fetch(self) -> Iterator[Either[ProfilerSourceAndEntity]]:
+        """Fetch entity"""
+        raise NotImplementedError
+
+
+class DatabaseFetcherStrategy(FetcherStrategy):
+    """Database fetcher strategy"""
+
+    def __init__(
+        self,
+        config: OpenMetadataWorkflowConfig,
+        metadata: OpenMetadata,
+        global_profiler_config: Optional[Settings],
+        status: Status,
+    ) -> None:
+        super().__init__(config, metadata, global_profiler_config, status)
+        self.source_config = cast(
+            DatabaseServiceProfilerPipeline, self.source_config
+        )  # Satisfy typchecker
+
+    def _filter_databases(self, databases: Iterable[Database]) -> Iterable[Database]:
+        """Filter databases based on the filter pattern
+
+        Args:
+            database (Database): Database to filter
+
+        Returns:
+            bool
+        """
+        filtered_databases = []
+
+        for database in databases:
+            database_name = database.name.root
+            if database.fullyQualifiedName and self.source_config.useFqnForFiltering:
+                database_name = database.fullyQualifiedName.root
+            if filter_by_database(
+                self.source_config.databaseFilterPattern, database_name
+            ):
+                self.status.filter(
+                    database_name,
+                    f"Database pattern not allowed for database {database_name}",
+                )
+                continue
+            filtered_databases.append(database)
+
+        return filtered_databases
+
+    def _filter_schemas(self, table: Table) -> bool:
+        """Filter tables based on the schema filter pattern
+
+        Args:
+            tables (List[Table]): Tables to filter
+
+        Returns:
+            List[Table]
+        """
+
+        if not table.databaseSchema:
+            return False
+        schema_name = (
+            table.databaseSchema.fullyQualifiedName
+            if self.source_config.useFqnForFiltering
+            else table.databaseSchema.name
+        )
+        if schema_name and filter_by_schema(
+            self.source_config.schemaFilterPattern, schema_name
+        ):
+            self.status.filter(
+                schema_name, f"Schema pattern not allowed for schema {schema_name}"
+            )
+            return True
+
+        return False
+
+    def _filter_tables(self, table: Table) -> bool:
+        """Filter tables based on the table filter pattern
+
+        Args:
+            tables (Iterable[Table]):
+
+        Returns:
+            Iterable[Table]:
+        """
+        table_name = table.name.root
+        if table.fullyQualifiedName and self.source_config.useFqnForFiltering:
+            table_name = table.fullyQualifiedName.root
+
+        if filter_by_table(self.source_config.tableFilterPattern, table_name):
+            self.status.filter(
+                table_name, f"Table pattern not allowed for table {table_name}"
+            )
+            return True
+
+        return False
+
+    def _filter_column_metrics_computation(self):
+        """Filter"""
+
+    def _get_database_entities(self) -> Iterable[Database]:
+        """Get database entities"""
+        if not self.config.source.serviceName:
+            raise ValueError("serviceName must be provided in the source configuration")
+
+        databases = self.metadata.list_all_entities(
+            entity=Database,
+            params={"service": self.config.source.serviceName},
+        )
+        databases = cast(Iterable[Database], databases)
+
+        if self.source_config.databaseFilterPattern:
+            databases = self._filter_databases(databases)
+
+        if not databases:
+            raise ValueError(
+                "databaseFilterPattern returned 0 result. At least 1 database must be returned by the filter pattern."
+                f"\n\t- includes: {self.source_config.databaseFilterPattern.includes if self.source_config.databaseFilterPattern else None}"  # pylint: disable=line-too-long
+                f"\n\t- excludes: {self.source_config.databaseFilterPattern.excludes if self.source_config.databaseFilterPattern else None}"  # pylint: disable=line-too-long
+            )
+
+        return cast(Iterable[Database], databases)
+
+    def _filter_entities(self, tables: Iterable[Table]) -> Iterable[Table]:
+        """Filter entities based on the filter pattern
+
+        Args:
+            entities (Iterable[EntityInterfaceWithTags]): Entities to filter
+
+        Returns:
+            Iterable[EntityInterfaceWithTags]
+        """
+        tables = [
+            table
+            for table in tables
+            if (
+                not self.source_config.schemaFilterPattern
+                or not self._filter_schemas(table)
+            )
+            and (
+                not self.source_config.tableFilterPattern
+                or not self._filter_tables(table)
+            )
+            and (
+                not self.source_config.classificationFilterPattern
+                or not self.filter_classifications(table)
+            )
+        ]
+
+        return tables
+
+    def _get_table_entities(self, database: Database) -> Iterable[Table]:
+        """Given a database, get all table entities
+
+        Args:
+            database (Database): Database to get tables from
+
+        Returns:
+            Iterable[Table]
+        """
+        tables = self.metadata.list_all_entities(
+            entity=Table,
+            fields=FIELDS,
+            params={
+                "service": self.config.source.serviceName,
+                "database": database.fullyQualifiedName.root,  # type: ignore
+            },  # type: ignore
+        )
+        tables = cast(Iterable[Table], tables)
+        tables = self._filter_entities(tables)
+
+        return cast(Iterable[Table], tables)
+
+    def fetch(self) -> Iterator[Either[ProfilerSourceAndEntity]]:
+        """Fetch database entity"""
+        for database in self._get_database_entities():
+            try:
+                profiler_source = profiler_source_factory.create(
+                    self.config.source.type.lower(),
+                    self.config,
+                    database,
+                    self.metadata,
+                    self.global_profiler_config,
+                )
+
+                for table in self._get_table_entities(database):
+                    yield Either(
+                        left=None,
+                        right=ProfilerSourceAndEntity(
+                            profiler_source=profiler_source,
+                            entity=table,
+                        ),
+                    )
+            except Exception as exc:
+                yield Either(
+                    left=StackTraceError(
+                        name=database.fullyQualifiedName.root,  # type: ignore
+                        error=f"Error listing source and entities for database due to [{exc}]",
+                        stackTrace=traceback.format_exc(),
+                    ),
+                    right=None,
+                )