openmetadata-ingestion 1.5.13.2__py3-none-any.whl → 1.6.0.0__py3-none-any.whl

metadata/profiler/source/metadata_ext.py

@@ -21,7 +21,7 @@ and sample data for that identified entity.
 """
 import traceback
 from copy import deepcopy
-from typing import Iterable, cast
+from typing import Iterable, Type, cast
 
 from sqlalchemy.inspection import inspect
 
@@ -29,6 +29,7 @@ from metadata.generated.schema.entity.data.database import Database
 from metadata.generated.schema.entity.services.ingestionPipelines.status import (
     StackTraceError,
 )
+from metadata.generated.schema.entity.services.serviceType import ServiceType
 from metadata.generated.schema.metadataIngestion.databaseServiceMetadataPipeline import (
     DatabaseServiceMetadataPipeline,
 )
@@ -40,16 +41,16 @@ from metadata.generated.schema.metadataIngestion.workflow import (
 )
 from metadata.ingestion.api.models import Either
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
-from metadata.profiler.source.metadata import (
-    OpenMetadataSource,
-    ProfilerSourceAndEntity,
-)
-from metadata.profiler.source.profiler_source_factory import profiler_source_factory
+from metadata.profiler.interface.profiler_interface import ProfilerInterface
+from metadata.profiler.source.metadata import OpenMetadataSource
+from metadata.profiler.source.model import ProfilerSourceAndEntity
 from metadata.utils import fqn
 from metadata.utils.class_helper import get_service_type_from_source_type
 from metadata.utils.filters import filter_by_database, filter_by_schema, filter_by_table
-from metadata.utils.importer import import_source_class
+from metadata.utils.importer import import_from_module
 from metadata.utils.logger import profiler_logger
+from metadata.utils.service_spec import BaseSpec
+from metadata.utils.service_spec.service_spec import import_source_class
 from metadata.utils.ssl_manager import get_ssl_connection
 
 logger = profiler_logger()
@@ -65,17 +66,12 @@ class OpenMetadataSourceExt(OpenMetadataSource):
     We do this here as well.
     """
 
-    # pylint: disable=super-init-not-called
     def __init__(
         self,
         config: OpenMetadataWorkflowConfig,
         metadata: OpenMetadata,
     ):
-        self.init_steps()
-
-        self.config = config
-        self.metadata = metadata
-        self.test_connection()
+        super().__init__(config, metadata)
 
         # Init and type the source config
         self.service_connection = self.config.source.serviceConnection.root.config
@@ -145,8 +141,7 @@ class OpenMetadataSourceExt(OpenMetadataSource):
                             )
                             continue
 
-                        profiler_source = profiler_source_factory.create(
-                            self.config.source.type.lower(),
+                        profiler_source = self.import_profiler_interface()(
                             self.config,
                             database_entity,
                             self.metadata,
@@ -174,6 +169,14 @@ class OpenMetadataSourceExt(OpenMetadataSource):
                 continue
             yield table_name
 
+    def import_profiler_interface(self) -> Type[ProfilerInterface]:
+        class_path = BaseSpec.get_for_source(
+            ServiceType.Database,
+            source_type=self.config.source.type.lower(),
+        ).profiler_class
+        profiler_source_class = import_from_module(class_path)
+        return cast(Type[ProfilerInterface], profiler_source_class)
+
     def get_schema_names(self) -> Iterable[str]:
         if self.service_connection.__dict__.get("databaseSchema"):
             yield self.service_connection.databaseSchema
@@ -206,9 +209,11 @@ class OpenMetadataSourceExt(OpenMetadataSource):
                         )
                         if filter_by_database(
                             self.source_config.databaseFilterPattern,
-                            database_fqn
-                            if self.source_config.useFqnForFiltering
-                            else database,
+                            (
+                                database_fqn
+                                if self.source_config.useFqnForFiltering
+                                else database
+                            ),
                         ):
                             self.status.filter(database, "Database pattern not allowed")
                             continue

metadata/profiler/source/model.py

@@ -0,0 +1,29 @@
+#  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.
+"""Model for the OpenMetadata Profiler Source"""
+from pydantic import ConfigDict
+
+from metadata.ingestion.models.custom_pydantic import BaseModel
+from metadata.ingestion.models.entity_interface import EntityInterface
+from metadata.profiler.source.profiler_source_interface import ProfilerSourceInterface
+
+
+class ProfilerSourceAndEntity(BaseModel):
+    """Return class for the OpenMetadata Profiler Source"""
+
+    model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True)
+
+    profiler_source: ProfilerSourceInterface
+    entity: EntityInterface
+
+    def __str__(self):
+        """Return the information of the table being profiler"""
+        return f"Table [{self.entity.name.root}]"

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

