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

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

@@ -14,19 +14,14 @@ Base source for the profiler used to instantiate a profiler runner with
 its interface
 """
 from copy import deepcopy
-from typing import List, Optional, Tuple, Type, cast
-
-from sqlalchemy import MetaData
+from typing import Optional, cast
 
 from metadata.generated.schema.configuration.profilerConfiguration import (
     ProfilerConfiguration,
 )
 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
-from metadata.generated.schema.entity.services.connections.database.datalakeConnection import (
-    DatalakeConnection,
-)
+from metadata.generated.schema.entity.data.table import Table
 from metadata.generated.schema.entity.services.databaseService import (
     DatabaseConnection,
     DatabaseService,
@@ -45,11 +40,19 @@ 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.sampler.config import (
+    get_config_for_table,
+    get_exclude_columns,
+    get_include_columns,
+)
+from metadata.sampler.models import SampleConfig
+from metadata.sampler.sampler_interface import SamplerInterface
 from metadata.utils.logger import profiler_logger
-from metadata.utils.service_spec.service_spec import BaseSpec
-
-NON_SQA_DATABASE_CONNECTIONS = (DatalakeConnection,)
+from metadata.utils.profiler_utils import get_context_entities
+from metadata.utils.service_spec.service_spec import (
+    import_profiler_class,
+    import_sampler_class,
+)
 
 logger = profiler_logger()
 
@@ -62,21 +65,21 @@ class ProfilerSource(ProfilerSourceInterface):
     def __init__(
         self,
         config: OpenMetadataWorkflowConfig,
-        database: DatabaseService,
+        database: Database,
         ometa_client: OpenMetadata,
         global_profiler_configuration: ProfilerConfiguration,
     ):
+        self.config = config
         self.service_conn_config = self._copy_service_config(config, database)
-        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 = config.source.type.lower()
-        self.sqa_metadata = self._set_sqa_metadata()
+        self._interface_type: str = config.source.type.lower()
         self._interface = None
+        # We define this in create_profiler_interface to help us reuse
+        # this method for the sampler, which does not have a DatabaseServiceProfilerPipeline
+        self.source_config = None
         self.global_profiler_configuration = global_profiler_configuration
 
     @property
@@ -91,68 +94,8 @@ class ProfilerSource(ProfilerSourceInterface):
         """Set the interface"""
         self._interface = interface
 
-    def _set_sqa_metadata(self):
-        """Set sqlalchemy metadata"""
-        if not isinstance(self.service_conn_config, NON_SQA_DATABASE_CONNECTIONS):
-            return MetaData()
-        return None
-
-    @staticmethod
-    def get_config_for_table(entity: Table, profiler_config) -> Optional[TableConfig]:
-        """Get config for a specific entity
-
-        Args:
-            entity: table entity
-        """
-        for table_config in profiler_config.tableConfig or []:
-            if table_config.fullyQualifiedName.root == entity.fullyQualifiedName.root:
-                return table_config
-
-        for schema_config in profiler_config.schemaConfig or []:
-            if (
-                schema_config.fullyQualifiedName.root
-                == entity.databaseSchema.fullyQualifiedName
-            ):
-                return TableConfig.from_database_and_schema_config(
-                    schema_config, entity.fullyQualifiedName.root
-                )
-        for database_config in profiler_config.databaseConfig or []:
-            if (
-                database_config.fullyQualifiedName.root
-                == entity.database.fullyQualifiedName
-            ):
-                return TableConfig.from_database_and_schema_config(
-                    database_config, entity.fullyQualifiedName.root
-                )
-
-        return None
-
-    def _get_include_columns(
-        self, entity, entity_config: Optional[TableConfig]
-    ) -> Optional[List[ColumnProfilerConfig]]:
-        """get included columns"""
-        if entity_config and entity_config.columnConfig:
-            return entity_config.columnConfig.includeColumns
-
-        if entity.tableProfilerConfig:
-            return entity.tableProfilerConfig.includeColumns
-
-        return None
-
-    def _get_exclude_columns(
-        self, entity, entity_config: Optional[TableConfig]
-    ) -> Optional[List[str]]:
-        """get included columns"""
-        if entity_config and entity_config.columnConfig:
-            return entity_config.columnConfig.excludeColumns
-
-        if entity.tableProfilerConfig:
-            return entity.tableProfilerConfig.excludeColumns
-
-        return None
-
     def _copy_service_config(
-        self, config: OpenMetadataWorkflowConfig, database: DatabaseService
+        self, config: OpenMetadataWorkflowConfig, database: Database
     ) -> DatabaseConnection:
         """Make a copy of the service config and update the database name
 
