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

metadata/ingestion/source/search/elasticsearch/metadata.py

@@ -12,6 +12,7 @@
 Elasticsearch source to extract metadata
 """
 import shutil
+import traceback
 from pathlib import Path
 from typing import Any, Iterable, Optional
 
@@ -21,6 +22,7 @@ from metadata.generated.schema.api.data.createSearchIndex import (
     CreateSearchIndexRequest,
 )
 from metadata.generated.schema.entity.data.searchIndex import (
+    IndexType,
     SearchIndex,
     SearchIndexSampleData,
 )
@@ -103,6 +105,7 @@ class ElasticsearchSource(SearchServiceSource):
                 fields=parse_es_index_mapping(
                     search_index_details.get(index_name, {}).get("mappings")
                 ),
+                indexType=IndexType.Index,
             )
             yield Either(right=search_index_request)
             self.register_record(search_index_request=search_index_request)
@@ -143,6 +146,56 @@ class ElasticsearchSource(SearchServiceSource):
                 )
             )
 
+    def get_search_index_template_list(self) -> Iterable[dict]:
+        """
+        Get List of all search index template
+        """
+        yield from self.client.indices.get_index_template().get("index_templates", [])
+
+    def get_search_index_template_name(
+        self, search_index_template_details: dict
+    ) -> Optional[str]:
+        """
+        Get Search Index Template Name
+        """
+        return search_index_template_details and search_index_template_details["name"]
+
+    def yield_search_index_template(
+        self, search_index_template_details: Any
+    ) -> Iterable[Either[CreateSearchIndexRequest]]:
+        """
+        Method to Get Search Index Template Entity
+        """
+        try:
+            if self.source_config.includeIndexTemplate:
+                index_name = self.get_search_index_template_name(
+                    search_index_template_details
+                )
+                index_template = search_index_template_details["index_template"]
+                if index_name:
+                    search_index_template_request = CreateSearchIndexRequest(
+                        name=EntityName(index_name),
+                        displayName=index_name,
+                        searchIndexSettings=index_template.get("template", {}).get(
+                            "settings", {}
+                        ),
+                        service=FullyQualifiedEntityName(
+                            self.context.get().search_service
+                        ),
+                        fields=parse_es_index_mapping(
+                            index_template.get("template", {}).get("mappings")
+                        ),
+                        indexType=IndexType.IndexTemplate,
+                        description=index_template.get("_meta", {}).get("description"),
+                    )
+                    yield Either(right=search_index_template_request)
+                    self.register_record(
+                        search_index_request=search_index_template_request
+                    )
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.error(f"Could not include index templates due to {exc}")
+
     def close(self):
         try:
             if Path(self.service_connection.sslConfig.certificates.stagingDir).exists():

metadata/ingestion/source/search/search_service.py

@@ -83,7 +83,7 @@ class SearchServiceTopology(ServiceTopology):
                 cache_entities=True,
             ),
         ],
-        children=["search_index"],
+        children=["search_index", "search_index_template"],
         post_process=["mark_search_indexes_as_deleted"],
     )
     search_index: Annotated[
@@ -107,6 +107,21 @@ class SearchServiceTopology(ServiceTopology):
         ],
     )
 
+    search_index_template: Annotated[
+        TopologyNode, Field(description="Search Index Template Processing Node")
+    ] = TopologyNode(
+        producer="get_search_index_template",
+        stages=[
+            NodeStage(
+                type_=SearchIndex,
+                context="search_index_template",
+                processor="yield_search_index_template",
+                consumer=["search_service"],
+                use_cache=True,
+            )
+        ],
+    )
+
 
 class SearchServiceSource(TopologyRunnerMixin, Source, ABC):
     """
@@ -178,6 +193,34 @@ class SearchServiceSource(TopologyRunnerMixin, Source, ABC):
                 continue
             yield index_details
 
+    def yield_search_index_template(
+        self, search_index_template_details: Any
+    ) -> Iterable[Either[CreateSearchIndexRequest]]:
+        """Method to Get Search Index Templates"""
+
+    def get_search_index_template_list(self) -> Optional[List[Any]]:
+        """Get list of all search index templates"""
+
+    def get_search_index_template_name(self, search_index_template_details: Any) -> str:
+        """Get Search Index Template Name"""
+
+    def get_search_index_template(self) -> Any:
+        if self.source_config.includeIndexTemplate:
+            for index_template_details in self.get_search_index_template_list():
+                if search_index_template_name := self.get_search_index_template_name(
+                    index_template_details
+                ):
+                    if filter_by_search_index(
+                        self.source_config.searchIndexFilterPattern,
+                        search_index_template_name,
+                    ):
+                        self.status.filter(
+                            search_index_template_name,
+                            "Search Index Template Filtered Out",
+                        )
+                        continue
+                    yield index_template_details
+
     def yield_create_request_search_service(
         self, config: WorkflowSource
     ) -> Iterable[Either[CreateSearchServiceRequest]]:

metadata/mixins/pandas/pandas_mixin.py

@@ -13,7 +13,6 @@
 Interfaces with database for all database engine
 supporting sqlalchemy abstraction layer
 """
-import math
 import random
 from typing import cast
 
@@ -23,7 +22,6 @@ from metadata.data_quality.validations.table.pandas.tableRowInsertedCountToBeBet
 from metadata.generated.schema.entity.data.table import (
     PartitionIntervalTypes,
     PartitionProfilerConfig,
-    ProfileSampleType,
 )
 from metadata.readers.dataframe.models import DatalakeTableSchemaWrapper
 from metadata.utils.datalake.datalake_utils import fetch_dataframe
@@ -81,9 +79,7 @@ class PandasInterfaceMixin:
             for df in dfs
         ]
 
-    def return_ometa_dataframes_sampled(
-        self, service_connection_config, client, table, profile_sample_config
-    ):
+    def get_dataframes(self, service_connection_config, client, table):
         """
         returns sampled ometa dataframes
         """
@@ -94,35 +90,10 @@ class PandasInterfaceMixin:
                 key=table.name.root,
                 bucket_name=table.databaseSchema.name,
                 file_extension=table.fileFormat,
+                separator=None,
             ),
         )
         if data:
             random.shuffle(data)
-            # sampling data based on profiler config (if any)
-            if hasattr(profile_sample_config, "profile_sample"):
-                if (
-                    profile_sample_config.profile_sample_type
-                    == ProfileSampleType.PERCENTAGE
-                ):
-                    return [
-                        df.sample(
-                            frac=profile_sample_config.profile_sample / 100,
-                            random_state=random.randint(0, 100),
-                            replace=True,
-                        )
-                        for df in data
-                    ]
-                if profile_sample_config.profile_sample_type == ProfileSampleType.ROWS:
-                    sample_rows_per_chunk: int = math.floor(
-                        profile_sample_config.profile_sample / len(data)
-                    )
-                    return [
-                        df.sample(
-                            n=sample_rows_per_chunk,
-                            random_state=random.randint(0, 100),
-                            replace=True,
-                        )
-                        for df in data
-                    ]
             return data
         raise TypeError(f"Couldn't fetch {table.name.root}")

metadata/mixins/sqalchemy/sqa_mixin.py

@@ -20,6 +20,7 @@ from typing import List, Optional
 from sqlalchemy import Column, MetaData, inspect
 from sqlalchemy.orm import DeclarativeMeta
 
+from metadata.generated.schema.entity.data.table import Table
 from metadata.generated.schema.entity.services.connections.database.databricksConnection import (
     DatabricksConnection,
 )
@@ -29,16 +30,19 @@ from metadata.generated.schema.entity.services.connections.database.snowflakeCon
 from metadata.generated.schema.entity.services.connections.database.unityCatalogConnection import (
     UnityCatalogConnection,
 )
+from metadata.ingestion.models.custom_pydantic import BaseModel
+from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.connections import get_connection
 from metadata.ingestion.source.database.snowflake.queries import (
     SNOWFLAKE_SESSION_TAG_QUERY,
 )
 from metadata.profiler.orm.converter.base import ometa_to_sqa_orm
 from metadata.utils.collaborative_super import Root
+from metadata.utils.constants import NON_SQA_DATABASE_CONNECTIONS
 
 
 class SQAInterfaceMixin(Root):
-    """SQLAlchemy inteface mixin grouping shared methods between sequential and threaded executor"""
+    """SQLAlchemy interface mixin grouping shared methods between sequential and threaded executor"""
 
     def _get_engine(self):
         """Get engine for database
@@ -48,26 +52,13 @@ class SQAInterfaceMixin(Root):
         Returns:
             sqlalchemy engine
         """
-        engine = get_connection(self.service_connection_config)
+        engine = get_connection(super().service_connection_config)
 
         return engine
 
-    def _convert_table_to_orm_object(
-        self,
-        sqa_metadata_obj: Optional[MetaData] = None,
-    ) -> DeclarativeMeta:
-        """Given a table entity return a SQA ORM object
-
-        Args:
-            sqa_metadata_obj: sqa metadata registry
-        Returns:
-            DeclarativeMeta
-        """
-        return ometa_to_sqa_orm(self.table_entity, self.ometa_client, sqa_metadata_obj)
-
     def get_columns(self) -> Column:
         """get columns from an orm object"""
-        return inspect(self.table).c
+        return inspect(super().table).c
 
     def set_session_tag(self, session) -> None:
         """