metadata/sampler/models.py

@@ -0,0 +1,106 @@
+#  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.
+"""
+Sampling Models
+"""
+from typing import List, Optional, Union
+
+from pydantic import Field
+from typing_extensions import Annotated
+
+from metadata.config.common import ConfigModel
+from metadata.generated.schema.entity.data.table import (
+    ColumnProfilerConfig,
+    PartitionProfilerConfig,
+    ProfileSampleType,
+    SamplingMethodType,
+    Table,
+    TableData,
+)
+from metadata.generated.schema.entity.services.connections.connectionBasicType import (
+    SampleDataStorageConfig,
+)
+from metadata.generated.schema.type.basic import FullyQualifiedEntityName
+from metadata.ingestion.models.custom_pydantic import BaseModel
+from metadata.ingestion.models.table_metadata import ColumnTag
+
+
+class BaseProfileConfig(ConfigModel):
+    """base profile config"""
+
+    fullyQualifiedName: FullyQualifiedEntityName
+    profileSample: Optional[Union[float, int]] = None
+    profileSampleType: Optional[ProfileSampleType] = None
+    samplingMethodType: Optional[SamplingMethodType] = None
+    sampleDataCount: Optional[int] = 100
+
+
+class ColumnConfig(ConfigModel):
+    """Column config for profiler"""
+
+    excludeColumns: Optional[List[str]] = None
+    includeColumns: Optional[List[ColumnProfilerConfig]] = None
+
+
+class TableConfig(BaseProfileConfig):
+    """table profile config"""
+
+    profileQuery: Optional[str] = None
+    partitionConfig: Optional[PartitionProfilerConfig] = None
+    columnConfig: Optional[ColumnConfig] = None
+
+    @classmethod
+    def from_database_and_schema_config(
+        cls, config: "DatabaseAndSchemaConfig", table_fqn: str
+    ):
+        table_config = TableConfig(
+            fullyQualifiedName=table_fqn,
+            profileSample=config.profileSample,
+            profileSampleType=config.profileSampleType,
+            sampleDataCount=config.sampleDataCount,
+            samplingMethodType=config.samplingMethodType,
+        )
+        return table_config
+
+
+class DatabaseAndSchemaConfig(BaseProfileConfig):
+    """schema profile config"""
+
+    sampleDataStorageConfig: Optional[SampleDataStorageConfig] = None
+
+
+class SampleData(BaseModel):
+    """TableData wrapper to handle ephemeral SampleData"""
+
+    data: Annotated[TableData, Field(None, description="Table Sample Data")]
+    store: Annotated[
+        bool, Field(False, description="Is the sample data should be stored or not")
+    ]
+
+
+class SamplerResponse(ConfigModel):
+    """PII & Sampler Workflow Response. For a given table, return all the tags and sample data"""
+
+    table: Table
+    sample_data: Optional[SampleData] = None
+    column_tags: Optional[List[ColumnTag]] = None
+
+    def __str__(self):
+        """Return the table name being processed"""
+        return f"Table [{self.table.name.root}]"
+
+
+class SampleConfig(ConfigModel):
+    """Profile Sample Config"""
+
+    profile_sample: Optional[Union[float, int]] = None
+    profile_sample_type: Optional[ProfileSampleType] = ProfileSampleType.PERCENTAGE
+    sampling_method_type: Optional[SamplingMethodType] = None

metadata/{profiler/processor/sampler → sampler}/nosql/sampler.py

