openmetadata-ingestion 1.4.6.0rc1__py3-none-any.whl → 1.5.0.0rc1__py3-none-any.whl

metadata/ingestion/source/storage/storage_service.py

@@ -14,6 +14,9 @@ Base class for ingesting Object Storage services
 from abc import ABC, abstractmethod
 from typing import Any, Iterable, List, Optional, Set
 
+from pydantic import Field
+from typing_extensions import Annotated
+
 from metadata.generated.schema.api.data.createContainer import CreateContainerRequest
 from metadata.generated.schema.entity.data.container import Container
 from metadata.generated.schema.entity.services.storageService import (
@@ -38,6 +41,7 @@ from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import Source
 from metadata.ingestion.api.topology_runner import TopologyRunnerMixin
 from metadata.ingestion.models.delete_entity import DeleteEntity
+from metadata.ingestion.models.ometa_classification import OMetaTagAndClassification
 from metadata.ingestion.models.topology import (
     NodeStage,
     ServiceTopology,
@@ -68,7 +72,14 @@ OPENMETADATA_TEMPLATE_FILE_NAME = "openmetadata.json"
 
 
 class StorageServiceTopology(ServiceTopology):
-    root = TopologyNode(
+    """
+    Defines the hierarchy in Messaging Services.
+    service -> container -> container -> container...
+    """
+
+    root: Annotated[
+        TopologyNode, Field(description="Root node for the topology")
+    ] = TopologyNode(
         producer="get_services",
         stages=[
             NodeStage(
@@ -84,9 +95,18 @@ class StorageServiceTopology(ServiceTopology):
         post_process=["mark_containers_as_deleted"],
     )
 
-    container = TopologyNode(
+    container: Annotated[
+        TopologyNode, Field(description="Container Processing Node")
+    ] = TopologyNode(
         producer="get_containers",
         stages=[
+            NodeStage(
+                type_=OMetaTagAndClassification,
+                context="tags",
+                processor="yield_tag_details",
+                nullable=True,
+                store_all_in_context=True,
+            ),
             NodeStage(
                 type_=Container,
                 context="container",
@@ -94,7 +114,7 @@ class StorageServiceTopology(ServiceTopology):
                 consumer=["objectstore_service"],
                 nullable=True,
                 use_cache=True,
-            )
+            ),
         ],
     )
 
@@ -109,7 +129,7 @@ class StorageServiceSource(TopologyRunnerMixin, Source, ABC):
     config: WorkflowSource
     metadata: OpenMetadata
     # Big union of types we want to fetch dynamically
-    service_connection: StorageConnection.__fields__["config"].type_
+    service_connection: StorageConnection.model_fields["config"].annotation
 
     topology = StorageServiceTopology()
     context = TopologyContextManager(topology)
@@ -125,7 +145,7 @@ class StorageServiceSource(TopologyRunnerMixin, Source, ABC):
         super().__init__()
         self.config = config
         self.metadata = metadata
-        self.service_connection = self.config.serviceConnection.__root__.config
+        self.service_connection = self.config.serviceConnection.root.config
         self.source_config: StorageServiceMetadataPipeline = (
             self.config.sourceConfig.config
         )
@@ -176,6 +196,22 @@ class StorageServiceSource(TopologyRunnerMixin, Source, ABC):
     def prepare(self):
         """By default, nothing needs to be taken care of when loading the source"""
 
+    def yield_container_tags(
+        self, container_details: Any
+    ) -> Iterable[Either[OMetaTagAndClassification]]:
+        """
+        From topology. To be run for each container
+        """
+
+    def yield_tag_details(
+        self, container_details: Any
+    ) -> Iterable[Either[OMetaTagAndClassification]]:
+        """
+        From topology. To be run for each container
+        """
+        if self.source_config.includeTags:
+            yield from self.yield_container_tags(container_details) or []
+
     def register_record(self, container_request: CreateContainerRequest) -> None:
         """
         Mark the container record as scanned and update
@@ -184,7 +220,7 @@ class StorageServiceSource(TopologyRunnerMixin, Source, ABC):
         parent_container = (
             self.metadata.get_by_id(
                 entity=Container, entity_id=container_request.parent.id
-            ).fullyQualifiedName.__root__
+            ).fullyQualifiedName.root
             if container_request.parent
             else None
         )
@@ -193,7 +229,7 @@ class StorageServiceSource(TopologyRunnerMixin, Source, ABC):
             entity_type=Container,
             service_name=self.context.get().objectstore_service,
             parent_container=parent_container,
-            container_name=container_request.name.__root__,
+            container_name=fqn.quote_name(container_request.name.root),
         )
 
         self.container_source_state.add(container_fqn)

metadata/ingestion/stage/table_usage.py

@@ -76,7 +76,7 @@ class TableUsageStage(Stage):
         metadata: OpenMetadata,
         pipeline_name: Optional[str] = None,
     ):
-        config = TableStageConfig.parse_obj(config_dict)
+        config = TableStageConfig.model_validate(config_dict)
         return cls(config, metadata)
 
     def init_location(self) -> None:
@@ -100,7 +100,7 @@ class TableUsageStage(Stage):
         if username:
             user = self.metadata.get_by_name(entity=User, fqn=username)
             if user:
-                return [user.fullyQualifiedName.__root__], [username]
+                return [user.fullyQualifiedName.root], [username]
             return None, [username]
         return None, None
 
@@ -191,7 +191,7 @@ class TableUsageStage(Stage):
         for key, value in self.table_usage.items():
             if value:
                 value.sqlQueries = self.table_queries.get(key, [])
-                data = value.json()
+                data = value.model_dump_json()
                 with open(
                     os.path.join(self.config.filename, f"{value.serviceName}_{key[1]}"),
                     "a+",

metadata/mixins/pandas/pandas_mixin.py

@@ -91,7 +91,7 @@ class PandasInterfaceMixin:
             config_source=service_connection_config.configSource,
             client=client,
             file_fqn=DatalakeTableSchemaWrapper(
-                key=table.name.__root__,
+                key=table.name.root,
                 bucket_name=table.databaseSchema.name,
                 file_extension=table.fileFormat,
             ),
@@ -125,4 +125,4 @@ class PandasInterfaceMixin:
                         for df in data
                     ]
             return data
-        raise TypeError(f"Couldn't fetch {table.name.__root__}")
+        raise TypeError(f"Couldn't fetch {table.name.root}")

metadata/mixins/sqalchemy/sqa_mixin.py

@@ -112,5 +112,5 @@ class SQAInterfaceMixin:
         return [
             column.name
             for column in self.table.__table__.columns
-            if column.name in {col.name.__root__ for col in self.table_entity.columns}
+            if column.name in {col.name.root for col in self.table_entity.columns}
         ]

metadata/parsers/avro_parser.py

@@ -14,11 +14,11 @@ Utils module to parse the avro schema
 """
 
 import traceback
-from typing import List, Optional, Tuple, Union
+from typing import List, Optional, Tuple, Type, Union
 
 import avro.schema as avroschema
 from avro.schema import ArraySchema, RecordSchema, Schema, UnionSchema
-from pydantic.main import ModelMetaclass
+from pydantic import BaseModel
 
 from metadata.generated.schema.entity.data.table import Column
 from metadata.generated.schema.type.schema import FieldModel
@@ -30,7 +30,7 @@ RECORD_DATATYPE_NAME = "RECORD"
 
 
 def _parse_array_children(
-    arr_item: Schema, cls: ModelMetaclass = FieldModel
+    arr_item: Schema, cls: Type[BaseModel] = FieldModel
 ) -> Tuple[str, Optional[Union[FieldModel, Column]]]:
     if isinstance(arr_item, ArraySchema):
         display_type, children = _parse_array_children(arr_item.items, cls=cls)
@@ -55,7 +55,7 @@ def _parse_array_children(
 
 
 def parse_array_fields(
-    field: ArraySchema, cls: ModelMetaclass = FieldModel
+    field: ArraySchema, cls: Type[BaseModel] = FieldModel
 ) -> Optional[List[Union[FieldModel, Column]]]:
     """
     Parse array field for avro schema
@@ -106,7 +106,9 @@ def parse_array_fields(
 
 
 def _parse_union_children(
-    parent: Optional[Schema], union_field: UnionSchema, cls: ModelMetaclass = FieldModel
+    parent: Optional[Schema],
+    union_field: UnionSchema,
+    cls: Type[BaseModel] = FieldModel,
 ) -> Tuple[str, Optional[Union[FieldModel, Column]]]:
     non_null_schema = [
         (i, schema)
@@ -137,7 +139,7 @@ def _parse_union_children(
     return sub_type, None
 
 
-def parse_record_fields(field: RecordSchema, cls: ModelMetaclass = FieldModel):
+def parse_record_fields(field: RecordSchema, cls: Type[BaseModel] = FieldModel):
     """
     Parse the nested record fields for avro
     """
@@ -160,7 +162,7 @@ def parse_record_fields(field: RecordSchema, cls: ModelMetaclass = FieldModel):
 def parse_union_fields(
     parent: Optional[Schema],
     union_field: Schema,
-    cls: ModelMetaclass = FieldModel,
+    cls: Type[BaseModel] = FieldModel,
 ) -> Optional[List[Union[FieldModel, Column]]]:
     """
     Parse union field for avro schema
@@ -209,7 +211,7 @@ def parse_union_fields(
 
 
 def parse_single_field(
-    field: Schema, cls: ModelMetaclass = FieldModel
+    field: Schema, cls: Type[BaseModel] = FieldModel
 ) -> Optional[List[Union[FieldModel, Column]]]:
     """
     Parse primitive field for avro schema
@@ -224,7 +226,7 @@ def parse_single_field(
 
 
 def parse_avro_schema(
-    schema: str, cls: ModelMetaclass = FieldModel
+    schema: str, cls: Type[BaseModel] = FieldModel
 ) -> Optional[List[Union[FieldModel, Column]]]:
     """
     Method to parse the avro schema
@@ -247,7 +249,7 @@ def parse_avro_schema(
 
 
 def get_avro_fields(
-    parsed_schema: Schema, cls: ModelMetaclass = FieldModel
+    parsed_schema: Schema, cls: Type[BaseModel] = FieldModel
 ) -> Optional[List[Union[FieldModel, Column]]]:
     """
     Recursively convert the parsed schema into required models

metadata/parsers/json_schema_parser.py

@@ -16,9 +16,9 @@ Utils module to parse the jsonschema
 import json
 import traceback
 from enum import Enum
-from typing import List, Optional
+from typing import List, Optional, Type
 
-from pydantic.main import ModelMetaclass
+from pydantic import BaseModel
 
 from metadata.generated.schema.type.schema import FieldModel
 from metadata.utils.logger import ingestion_logger
@@ -42,7 +42,7 @@ class JsonSchemaDataTypes(Enum):
 
 
 def parse_json_schema(
-    schema_text: str, cls: ModelMetaclass = FieldModel
+    schema_text: str, cls: Type[BaseModel] = FieldModel
 ) -> Optional[List[FieldModel]]:
     """
     Method to parse the jsonschema
@@ -67,7 +67,7 @@ def parse_json_schema(
 
 
 def get_json_schema_fields(
-    properties, cls: ModelMetaclass = FieldModel
+    properties, cls: Type[BaseModel] = FieldModel
 ) -> Optional[List[FieldModel]]:
     """
     Recursively convert the parsed schema into required models

metadata/parsers/protobuf_parser.py

@@ -20,11 +20,10 @@ import sys
 import traceback
 from enum import Enum
 from pathlib import Path
-from typing import List, Optional, Union
+from typing import List, Optional, Type, Union
 
 import grpc_tools.protoc
 from pydantic import BaseModel
-from pydantic.main import ModelMetaclass
 
 from metadata.generated.schema.entity.data.table import Column, DataType
 from metadata.generated.schema.type.schema import DataTypeTopic, FieldModel
@@ -168,7 +167,7 @@ class ProtobufParser:
         return None
 
     def parse_protobuf_schema(
-        self, cls: ModelMetaclass = FieldModel
+        self, cls: Type[BaseModel] = FieldModel
     ) -> Optional[List[Union[FieldModel, Column]]]:
         """
         Method to parse the protobuf schema
@@ -202,7 +201,7 @@ class ProtobufParser:
             )
         return None
 
-    def _get_field_type(self, type_: int, cls: ModelMetaclass = FieldModel) -> str:
+    def _get_field_type(self, type_: int, cls: Type[BaseModel] = FieldModel) -> str:
         if type_ > 18:
             return DataType.UNKNOWN.value
         data_type = ProtobufDataTypes(type_).name
@@ -211,7 +210,7 @@ class ProtobufParser:
         return data_type
 
     def get_protobuf_fields(
-        self, fields, cls: ModelMetaclass = FieldModel
+        self, fields, cls: Type[BaseModel] = FieldModel
     ) -> Optional[List[Union[FieldModel, Column]]]:
         """
         Recursively convert the parsed schema into required models

metadata/pii/processor.py

@@ -120,9 +120,7 @@ class PIIProcessor(Processor):
 
         # First, check if the column we are about to process
         # already has PII tags or not
-        column_has_pii_tag = any(
-            (PII in tag.tagFQN.__root__ for tag in column.tags or [])
-        )
+        column_has_pii_tag = any((PII in tag.tagFQN.root for tag in column.tags or []))
 
         # If it has PII tags, we skip the processing
         # for the column
@@ -130,7 +128,7 @@ class PIIProcessor(Processor):
             return None
 
         # Scan by column name. If no results there, check the sample data, if any
-        tag_and_confidence = ColumnNameScanner.scan(column.name.__root__) or (
+        tag_and_confidence = ColumnNameScanner.scan(column.name.root) or (
             self.ner_scanner.scan([row[idx] for row in table_data.rows])
             if table_data
             else None
@@ -145,7 +143,7 @@ class PIIProcessor(Processor):
             return [
                 self.build_column_tag(
                     tag_fqn=tag_and_confidence.tag_fqn,
-                    column_fqn=column.fullyQualifiedName.__root__,
+                    column_fqn=column.fullyQualifiedName.root,
                 )
             ]
 
@@ -180,7 +178,7 @@ class PIIProcessor(Processor):
             except Exception as err:
                 self.status.failed(
                     StackTraceError(
-                        name=record.table.fullyQualifiedName.__root__,
+                        name=record.table.fullyQualifiedName.root,
                         error=f"Error computing PII tags for [{column}] - [{err}]",
                         stackTrace=traceback.format_exc(),
                     )

metadata/profiler/adaptors/dynamodb.py

@@ -29,12 +29,12 @@ class DynamoDB(NoSQLAdaptor):
         self.client = client
 
     def item_count(self, table: Table) -> int:
-        table = self.client.Table(table.name.__root__)
+        table = self.client.Table(table.name.root)
         return table.item_count
 
     def scan(
         self, table: Table, columns: List[Column], limit: int
     ) -> List[Dict[str, any]]:
-        table = self.client.Table(table.name.__root__)
+        table = self.client.Table(table.name.root)
         response = table.scan(Limit=limit)
         return response["Items"]

metadata/profiler/adaptors/mongodb.py

@@ -92,7 +92,7 @@ class MongoDB(NoSQLAdaptor):
 
     def item_count(self, table: Table) -> int:
         db = self.client[table.databaseSchema.name]
-        collection = db[table.name.__root__]
+        collection = db[table.name.root]
         return collection.count_documents({})
 
     def scan(
@@ -101,7 +101,7 @@ class MongoDB(NoSQLAdaptor):
         return self.execute(
             Query(
                 database=table.databaseSchema.name,
-                collection=table.name.__root__,
+                collection=table.name.root,
                 limit=limit,
             )
         )
@@ -116,7 +116,7 @@ class MongoDB(NoSQLAdaptor):
         return self.execute(
             Query(
                 database=table.databaseSchema.name,
-                collection=table.name.__root__,
+                collection=table.name.root,
                 filter=json_query,
             )
         )
@@ -143,7 +143,7 @@ class MongoDB(NoSQLAdaptor):
         row = self.execute(
             Aggregation(
                 database=table.databaseSchema.name,
-                collection=table.name.__root__,
+                collection=table.name.root,
                 column=column.name,
                 aggregations=aggregate_functions,
             )
@@ -163,16 +163,4 @@ class MongoDB(NoSQLAdaptor):
         return AggregationFunction.MIN
 
     def execute(self, query: Executable) -> List[Dict[str, any]]:
-        records = list(query.to_executable(self.client))
-        result = []
-        for r in records:
-            result.append({c: self._json_safe(r.get(c)) for c in r})
-        return result
-
-    @staticmethod
-    def _json_safe(data: any):
-        try:
-            json.dumps(data)
-            return data
-        except Exception:  # noqa
-            return str(data)
+        return list(query.to_executable(self.client))

metadata/profiler/api/models.py

@@ -17,6 +17,7 @@ multiple profilers per table and columns.
 """
 from typing import List, Optional, Type, Union
 
+from pydantic import ConfigDict
 from sqlalchemy import Column
 from sqlalchemy.orm import DeclarativeMeta
 
@@ -45,8 +46,8 @@ from metadata.utils.sqa_like_column import SQALikeColumn
 class ColumnConfig(ConfigModel):
     """Column config for profiler"""
 
-    excludeColumns: Optional[List[str]]
-    includeColumns: Optional[List[ColumnProfilerConfig]]
+    excludeColumns: Optional[List[str]] = None
+    includeColumns: Optional[List[ColumnProfilerConfig]] = None
 
 
 class BaseProfileConfig(ConfigModel):
@@ -62,8 +63,8 @@ class TableConfig(BaseProfileConfig):
     """table profile config"""
 
     profileQuery: Optional[str] = None
-    partitionConfig: Optional[PartitionProfilerConfig]
-    columnConfig: Optional[ColumnConfig]
+    partitionConfig: Optional[PartitionProfilerConfig] = None
+    columnConfig: Optional[ColumnConfig] = None
 
     @classmethod
     def from_database_and_schema_config(
@@ -118,16 +119,15 @@ class ProfilerResponse(ConfigModel):
 
     def __str__(self):
         """Return the table name being processed"""
-        return f"Table [{self.table.name.__root__}]"
+        return f"Table [{self.table.name.root}]"
 
 
 class ThreadPoolMetrics(ConfigModel):
     """A container for all metrics to be computed on the same thread."""
 
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
     metrics: Union[List[Union[Type[Metric], CustomMetric]], Type[Metric]]
     metric_type: MetricTypes
-    column: Optional[Union[Column, SQALikeColumn]]
+    column: Optional[Union[Column, SQALikeColumn]] = None
     table: Union[Table, DeclarativeMeta]
-
-    class Config:
-        arbitrary_types_allowed = True

metadata/profiler/interface/nosql/profiler_interface.py

@@ -16,7 +16,7 @@ supporting sqlalchemy abstraction layer
 """
 import traceback
 from collections import defaultdict
-from datetime import datetime, timezone
+from datetime import datetime
 from typing import Dict, List, Optional, Type
 
 from sqlalchemy import Column
@@ -148,9 +148,9 @@ class NoSQLProfilerInterface(ProfilerInterface):
             row = None
         if metric_func.column is not None:
             column = metric_func.column.name
-            self.status.scanned(f"{metric_func.table.name.__root__}.{column}")
+            self.status.scanned(f"{metric_func.table.name.root}.{column}")
         else:
-            self.status.scanned(metric_func.table.name.__root__)
+            self.status.scanned(metric_func.table.name.root)
             column = None
         return row, column, metric_func.metric_type.value
 
@@ -212,9 +212,7 @@ class NoSQLProfilerInterface(ProfilerInterface):
                     profile_results["columns"][column].update(
                         {
                             "name": column,
-                            "timestamp": int(
-                                datetime.now(tz=timezone.utc).timestamp() * 1000
-                            ),
+                            "timestamp": int(datetime.now().timestamp() * 1000),
                             **profile,
                         }
                     )
@@ -227,8 +225,7 @@ class NoSQLProfilerInterface(ProfilerInterface):
 
     def get_columns(self) -> List[Optional[SQALikeColumn]]:
         return [
-            SQALikeColumn(name=c.name.__root__, type=c.dataType)
-            for c in self.table.columns
+            SQALikeColumn(name=c.name.root, type=c.dataType) for c in self.table.columns
         ]
 
     def close(self):

metadata/profiler/interface/pandas/profiler_interface.py

@@ -17,7 +17,7 @@ supporting sqlalchemy abstraction layer
 import traceback
 from collections import defaultdict
 from copy import deepcopy
-from datetime import datetime, timezone
+from datetime import datetime
 from typing import Dict, List, Optional
 
 from sqlalchemy import Column
@@ -88,7 +88,7 @@ class PandasProfilerInterface(ProfilerInterface, PandasInterfaceMixin):
         self.client = self.connection.client
         self.dfs = self.return_ometa_dataframes_sampled(
             service_connection_config=self.service_connection_config,
-            client=self.client,
+            client=self.client._client,
             table=self.table_entity,
             profile_sample_config=profile_sample_config,
         )
@@ -137,7 +137,7 @@ class PandasProfilerInterface(ProfilerInterface, PandasInterfaceMixin):
 
         return sampler_factory_.create(
             DatalakeConnection.__name__,
-            client=self.client,
+            client=self.client._client,  # pylint: disable=W0212
             table=self.dfs,
             profile_sample_config=self.profile_sample_config,
             partition_details=self.partition_details,
@@ -287,7 +287,7 @@ class PandasProfilerInterface(ProfilerInterface, PandasInterfaceMixin):
                     if len(df.query(metric.expression).index)
                 )
                 custom_metrics.append(
-                    CustomMetricProfile(name=metric.name.__root__, value=row)
+                    CustomMetricProfile(name=metric.name.root, value=row)
                 )
 
             except Exception as exc:
@@ -303,7 +303,7 @@ class PandasProfilerInterface(ProfilerInterface, PandasInterfaceMixin):
         metric_func: ThreadPoolMetrics,
     ):
         """Run metrics in processor worker"""
-        logger.debug(f"Running profiler for {metric_func.table.name.__root__}")
+        logger.debug(f"Running profiler for {metric_func.table.name.root}")
         try:
             row = None
             if self.complex_dataframe_sample:
@@ -320,9 +320,9 @@ class PandasProfilerInterface(ProfilerInterface, PandasInterfaceMixin):
             row = None
         if metric_func.column is not None:
             column = metric_func.column.name
-            self.status.scanned(f"{metric_func.table.name.__root__}.{column}")
+            self.status.scanned(f"{metric_func.table.name.root}.{column}")
         else:
-            self.status.scanned(metric_func.table.name.__root__)
+            self.status.scanned(metric_func.table.name.root)
             column = None
         return row, column, metric_func.metric_type.value
 
@@ -403,9 +403,7 @@ class PandasProfilerInterface(ProfilerInterface, PandasInterfaceMixin):
                         profile_results["columns"][column].update(
                             {
                                 "name": column,
-                                "timestamp": int(
-                                    datetime.now(tz=timezone.utc).timestamp() * 1000
-                                ),
+                                "timestamp": int(datetime.now().timestamp() * 1000),
                                 **profile,
                             }
                         )

metadata/profiler/interface/profiler_interface.py

@@ -117,7 +117,7 @@ class ProfilerInterface(ABC):
         except AttributeError:
             self.status.entity = None
         else:
-            self.status.entity = fqn.__root__ if fqn else None
+            self.status.entity = fqn.root if fqn else None
         self.profile_sample_config = profile_sample_config
         self.profile_query = sample_query
         self.partition_details = (
@@ -272,7 +272,7 @@ class ProfilerInterface(ABC):
         """
         for schema_config in profiler_config.schemaConfig:
             if (
-                schema_config.fullyQualifiedName.__root__
+                schema_config.fullyQualifiedName.root
                 == entity.databaseSchema.fullyQualifiedName
                 and ProfilerInterface._get_sample_storage_config(schema_config)
             ):
@@ -280,7 +280,7 @@ class ProfilerInterface(ABC):
 
         for database_config in profiler_config.databaseConfig:
             if (
-                database_config.fullyQualifiedName.__root__
+                database_config.fullyQualifiedName.root
                 == entity.database.fullyQualifiedName
                 and ProfilerInterface._get_sample_storage_config(database_config)
             ):