@@ -189,75 +132,51 @@ class ProfilerSource(ProfilerSourceInterface):
         db_service: Optional[DatabaseService],
     ) -> ProfilerInterface:
         """Create sqlalchemy profiler interface"""
-        profiler_class = self.import_profiler_class(
-            ServiceType.Database, source_type=self.profiler_interface_type
+        self.source_config = DatabaseServiceProfilerPipeline.model_validate(
+            self.config.source.sourceConfig.config
+        )
+        profiler_class = import_profiler_class(
+            ServiceType.Database, source_type=self._interface_type
+        )
+        sampler_class = import_sampler_class(
+            ServiceType.Database, source_type=self._interface_type
         )
+        # This is shared between the sampler and profiler interfaces
+        sampler_interface: SamplerInterface = sampler_class.create(
+            service_connection_config=self.service_conn_config,
+            ometa_client=self.ometa_client,
+            entity=entity,
+            schema_entity=schema_entity,
+            database_entity=database_entity,
+            table_config=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,
+        )
+
         profiler_interface: ProfilerInterface = profiler_class.create(
-            entity,
-            schema_entity,
-            database_entity,
-            db_service,
-            config,
-            profiler_config,
-            self.source_config,
-            self.service_conn_config,
-            self.ometa_client,
-            sqa_metadata=self.sqa_metadata,
+            entity=entity,
+            source_config=self.source_config,
+            service_connection_config=self.service_conn_config,
+            sampler=sampler_interface,
+            ometa_client=self.ometa_client,
         )  # type: ignore
 
         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]:
-        schema_entity = None
-        database_entity = None
-        db_service = None
-
-        if entity.databaseSchema:
-            schema_entity_list = self.ometa_client.es_search_from_fqn(
-                entity_type=DatabaseSchema,
-                fqn_search_string=entity.databaseSchema.fullyQualifiedName,
-                fields="databaseSchemaProfilerConfig",
-            )
-            if schema_entity_list:
-                schema_entity = schema_entity_list[0]
-
-        if entity.database:
-            database_entity_list = self.ometa_client.es_search_from_fqn(
-                entity_type=Database,
-                fqn_search_string=entity.database.fullyQualifiedName,
-                fields="databaseProfilerConfig",
-            )
-            if database_entity_list:
-                database_entity = database_entity_list[0]
-
-        if entity.service:
-            db_service_list = self.ometa_client.es_search_from_fqn(
-                entity_type=DatabaseService,
-                fqn_search_string=entity.service.fullyQualifiedName,
-            )
-            if db_service_list:
-                db_service = db_service_list[0]
-
-        return schema_entity, database_entity, db_service
-
     def get_profiler_runner(
         self, entity: Table, profiler_config: ProfilerProcessorConfig
     ) -> Profiler:
         """
         Returns the runner for the profiler
         """
