openmetadata-ingestion 1.5.13.0__py3-none-any.whl → 1.6.0.0rc2__py3-none-any.whl

metadata/parsers/avro_parser.py

@@ -30,15 +30,19 @@ RECORD_DATATYPE_NAME = "RECORD"
 
 
 def _parse_array_children(
-    arr_item: Schema, cls: Type[BaseModel] = FieldModel
+    arr_item: Schema,
+    cls: Type[BaseModel] = FieldModel,
+    already_parsed: Optional[dict] = None,
 ) -> Tuple[str, Optional[Union[FieldModel, Column]]]:
     if isinstance(arr_item, ArraySchema):
-        display_type, children = _parse_array_children(arr_item.items, cls=cls)
+        display_type, children = _parse_array_children(
+            arr_item.items, cls=cls, already_parsed=already_parsed
+        )
         return f"ARRAY<{display_type}>", children
 
     if isinstance(arr_item, UnionSchema):
         display_type, children = _parse_union_children(
-            parent=None, union_field=arr_item, cls=cls
+            parent=None, union_field=arr_item, cls=cls, already_parsed=already_parsed
         )
         return f"UNION<{display_type}>", children
 
@@ -46,7 +50,7 @@ def _parse_array_children(
         child_obj = cls(
             name=arr_item.name,
             dataType=str(arr_item.type).upper(),
-            children=get_avro_fields(arr_item, cls),
+            children=get_avro_fields(arr_item, cls, already_parsed=already_parsed),
             description=arr_item.doc,
         )
         return str(arr_item.type), child_obj
@@ -55,7 +59,9 @@ def _parse_array_children(
 
 
 def parse_array_fields(
-    field: ArraySchema, cls: Type[BaseModel] = FieldModel
+    field: ArraySchema,
+    cls: Type[BaseModel] = FieldModel,
+    already_parsed: Optional[dict] = None,
 ) -> Optional[List[Union[FieldModel, Column]]]:
     """
     Parse array field for avro schema
@@ -93,7 +99,9 @@ def parse_array_fields(
         description=field.doc,
     )
 
-    display, children = _parse_array_children(arr_item=field.type.items, cls=cls)
+    display, children = _parse_array_children(
+        arr_item=field.type.items, cls=cls, already_parsed=already_parsed
+    )
 
     obj.dataTypeDisplay = f"ARRAY<{display}>"
     if cls == Column:
@@ -109,6 +117,7 @@ def _parse_union_children(
     parent: Optional[Schema],
     union_field: UnionSchema,
     cls: Type[BaseModel] = FieldModel,
+    already_parsed: Optional[dict] = None,
 ) -> Tuple[str, Optional[Union[FieldModel, Column]]]:
     non_null_schema = [
         (i, schema)
@@ -120,7 +129,9 @@ def _parse_union_children(
         field = non_null_schema[0][1]
 
         if isinstance(field, ArraySchema):
-            display, children = _parse_array_children(arr_item=field.items, cls=cls)
+            display, children = _parse_array_children(
+                arr_item=field.items, cls=cls, already_parsed=already_parsed
+            )
             sub_type = [None, None]
             sub_type[non_null_schema[0][0]] = f"ARRAY<{display}>"
             sub_type[non_null_schema[0][0] ^ 1] = "null"
@@ -131,7 +142,9 @@ def _parse_union_children(
             children = cls(
                 name=field.name,
                 dataType=str(field.type).upper(),
-                children=None if field == parent else get_avro_fields(field, cls),
+                children=None
+                if field == parent
+                else get_avro_fields(field, cls, already_parsed),
                 description=field.doc,
             )
             return sub_type, children
@@ -139,7 +152,11 @@ def _parse_union_children(
     return sub_type, None
 
 
-def parse_record_fields(field: RecordSchema, cls: Type[BaseModel] = FieldModel):
+def parse_record_fields(
+    field: RecordSchema,
+    cls: Type[BaseModel] = FieldModel,
+    already_parsed: Optional[dict] = None,
+):
     """
     Parse the nested record fields for avro
     """
@@ -150,7 +167,7 @@ def parse_record_fields(field: RecordSchema, cls: Type[BaseModel] = FieldModel):
             cls(
                 name=field.type.name,
                 dataType=RECORD_DATATYPE_NAME,
-                children=get_avro_fields(field.type, cls),
+                children=get_avro_fields(field.type, cls, already_parsed),
                 description=field.type.doc,
             )
         ],
@@ -163,6 +180,7 @@ def parse_union_fields(
     parent: Optional[Schema],
     union_field: Schema,
     cls: Type[BaseModel] = FieldModel,
+    already_parsed: Optional[dict] = None,
 ) -> Optional[List[Union[FieldModel, Column]]]:
     """
     Parse union field for avro schema
@@ -202,7 +220,7 @@ def parse_union_fields(
         description=union_field.doc,
     )
     sub_type, children = _parse_union_children(
-        union_field=field_type, cls=cls, parent=parent
+        union_field=field_type, cls=cls, parent=parent, already_parsed=already_parsed
     )
     obj.dataTypeDisplay = f"UNION<{sub_type}>"
     if children and cls == FieldModel:
@@ -237,7 +255,7 @@ def parse_avro_schema(
             cls(
                 name=parsed_schema.name,
                 dataType=str(parsed_schema.type).upper(),
-                children=get_avro_fields(parsed_schema, cls),
+                children=get_avro_fields(parsed_schema, cls, {}),
                 description=parsed_schema.doc,
             )
         ]
@@ -249,23 +267,40 @@ def parse_avro_schema(
 
 
 def get_avro_fields(
-    parsed_schema: Schema, cls: Type[BaseModel] = FieldModel
+    parsed_schema: Schema,
+    cls: Type[BaseModel] = FieldModel,
+    already_parsed: Optional[dict] = None,
 ) -> Optional[List[Union[FieldModel, Column]]]:
     """
     Recursively convert the parsed schema into required models
     """
     field_models = []
 
+    if parsed_schema.name in already_parsed:
+        if already_parsed[parsed_schema.name] == parsed_schema.type:
+            return None
+    else:
+        already_parsed.update({parsed_schema.name: parsed_schema.type})
+
     for field in parsed_schema.fields:
         try:
             if isinstance(field.type, ArraySchema):
-                field_models.append(parse_array_fields(field, cls=cls))
+                field_models.append(
+                    parse_array_fields(field, cls=cls, already_parsed=already_parsed)
+                )
             elif isinstance(field.type, UnionSchema):
                 field_models.append(
-                    parse_union_fields(union_field=field, cls=cls, parent=parsed_schema)
+                    parse_union_fields(
+                        union_field=field,
+                        cls=cls,
+                        parent=parsed_schema,
+                        already_parsed=already_parsed,
+                    )
                 )
             elif isinstance(field.type, RecordSchema):
-                field_models.append(parse_record_fields(field, cls=cls))
+                field_models.append(
+                    parse_record_fields(field, cls=cls, already_parsed=already_parsed)
+                )
             else:
                 field_models.append(parse_single_field(field, cls=cls))
         except Exception as exc:  # pylint: disable=broad-except

metadata/parsers/json_schema_parser.py

@@ -16,7 +16,7 @@ Utils module to parse the jsonschema
 import json
 import traceback
 from enum import Enum
-from typing import List, Optional, Type
+from typing import List, Optional, Tuple, Type
 
 from pydantic import BaseModel
 
@@ -66,6 +66,49 @@ def parse_json_schema(
     return None
 
 
+def get_child_models(key, value, field_models, cls: Type[BaseModel] = FieldModel):
+    """
+    Method to parse the child objects in the json schema
+    """
+    try:
+        cls_obj = cls(
+            name=key,
+            displayName=value.get("title"),
+            dataType=JsonSchemaDataTypes(value.get("type", "unknown")).name,
+            description=value.get("description"),
+        )
+        children = None
+        if value.get("type") == JsonSchemaDataTypes.RECORD.value:
+            children = get_json_schema_fields(value.get("properties"), cls=cls)
+        if value.get("type") == JsonSchemaDataTypes.ARRAY.value:
+            datatype_display, children = get_json_schema_array_fields(
+                value.get("items"), cls=cls
+            )
+            cls_obj.dataTypeDisplay = f"ARRAY<{datatype_display}>"
+        cls_obj.children = children
+        field_models.append(cls_obj)
+    except Exception as exc:  # pylint: disable=broad-except
+        logger.debug(traceback.format_exc())
+        logger.warning(f"Unable to parse the json schema into models: {exc}")
+
+
+def get_json_schema_array_fields(
+    array_items, cls: Type[BaseModel] = FieldModel
+) -> Optional[Tuple[str, List[FieldModel]]]:
+    """
+    Recursively convert the parsed array schema into required models
+    """
+    field_models = []
+    if array_items.get("type") == JsonSchemaDataTypes.RECORD.value:
+        for key, value in array_items.get("properties", {}).items():
+            get_child_models(key, value, field_models, cls)
+
+    return (
+        JsonSchemaDataTypes(array_items.get("type", "unknown")).name,
+        field_models or None,
+    )
+
+
 def get_json_schema_fields(
     properties, cls: Type[BaseModel] = FieldModel
 ) -> Optional[List[FieldModel]]:
@@ -74,20 +117,6 @@ def get_json_schema_fields(
     """
     field_models = []
     for key, value in properties.items():
-        try:
-            field_models.append(
-                cls(
-                    name=key,
-                    displayName=value.get("title"),
-                    dataType=JsonSchemaDataTypes(value.get("type", "unknown")).name,
-                    description=value.get("description"),
-                    children=get_json_schema_fields(value.get("properties"))
-                    if value.get("type") == "object"
-                    else None,
-                )
-            )
-        except Exception as exc:  # pylint: disable=broad-except
-            logger.debug(traceback.format_exc())
-            logger.warning(f"Unable to parse the json schema into models: {exc}")
+        get_child_models(key, value, field_models, cls)
 
     return field_models

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/pii/scanners/ner_scanner.py

@@ -19,7 +19,7 @@ import traceback
 from collections import defaultdict
 from typing import Any, Dict, List, Optional, Tuple, Union
 
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict
 
 from metadata.generated.schema.entity.classification.tag import Tag
 from metadata.pii.constants import PII, SPACY_EN_MODEL
@@ -46,6 +46,7 @@ class StringAnalysis(BaseModel):
 class NLPEngineModel(BaseModel):
     """Required to pass the nlp_engine as {"lang_code": "en", "model_name": "en_core_web_lg"}"""
 
+    model_config = ConfigDict(protected_namespaces=())
     lang_code: str
     model_name: str
 

metadata/profiler/api/models.py

@@ -15,85 +15,25 @@ Return types for Profiler workflow execution.
 We need to define this class as we end up having
 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,
-    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
-    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,
-        )
-        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
-
-
 class ProfilerProcessorConfig(ConfigModel):
     """
     Defines how we read the processor information
@@ -106,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.
@@ -125,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(