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

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

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
+        ]