-        table_config = self.get_config_for_table(entity, profiler_config)
-        schema_entity, database_entity, db_service = self._get_context_entities(
-            entity=entity
+        table_config = get_config_for_table(entity, profiler_config)
+        schema_entity, database_entity, db_service = get_context_entities(
+            entity=entity, metadata=self.ometa_client
         )
         profiler_interface = self.create_profiler_interface(
             entity,
@@ -271,8 +190,8 @@ class ProfilerSource(ProfilerSourceInterface):
         if not profiler_config.profiler:
             return DefaultProfiler(
                 profiler_interface=profiler_interface,
-                include_columns=self._get_include_columns(entity, table_config),
-                exclude_columns=self._get_exclude_columns(entity, table_config),
+                include_columns=get_include_columns(entity, table_config),
+                exclude_columns=get_exclude_columns(entity, table_config),
                 global_profiler_configuration=self.global_profiler_configuration,
                 db_service=db_service,
             )
@@ -290,7 +209,7 @@ class ProfilerSource(ProfilerSourceInterface):
         return Profiler(
             *metrics,  # type: ignore
             profiler_interface=profiler_interface,
-            include_columns=self._get_include_columns(entity, table_config),
-            exclude_columns=self._get_exclude_columns(entity, table_config),
+            include_columns=get_include_columns(entity, table_config),
+            exclude_columns=get_exclude_columns(entity, table_config),
             global_profiler_configuration=self.global_profiler_configuration,
         )

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

@@ -15,10 +15,10 @@ Bigquery Profiler source
 
 from copy import deepcopy
 
+from metadata.generated.schema.entity.data.database import Database
 from metadata.generated.schema.entity.services.connections.database.bigQueryConnection import (
     BigQueryConnection,
 )
-from metadata.generated.schema.entity.services.databaseService import DatabaseService
 from metadata.generated.schema.metadataIngestion.workflow import (
     OpenMetadataWorkflowConfig,
 )
@@ -34,14 +34,14 @@ class BigQueryProfilerSource(ProfilerSource):
     """override the base profiler source to handle BigQuery specific connection configs"""
 
     def _copy_service_config(
-        self, config: OpenMetadataWorkflowConfig, database: DatabaseService
+        self, config: OpenMetadataWorkflowConfig, database: Database
     ) -> BigQueryConnection:
         """Make a copy of the database connection config. If MultiProjectId is used, replace it
         with SingleProjectId with the database name being profiled. We iterate over all non filtered
         database in workflow.py `def execute`.
 
         Args:
-            database (DatabaseService): a database entity
+            database (Database): a database entity
 
         Returns:
             DatabaseConnection

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

@@ -1,9 +1,8 @@
 """Extend the ProfilerSource class to add support for Databricks is_disconnect SQA method"""
-
 from metadata.generated.schema.configuration.profilerConfiguration import (
     ProfilerConfiguration,
 )
-from metadata.generated.schema.entity.services.databaseService import DatabaseService
+from metadata.generated.schema.entity.data.database import Database
 from metadata.generated.schema.metadataIngestion.workflow import (
     OpenMetadataWorkflowConfig,
 )
@@ -25,7 +24,7 @@ class DataBricksProfilerSource(ProfilerSource):
     def __init__(
         self,
         config: OpenMetadataWorkflowConfig,
-        database: DatabaseService,
+        database: Database,
         ometa_client: OpenMetadata,
         global_profiler_config: ProfilerConfiguration,
     ):

metadata/profiler/source/fetcher/config.py

@@ -0,0 +1,44 @@
+#  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 Configuration Models
+"""
+
+from typing import Optional, Protocol, runtime_checkable
+
+from metadata.generated.schema.type.filterPattern import FilterPattern
+
+
+# We take the names from the JSON Schema
+# pylint: disable=invalid-name
+@runtime_checkable
+class EntityFilterConfigInterface(Protocol):
+    """Interface for the OM workflow source configs that allow filtering"""
+
+    @property
+    def classificationFilterPattern(self) -> Optional[FilterPattern]:
+        ...
+
+    @property
+    def databaseFilterPattern(self) -> Optional[FilterPattern]:
+        ...
+
+    @property
+    def schemaFilterPattern(self) -> Optional[FilterPattern]:
+        ...
+
+    @property
+    def tableFilterPattern(self) -> Optional[FilterPattern]:
+        ...
+
+    @property
+    def useFqnForFiltering(self) -> Optional[bool]:
+        ...

metadata/profiler/source/fetcher/fetcher_strategy.py

@@ -20,9 +20,6 @@ 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,
 )
@@ -31,6 +28,7 @@ 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.config import EntityFilterConfigInterface
 from metadata.profiler.source.fetcher.profiler_source_factory import (
     profiler_source_factory,
 )
@@ -116,7 +114,7 @@ class DatabaseFetcherStrategy(FetcherStrategy):
     ) -> None:
         super().__init__(config, metadata, global_profiler_config, status)
         self.source_config = cast(
-            DatabaseServiceProfilerPipeline, self.source_config
+            EntityFilterConfigInterface, self.source_config
         )  # Satisfy typchecker
 
     def _filter_databases(self, databases: Iterable[Database]) -> Iterable[Database]:

