openmetadata-ingestion 1.5.11.0__py3-none-any.whl → 1.6.0.0rc1__py3-none-any.whl

metadata/ingestion/source/database/vertica/service_spec.py

@@ -0,0 +1,10 @@
+from metadata.ingestion.source.database.vertica.lineage import VerticaLineageSource
+from metadata.ingestion.source.database.vertica.metadata import VerticaSource
+from metadata.ingestion.source.database.vertica.usage import VerticaUsageSource
+from metadata.utils.service_spec.default import DefaultDatabaseSpec
+
+ServiceSpec = DefaultDatabaseSpec(
+    metadata_source_class=VerticaSource,
+    lineage_source_class=VerticaLineageSource,
+    usage_source_class=VerticaUsageSource,
+)

metadata/ingestion/source/messaging/kafka/service_spec.py

@@ -0,0 +1,4 @@
+from metadata.ingestion.source.messaging.kafka.metadata import KafkaSource
+from metadata.utils.service_spec import BaseSpec
+
+ServiceSpec = BaseSpec(metadata_source_class=KafkaSource)

metadata/ingestion/source/messaging/kinesis/service_spec.py

@@ -0,0 +1,4 @@
+from metadata.ingestion.source.messaging.kinesis.metadata import KinesisSource
+from metadata.utils.service_spec import BaseSpec
+
+ServiceSpec = BaseSpec(metadata_source_class=KinesisSource)

metadata/ingestion/source/messaging/redpanda/service_spec.py

@@ -0,0 +1,4 @@
+from metadata.ingestion.source.messaging.redpanda.metadata import RedpandaSource
+from metadata.utils.service_spec import BaseSpec
+
+ServiceSpec = BaseSpec(metadata_source_class=RedpandaSource)

metadata/ingestion/source/metadata/alationsink/service_spec.py

@@ -0,0 +1,4 @@
+from metadata.ingestion.source.metadata.alationsink.metadata import AlationsinkSource
+from metadata.utils.service_spec import BaseSpec
+
+ServiceSpec = BaseSpec(metadata_source_class=AlationsinkSource)

metadata/ingestion/source/metadata/amundsen/service_spec.py

@@ -0,0 +1,4 @@
+from metadata.ingestion.source.metadata.amundsen.metadata import AmundsenSource
+from metadata.utils.service_spec import BaseSpec
+
+ServiceSpec = BaseSpec(metadata_source_class=AmundsenSource)

metadata/ingestion/source/metadata/atlas/service_spec.py

@@ -0,0 +1,4 @@
+from metadata.ingestion.source.metadata.atlas.metadata import AtlasSource
+from metadata.utils.service_spec import BaseSpec
+
+ServiceSpec = BaseSpec(metadata_source_class=AtlasSource)

metadata/ingestion/source/mlmodel/mlflow/service_spec.py

@@ -0,0 +1,4 @@
+from metadata.ingestion.source.mlmodel.mlflow.metadata import MlflowSource
+from metadata.utils.service_spec import BaseSpec
+
+ServiceSpec = BaseSpec(metadata_source_class=MlflowSource)

metadata/ingestion/source/mlmodel/sagemaker/service_spec.py

@@ -0,0 +1,4 @@
+from metadata.ingestion.source.mlmodel.sagemaker.metadata import SagemakerSource
+from metadata.utils.service_spec import BaseSpec
+
+ServiceSpec = BaseSpec(metadata_source_class=SagemakerSource)

metadata/ingestion/source/pipeline/airbyte/service_spec.py

@@ -0,0 +1,4 @@
+from metadata.ingestion.source.pipeline.airbyte.metadata import AirbyteSource
+from metadata.utils.service_spec import BaseSpec
+
+ServiceSpec = BaseSpec(metadata_source_class=AirbyteSource)

metadata/ingestion/source/pipeline/airflow/metadata.py

@@ -17,7 +17,7 @@ from datetime import datetime
 from enum import Enum
 from typing import Dict, Iterable, List, Optional, cast
 