@@ -1,22 +1,46 @@
+#  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.
+"""NoSQL Sampler"""
 from typing import Dict, List, Optional, Tuple
 
 from metadata.generated.schema.entity.data.table import ProfileSampleType, TableData
+from metadata.profiler.adaptors.factory import factory
 from metadata.profiler.adaptors.nosql_adaptor import NoSQLAdaptor
-from metadata.profiler.processor.sampler.sampler_interface import SamplerInterface
+from metadata.sampler.sampler_interface import SamplerInterface
 from metadata.utils.constants import SAMPLE_DATA_DEFAULT_COUNT
 from metadata.utils.sqa_like_column import SQALikeColumn
 
 
 class NoSQLSampler(SamplerInterface):
+    """NoSQL generic implementation for the sampler"""
+
     client: NoSQLAdaptor
 
+    @property
+    def raw_dataset(self):
+        return self.entity
+
+    def get_client(self):
+        return factory.create(
+            self.service_connection_config.__class__.__name__,
+            client=self.connection,
+        )
+
     def _rdn_sample_from_user_query(self) -> List[Dict[str, any]]:
         """
         Get random sample from user query
         """
         limit = self._get_limit()
         return self.client.query(
-            self.table, self.table.columns, self._profile_sample_query, limit
+            self.raw_dataset, self.raw_dataset.columns, self.sample_query, limit
         )
 
     def _fetch_sample_data_from_user_query(self) -> TableData:
@@ -27,18 +51,18 @@ class NoSQLSampler(SamplerInterface):
         records = self._rdn_sample_from_user_query()
         columns = [
             SQALikeColumn(name=column.name.root, type=column.dataType)
-            for column in self.table.columns
+            for column in self.raw_dataset.columns
         ]
         rows, cols = self.transpose_records(records, columns)
         return TableData(
             rows=[list(map(str, row)) for row in rows], columns=[c.name for c in cols]
         )
 
-    def random_sample(self):
-        pass
+    def get_dataset(self, **__):
+        """No randomization for NoSQL"""
 
     def fetch_sample_data(self, columns: List[SQALikeColumn]) -> TableData:
-        if self._profile_sample_query:
+        if self.sample_query:
             return self._fetch_sample_data_from_user_query()
         return self._fetch_sample_data(columns)
 
@@ -47,7 +71,9 @@ class NoSQLSampler(SamplerInterface):
         returns sampled ometa dataframes
         """
         limit = self._get_limit()
-        records = self.client.scan(self.table, self.table.columns, limit)
+        records = self.client.scan(
+            self.raw_dataset, self.raw_dataset.columns, int(limit)
+        )
         rows, cols = self.transpose_records(records, columns)
         return TableData(
             rows=[list(map(str, row)) for row in rows],
@@ -55,11 +81,11 @@ class NoSQLSampler(SamplerInterface):
         )
 
     def _get_limit(self) -> Optional[int]:
-        num_rows = self.client.item_count(self.table)
-        if self.profile_sample_type == ProfileSampleType.PERCENTAGE:
-            limit = num_rows * (self.profile_sample / 100)
-        elif self.profile_sample_type == ProfileSampleType.ROWS:
-            limit = self.profile_sample
+        num_rows = self.client.item_count(self.raw_dataset)
+        if self.sample_config.profile_sample_type == ProfileSampleType.PERCENTAGE:
+            limit = num_rows * (self.sample_config.profile_sample / 100)
+        elif self.sample_config.profile_sample_type == ProfileSampleType.ROWS:
+            limit = self.sample_config.profile_sample
         else:
             limit = SAMPLE_DATA_DEFAULT_COUNT
         return limit
@@ -75,3 +101,9 @@ class NoSQLSampler(SamplerInterface):
                 row.append(record.get(column.name))
             rows.append(row)
         return rows, columns
+
+    def get_columns(self) -> List[Optional[SQALikeColumn]]:
+        return [
+            SQALikeColumn(name=c.name.root, type=c.dataType)
+            for c in self.raw_dataset.columns
+        ]