metadata/sampler/config.py

@@ -0,0 +1,237 @@
+#  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.
+"""
+Sampler configuration helpers
+"""
+from typing import Any, Dict, List, Optional, Union
+
+from metadata.generated.schema.entity.data.database import (
+    Database,
+    DatabaseProfilerConfig,
+)
+from metadata.generated.schema.entity.data.databaseSchema import (
+    DatabaseSchema,
+    DatabaseSchemaProfilerConfig,
+)
+from metadata.generated.schema.entity.data.table import ColumnProfilerConfig, Table
+from metadata.generated.schema.entity.services.connections.connectionBasicType import (
+    DataStorageConfig,
+)
+from metadata.generated.schema.entity.services.databaseService import DatabaseService
+from metadata.profiler.api.models import ProfilerProcessorConfig
+from metadata.profiler.config import (
+    get_database_profiler_config,
+    get_schema_profiler_config,
+)
+from metadata.sampler.models import DatabaseAndSchemaConfig, SampleConfig, TableConfig
+
+
+def get_sample_storage_config(
+    config: Union[
+        DatabaseSchemaProfilerConfig,
+        DatabaseProfilerConfig,
+        DatabaseAndSchemaConfig,
+    ],
+) -> Optional[Union[DataStorageConfig, Dict[str, Any]]]:
+    """Get sample storage config"""
+    if (
+        config
+        and config.sampleDataStorageConfig
+        and config.sampleDataStorageConfig.config
+    ):
+        return config.sampleDataStorageConfig.config
+    return None
+
+
+def get_storage_config_for_table(
+    entity: Table,
+    schema_entity: DatabaseSchema,
+    database_entity: Database,
+    db_service: Optional[DatabaseService],
+    profiler_config: ProfilerProcessorConfig,
+) -> Optional[Union[DataStorageConfig, Dict[str, Any]]]:
+    """Get storage config for a specific entity"""
+    schema_profiler_config = get_schema_profiler_config(schema_entity=schema_entity)
+    database_profiler_config = get_database_profiler_config(
+        database_entity=database_entity
+    )
+
+    for schema_config in profiler_config.schemaConfig or []:
+        if (
+            entity.databaseSchema
+            and schema_config.fullyQualifiedName.root
+            == entity.databaseSchema.fullyQualifiedName
+            and get_sample_storage_config(schema_config)
+        ):
+            return get_sample_storage_config(schema_config)
+
+    for database_config in profiler_config.databaseConfig or []:
+        if (
+            entity.database
+            and database_config.fullyQualifiedName.root
+            == entity.database.fullyQualifiedName
+            and get_sample_storage_config(database_config)
+        ):
+            return get_sample_storage_config(database_config)
+
+    if schema_profiler_config and get_sample_storage_config(schema_profiler_config):
+        return get_sample_storage_config(schema_profiler_config)
+
+    if database_profiler_config and get_sample_storage_config(database_profiler_config):
+        return get_sample_storage_config(database_profiler_config)
+
+    try:
+        return db_service.connection.config.sampleDataStorageConfig.config
+    except AttributeError:
+        pass
+
+    return None
+
+
+def get_profile_sample_config(
+    entity: Table,
+    schema_entity: Optional[DatabaseSchema],
+    database_entity: Optional[Database],
+    entity_config: Optional[Union[TableConfig, DatabaseAndSchemaConfig]],
+    default_sample_config: Optional[SampleConfig],
+) -> SampleConfig:
+    """Get profile sample config for a specific entity"""
+    schema_profiler_config = get_schema_profiler_config(schema_entity=schema_entity)
+    database_profiler_config = get_database_profiler_config(
+        database_entity=database_entity
+    )
+
+    for config in (
+        entity_config,
+        entity.tableProfilerConfig,
+        schema_profiler_config,
+        database_profiler_config,
+        default_sample_config,
+    ):
+        try:
+            if config and config.profileSample:
+                return SampleConfig(
+                    profile_sample=config.profileSample,
+                    profile_sample_type=config.profileSampleType,
+                    sampling_method_type=config.samplingMethodType,
+                )
+        except AttributeError:
+            pass
+
+    return SampleConfig()
+
+
+def get_sample_query(
+    entity: Table, entity_config: Optional[TableConfig]
+) -> Optional[str]:
+    """get profile query for sampling
+
+    Args:
+        entity (Table): table entity object
+        entity_config (Optional[TableConfig]): entity configuration
+
+    Returns:
+        Optional[str]:
+    """
+    if entity_config:
+        return entity_config.profileQuery
+
+    if entity.tableProfilerConfig:
+        return entity.tableProfilerConfig.profileQuery
+
+    return None
+
+
+def get_sample_data_count_config(
+    entity: Table,
+    schema_entity: Optional[DatabaseSchema],
+    database_entity: Optional[Database],
+    entity_config: Optional[TableConfig],
+    default_sample_data_count: int,
+) -> Optional[int]:
+    """_summary_
+    Args:
+        entity_config (Optional[TableConfig]): table config object from yaml/json file
+        source_config DatabaseServiceProfilerPipeline: profiler pipeline details
+    Returns:
+        Optional[int]: int
+    """
+    schema_profiler_config = get_schema_profiler_config(schema_entity=schema_entity)
+    database_profiler_config = get_database_profiler_config(
+        database_entity=database_entity
+    )
+
+    for config in (
+        entity_config,
+        entity.tableProfilerConfig,
+        schema_profiler_config,
+        database_profiler_config,
+    ):
+        if config and config.sampleDataCount:
+            return config.sampleDataCount
+
+    return default_sample_data_count
+
+
+def get_config_for_table(entity: Table, profiler_config) -> Optional[TableConfig]:
+    """Get config for a specific entity
+
+    Args:
+        entity: table entity
+    """
+    for table_config in profiler_config.tableConfig or []:
+        if table_config.fullyQualifiedName.root == entity.fullyQualifiedName.root:
+            return table_config
+
+    for schema_config in profiler_config.schemaConfig or []:
+        if (
+            schema_config.fullyQualifiedName.root
+            == entity.databaseSchema.fullyQualifiedName
+        ):
+            return TableConfig.from_database_and_schema_config(
+                schema_config, entity.fullyQualifiedName.root
+            )
+    for database_config in profiler_config.databaseConfig or []:
+        if (
+            database_config.fullyQualifiedName.root
+            == entity.database.fullyQualifiedName
+        ):
+            return TableConfig.from_database_and_schema_config(
+                database_config, entity.fullyQualifiedName.root
+            )
+
+    return None
+
+
+def get_include_columns(
+    entity, entity_config: Optional[TableConfig]
+) -> Optional[List[ColumnProfilerConfig]]:
+    """get included columns"""
+    if entity_config and entity_config.columnConfig:
+        return entity_config.columnConfig.includeColumns
+
+    if entity.tableProfilerConfig:
+        return entity.tableProfilerConfig.includeColumns
+
+    return None
+
+
+def get_exclude_columns(
+    entity, entity_config: Optional[TableConfig]
+) -> Optional[List[str]]:
+    """get included columns"""
+    if entity_config and entity_config.columnConfig:
+        return entity_config.columnConfig.excludeColumns
+
+    if entity.tableProfilerConfig:
+        return entity.tableProfilerConfig.excludeColumns
+
+    return None