-from airflow.models import BaseOperator, DagRun, TaskInstance
+from airflow.models import BaseOperator, DagRun, DagTag, TaskInstance
 from airflow.models.dag import DagModel
 from airflow.models.serialized_dag import SerializedDagModel
 from airflow.serialization.serialized_objects import SerializedDAG
@@ -57,6 +57,7 @@ from metadata.generated.schema.type.entityReferenceList import EntityReferenceLi
 from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import InvalidSourceException
 from metadata.ingestion.connections.session import create_and_bind_session
+from metadata.ingestion.models.ometa_classification import OMetaTagAndClassification
 from metadata.ingestion.models.pipeline_status import OMetaPipelineStatus
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.pipeline.airflow.lineage_parser import (
@@ -74,9 +75,12 @@ from metadata.utils import fqn
 from metadata.utils.constants import ENTITY_REFERENCE_TYPE_MAP
 from metadata.utils.helpers import clean_uri, datetime_to_ts
 from metadata.utils.logger import ingestion_logger
+from metadata.utils.tag_utils import get_ometa_tag_and_classification, get_tag_labels
 
 logger = ingestion_logger()
 
+AIRFLOW_TAG_CATEGORY = "AirflowTags"
+
 
 class AirflowTaskStatus(Enum):
     SUCCESS = "success"
@@ -152,6 +156,31 @@ class AirflowSource(PipelineServiceSource):
             return task["__var"]
         return task
 
+    def get_all_tags(self, dag_id: str) -> List[str]:
+        try:
+            tag_query = (
+                self.session.query(DagTag.name)
+                .filter(DagTag.dag_id == dag_id)
+                .distinct()
+                .all()
+            )
+            return [tag[0] for tag in tag_query]
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.warning(f"Could not extract tags details due to {exc}")
+        return []
+
+    def yield_tag(
+        self, pipeline_details: AirflowDagDetails
+    ) -> Iterable[Either[OMetaTagAndClassification]]:
+        yield from get_ometa_tag_and_classification(
+            tags=self.get_all_tags(dag_id=pipeline_details.dag_id),
+            classification_name=AIRFLOW_TAG_CATEGORY,
+            tag_description="Airflow Tag",
+            classification_description="Tags associated with airflow entities.",
+            include_tags=self.source_config.includeTags,
+        )
+
     def get_pipeline_status(self, dag_id: str) -> List[DagRun]:
         """
         Return the DagRuns of given dag
@@ -441,7 +470,6 @@ class AirflowSource(PipelineServiceSource):
         try:
             # Airflow uses /dags/dag_id/grid to show pipeline / dag
             source_url = f"{clean_uri(self.service_connection.hostPort)}/dags/{pipeline_details.dag_id}/grid"
-
             pipeline_request = CreatePipelineRequest(
                 name=EntityName(pipeline_details.dag_id),
                 description=Markdown(pipeline_details.description)
@@ -459,6 +487,12 @@ class AirflowSource(PipelineServiceSource):
                 service=FullyQualifiedEntityName(self.context.get().pipeline_service),
                 owners=self.get_owner(pipeline_details.owner),
                 scheduleInterval=pipeline_details.schedule_interval,
+                tags=get_tag_labels(
+                    metadata=self.metadata,
+                    tags=pipeline_details.data.dag.tags,
+                    classification_name=AIRFLOW_TAG_CATEGORY,
+                    include_tags=self.source_config.includeTags,
+                ),
             )
             yield Either(right=pipeline_request)
             self.register_record(pipeline_request=pipeline_request)

metadata/ingestion/source/pipeline/airflow/service_spec.py

@@ -0,0 +1,4 @@
+from metadata.ingestion.source.pipeline.airflow.metadata import AirflowSource
+from metadata.utils.service_spec import BaseSpec
+
+ServiceSpec = BaseSpec(metadata_source_class=AirflowSource)

metadata/ingestion/source/pipeline/dagster/metadata.py

@@ -161,7 +161,9 @@ class DagsterSource(PipelineServiceSource):
                 )
             )
 
-    def yield_tag(self, *_, **__) -> Iterable[Either[OMetaTagAndClassification]]:
+    def yield_tag(
+        self, pipeline_details: DagsterPipeline
+    ) -> Iterable[Either[OMetaTagAndClassification]]:
         yield from get_ometa_tag_and_classification(
             tags=[self.context.get().repository_name],
             classification_name=DAGSTER_TAG_CATEGORY,

metadata/ingestion/source/pipeline/dagster/service_spec.py

@@ -0,0 +1,4 @@
+from metadata.ingestion.source.pipeline.dagster.metadata import DagsterSource
+from metadata.utils.service_spec import BaseSpec
+
+ServiceSpec = BaseSpec(metadata_source_class=DagsterSource)

metadata/ingestion/source/pipeline/databrickspipeline/service_spec.py

@@ -0,0 +1,6 @@
+from metadata.ingestion.source.pipeline.databrickspipeline.metadata import (
+    DatabrickspipelineSource,
+)
+from metadata.utils.service_spec import BaseSpec
+
+ServiceSpec = BaseSpec(metadata_source_class=DatabrickspipelineSource)

metadata/ingestion/source/pipeline/dbtcloud/service_spec.py

@@ -0,0 +1,4 @@
+from metadata.ingestion.source.pipeline.dbtcloud.metadata import DbtcloudSource
+from metadata.utils.service_spec import BaseSpec
+
+ServiceSpec = BaseSpec(metadata_source_class=DbtcloudSource)

metadata/ingestion/source/pipeline/domopipeline/service_spec.py

@@ -0,0 +1,4 @@
+from metadata.ingestion.source.pipeline.domopipeline.metadata import DomopipelineSource
+from metadata.utils.service_spec import BaseSpec
+
+ServiceSpec = BaseSpec(metadata_source_class=DomopipelineSource)

metadata/ingestion/source/pipeline/fivetran/service_spec.py

@@ -0,0 +1,4 @@
+from metadata.ingestion.source.pipeline.fivetran.metadata import FivetranSource
+from metadata.utils.service_spec import BaseSpec
+
+ServiceSpec = BaseSpec(metadata_source_class=FivetranSource)

metadata/ingestion/source/pipeline/flink/service_spec.py

@@ -0,0 +1,4 @@
+from metadata.ingestion.source.pipeline.flink.metadata import FlinkSource
+from metadata.utils.service_spec import BaseSpec
+
+ServiceSpec = BaseSpec(metadata_source_class=FlinkSource)

metadata/ingestion/source/pipeline/gluepipeline/service_spec.py

@@ -0,0 +1,4 @@
+from metadata.ingestion.source.pipeline.gluepipeline.metadata import GluepipelineSource
+from metadata.utils.service_spec import BaseSpec
+
+ServiceSpec = BaseSpec(metadata_source_class=GluepipelineSource)

metadata/ingestion/source/pipeline/kafkaconnect/service_spec.py

@@ -0,0 +1,4 @@
+from metadata.ingestion.source.pipeline.kafkaconnect.metadata import KafkaconnectSource
+from metadata.utils.service_spec import BaseSpec
+
+ServiceSpec = BaseSpec(metadata_source_class=KafkaconnectSource)

metadata/ingestion/source/pipeline/nifi/service_spec.py

@@ -0,0 +1,4 @@
+from metadata.ingestion.source.pipeline.nifi.metadata import NifiSource
+from metadata.utils.service_spec import BaseSpec
+
+ServiceSpec = BaseSpec(metadata_source_class=NifiSource)

metadata/ingestion/source/pipeline/openlineage/service_spec.py

@@ -0,0 +1,4 @@
+from metadata.ingestion.source.pipeline.openlineage.metadata import OpenlineageSource
+from metadata.utils.service_spec import BaseSpec
+
+ServiceSpec = BaseSpec(metadata_source_class=OpenlineageSource)

metadata/ingestion/source/pipeline/pipeline_service.py

@@ -222,7 +222,9 @@ class PipelineServiceSource(TopologyRunnerMixin, Source, ABC):
             f"Table FQN not found for table: {table_details} within services: {services}"
         )
 
-    def yield_tag(self, *args, **kwargs) -> Iterable[Either[OMetaTagAndClassification]]:
+    def yield_tag(
+        self, pipeline_details: Any
+    ) -> Iterable[Either[OMetaTagAndClassification]]:
         """Method to fetch pipeline tags"""
 
     def close(self):

metadata/ingestion/source/pipeline/spline/service_spec.py

@@ -0,0 +1,4 @@
+from metadata.ingestion.source.pipeline.spline.metadata import SplineSource
+from metadata.utils.service_spec import BaseSpec
+
+ServiceSpec = BaseSpec(metadata_source_class=SplineSource)

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

@@ -0,0 +1,4 @@
+from metadata.ingestion.source.search.elasticsearch.metadata import ElasticsearchSource
+from metadata.utils.service_spec import BaseSpec
+
+ServiceSpec = BaseSpec(metadata_source_class=ElasticsearchSource)

metadata/ingestion/source/storage/gcs/metadata.py

@@ -87,7 +87,7 @@ class GcsSource(StorageServiceSource):
     def create(
         cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
     ):
-        config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
+        config: WorkflowSource = WorkflowSource.model_validate(config_dict)
         connection: GcsConnection = config.serviceConnection.root.config
         if not isinstance(connection, GcsConnection):
             raise InvalidSourceException(
@@ -451,7 +451,7 @@ class GcsSource(StorageServiceSource):
                 verbose=False,
             )
             content = json.loads(response_object)
-            metadata_config = StorageContainerConfig.parse_obj(content)
+            metadata_config = StorageContainerConfig.model_validate(content)
             return metadata_config
         except ReadException:
             logger.warning(

metadata/ingestion/source/storage/gcs/service_spec.py

@@ -0,0 +1,4 @@
+from metadata.ingestion.source.storage.gcs.metadata import GcsSource
+from metadata.utils.service_spec import BaseSpec
+
+ServiceSpec = BaseSpec(metadata_source_class=GcsSource)

metadata/ingestion/source/storage/s3/service_spec.py

@@ -0,0 +1,4 @@
+from metadata.ingestion.source.storage.s3.metadata import S3Source
+from metadata.utils.service_spec import BaseSpec
+
+ServiceSpec = BaseSpec(metadata_source_class=S3Source)

metadata/mixins/sqalchemy/sqa_mixin.py

@@ -34,9 +34,10 @@ 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
 
 
-class SQAInterfaceMixin:
+class SQAInterfaceMixin(Root):
     """SQLAlchemy inteface mixin grouping shared methods between sequential and threaded executor"""
 
     def _get_engine(self):

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/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,6 +15,7 @@ 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
@@ -30,6 +31,7 @@ from metadata.generated.schema.entity.data.table import (
     ColumnProfilerConfig,
     PartitionProfilerConfig,
     ProfileSampleType,
+    SamplingMethodType,
     Table,
     TableData,
 )
@@ -58,6 +60,7 @@ class BaseProfileConfig(ConfigModel):
     fullyQualifiedName: FullyQualifiedEntityName
     profileSample: Optional[Union[float, int]] = None
     profileSampleType: Optional[ProfileSampleType] = None
+    samplingMethodType: Optional[SamplingMethodType] = None
     sampleDataCount: Optional[int] = 100
 
 
@@ -77,6 +80,7 @@ class TableConfig(BaseProfileConfig):
             profileSample=config.profileSample,
             profileSampleType=config.profileSampleType,
             sampleDataCount=config.sampleDataCount,
+            samplingMethodType=config.samplingMethodType,
         )
         return table_config
 
@@ -92,6 +96,7 @@ class ProfileSampleConfig(ConfigModel):
 
     profile_sample: Optional[Union[float, int]] = None
     profile_sample_type: Optional[ProfileSampleType] = ProfileSampleType.PERCENTAGE
+    sampling_method_type: Optional[SamplingMethodType] = None
 
 
 class ProfilerProcessorConfig(ConfigModel):