@@ -115,3 +106,12 @@ class SQAInterfaceMixin(Root):
             for column in self.table.__table__.columns
             if column.name in {col.name.root for col in self.table_entity.columns}
         ]
+
+    def build_table_orm(
+        self, table: Table, service_conn_config: BaseModel, ometa_client: OpenMetadata
+    ) -> Optional[DeclarativeMeta]:
+        """Build the ORM table if needed for the sampler and profiler interfaces"""
+        if service_conn_config.type.value not in NON_SQA_DATABASE_CONNECTIONS:
+            orm_obj = ometa_to_sqa_orm(table, ometa_client, MetaData())
+            return orm_obj
+        return None

metadata/pii/processor.py

@@ -19,8 +19,8 @@ from metadata.generated.schema.entity.data.table import Column, TableData
 from metadata.generated.schema.entity.services.ingestionPipelines.status import (
     StackTraceError,
 )
-from metadata.generated.schema.metadataIngestion.databaseServiceProfilerPipeline import (
-    DatabaseServiceProfilerPipeline,
+from metadata.generated.schema.metadataIngestion.databaseServiceAutoClassificationPipeline import (
+    DatabaseServiceAutoClassificationPipeline,
 )
 from metadata.generated.schema.metadataIngestion.workflow import (
     OpenMetadataWorkflowConfig,
@@ -40,7 +40,7 @@ from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.pii.constants import PII
 from metadata.pii.scanners.column_name_scanner import ColumnNameScanner
 from metadata.pii.scanners.ner_scanner import NERScanner
-from metadata.profiler.api.models import ProfilerResponse
+from metadata.sampler.models import SamplerResponse
 from metadata.utils.logger import profiler_logger
 
 logger = profiler_logger()
@@ -61,8 +61,9 @@ class PIIProcessor(Processor):
         self.metadata = metadata
 
         # Init and type the source config
-        self.source_config: DatabaseServiceProfilerPipeline = cast(
-            DatabaseServiceProfilerPipeline, self.config.source.sourceConfig.config
+        self.source_config: DatabaseServiceAutoClassificationPipeline = cast(
+            DatabaseServiceAutoClassificationPipeline,
+            self.config.source.sourceConfig.config,
         )  # Used to satisfy type checked
 
         self._ner_scanner = None
@@ -71,7 +72,7 @@ class PIIProcessor(Processor):
 
     @property
     def name(self) -> str:
-        return "PII Processor"
+        return "Auto Classification Processor"
 
     @property
     def ner_scanner(self) -> NERScanner:
@@ -152,8 +153,8 @@ class PIIProcessor(Processor):
 
     def _run(
         self,
-        record: ProfilerResponse,
-    ) -> Either[ProfilerResponse]:
+        record: SamplerResponse,
+    ) -> Either[SamplerResponse]:
         """
         Main entrypoint for the scanner.
 
@@ -162,7 +163,7 @@ class PIIProcessor(Processor):
         """
 
         # We don't always need to process
-        if not self.source_config.processPiiSensitive:
+        if not self.source_config.enableAutoClassification:
             return Either(right=record)
 
         column_tags = []

metadata/profiler/api/models.py

@@ -18,87 +18,22 @@ multiple profilers per table and columns.
 
 from typing import List, Optional, Type, Union
 
-from pydantic import ConfigDict, Field
+from pydantic import ConfigDict
 from sqlalchemy import Column
 from sqlalchemy.orm import DeclarativeMeta
-from typing_extensions import Annotated
 
 from metadata.config.common import ConfigModel
 from metadata.generated.schema.api.data.createTableProfile import (
     CreateTableProfileRequest,
 )
-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.entity.data.table import Table
 from metadata.generated.schema.tests.customMetric import CustomMetric
-from metadata.generated.schema.type.basic import FullyQualifiedEntityName
-from metadata.ingestion.models.custom_pydantic import BaseModel
-from metadata.ingestion.models.table_metadata import ColumnTag
 from metadata.profiler.metrics.core import Metric, MetricTypes
 from metadata.profiler.processor.models import ProfilerDef
+from metadata.sampler.models import DatabaseAndSchemaConfig, TableConfig
 from metadata.utils.sqa_like_column import SQALikeColumn
 
 
-class ColumnConfig(ConfigModel):
-    """Column config for profiler"""
-
-    excludeColumns: Optional[List[str]] = None
-    includeColumns: Optional[List[ColumnProfilerConfig]] = None
-
-
-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 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 ProfileSampleConfig(ConfigModel):
-    """Profile Sample Config"""
-
-    profile_sample: Optional[Union[float, int]] = None
-    profile_sample_type: Optional[ProfileSampleType] = ProfileSampleType.PERCENTAGE
-    sampling_method_type: Optional[SamplingMethodType] = None
-
-
 class ProfilerProcessorConfig(ConfigModel):
     """
     Defines how we read the processor information
@@ -111,15 +46,6 @@ class ProfilerProcessorConfig(ConfigModel):
     databaseConfig: Optional[List[DatabaseAndSchemaConfig]] = []
 
 
-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 ProfilerResponse(ConfigModel):
     """
     ORM Profiler processor response.
@@ -130,8 +56,6 @@ class ProfilerResponse(ConfigModel):
 
     table: Table
     profile: CreateTableProfileRequest
-    sample_data: Optional[SampleData] = None
-    column_tags: Optional[List[ColumnTag]] = None
 
     def __str__(self):
         """Return the table name being processed"""

metadata/profiler/config.py

@@ -0,0 +1,39 @@
+#  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.
+"""
+Profiler configuration helpers
+"""
+from typing import Optional
+
+from metadata.generated.schema.entity.data.database import (
+    Database,
+    DatabaseProfilerConfig,
+)
+from metadata.generated.schema.entity.data.databaseSchema import (
+    DatabaseSchema,
+    DatabaseSchemaProfilerConfig,
+)
+
+
+def get_database_profiler_config(
+    database_entity: Optional[Database],
+) -> Optional[DatabaseProfilerConfig]:
+    if database_entity and database_entity.databaseProfilerConfig:
+        return database_entity.databaseProfilerConfig
+    return None
+
+
+def get_schema_profiler_config(
+    schema_entity: Optional[DatabaseSchema],
+) -> Optional[DatabaseSchemaProfilerConfig]:
+    if schema_entity and schema_entity.databaseSchemaProfilerConfig:
+        return schema_entity.databaseSchemaProfilerConfig
+    return None

metadata/profiler/interface/nosql/profiler_interface.py

@@ -30,7 +30,6 @@ from metadata.profiler.interface.profiler_interface import ProfilerInterface
 from metadata.profiler.metrics.core import Metric, MetricTypes
 from metadata.profiler.metrics.registry import Metrics
 from metadata.profiler.processor.metric_filter import MetricFilter
-from metadata.profiler.processor.sampler.nosql.sampler import NoSQLSampler
 from metadata.utils.logger import profiler_interface_registry_logger
 from metadata.utils.sqa_like_column import SQALikeColumn
 
@@ -45,10 +44,6 @@ class NoSQLProfilerInterface(ProfilerInterface):
 
     # pylint: disable=too-many-arguments
 
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.sampler = self._get_sampler()
-
     def _compute_table_metrics(
         self,
         metrics: List[Type[Metric]],
@@ -157,32 +152,12 @@ class NoSQLProfilerInterface(ProfilerInterface):
     def fetch_sample_data(self, table, columns: List[SQALikeColumn]) -> TableData:
         return self.sampler.fetch_sample_data(columns)
 
-    def _get_sampler(self) -> NoSQLSampler:
-        """Get NoSQL sampler from config"""
-        from metadata.profiler.processor.sampler.sampler_factory import (  # pylint: disable=import-outside-toplevel
-            sampler_factory_,
-        )
-
-        return sampler_factory_.create(
-            self.service_connection_config.__class__.__name__,
-            table=self.table,
-            client=factory.create(
-                self.service_connection_config.__class__.__name__,
-                client=self.connection,
-            ),
-            profile_sample_config=self.profile_sample_config,
-            partition_details=self.partition_details,
-            profile_sample_query=self.profile_query,
-        )
-
     def get_composed_metrics(
         self, column: Column, metric: Metrics, column_results: Dict
     ):
         return None
 
-    def get_hybrid_metrics(
-        self, column: Column, metric: Metrics, column_results: Dict, **kwargs
-    ):
+    def get_hybrid_metrics(self, column: Column, metric: Metrics, column_results: Dict):
         return None
 
     def get_all_metrics(