openmetadata-ingestion 1.2.4.0__py3-none-any.whl → 1.3.2.0rc1__py3-none-any.whl

metadata/ingestion/source/database/{databricks/unity_catalog → unitycatalog}/metadata.py

@@ -13,11 +13,10 @@ Databricks Unity Catalog Source source methods.
 """
 import json
 import traceback
-from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
+from typing import Any, Iterable, List, Optional, Tuple, Union
 
 from databricks.sdk.service.catalog import ColumnInfo
 from databricks.sdk.service.catalog import TableConstraint as DBTableConstraint
-from databricks.sdk.service.catalog import TableConstraintList
 
 from metadata.generated.schema.api.data.createDatabase import CreateDatabaseRequest
 from metadata.generated.schema.api.data.createDatabaseSchema import (
@@ -38,8 +37,11 @@ from metadata.generated.schema.entity.data.table import (
     TableConstraint,
     TableType,
 )
-from metadata.generated.schema.entity.services.connections.database.databricksConnection import (
-    DatabricksConnection,
+from metadata.generated.schema.entity.services.connections.database.unityCatalogConnection import (
+    UnityCatalogConnection,
+)
+from metadata.generated.schema.entity.services.ingestionPipelines.status import (
+    StackTraceError,
 )
 from metadata.generated.schema.metadataIngestion.databaseServiceMetadataPipeline import (
     DatabaseServiceMetadataPipeline,
@@ -47,22 +49,22 @@ from metadata.generated.schema.metadataIngestion.databaseServiceMetadataPipeline
 from metadata.generated.schema.metadataIngestion.workflow import (
     Source as WorkflowSource,
 )
-from metadata.ingestion.api.models import Either, StackTraceError
+from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import InvalidSourceException
 from metadata.ingestion.lineage.sql_lineage import get_column_fqn
 from metadata.ingestion.models.ometa_classification import OMetaTagAndClassification
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.database.column_type_parser import ColumnTypeParser
 from metadata.ingestion.source.database.database_service import DatabaseServiceSource
-from metadata.ingestion.source.database.databricks.connection import get_connection
-from metadata.ingestion.source.database.databricks.models import (
+from metadata.ingestion.source.database.multi_db_source import MultiDBSource
+from metadata.ingestion.source.database.stored_procedures_mixin import QueryByProcedure
+from metadata.ingestion.source.database.unitycatalog.connection import get_connection
+from metadata.ingestion.source.database.unitycatalog.models import (
     ColumnJson,
     ElementType,
     ForeignConstrains,
     Type,
 )
-from metadata.ingestion.source.database.multi_db_source import MultiDBSource
-from metadata.ingestion.source.database.stored_procedures_mixin import QueryByProcedure
 from metadata.ingestion.source.models import TableView
 from metadata.utils import fqn
 from metadata.utils.db_utils import get_view_lineage
@@ -72,20 +74,7 @@ from metadata.utils.logger import ingestion_logger
 logger = ingestion_logger()
 
 
-# pylint: disable=not-callable
-@classmethod
-def from_dict(cls, dct: Dict[str, Any]) -> "TableConstraintList":
-    return cls(
-        table_constraints=[
-            DBTableConstraint.from_dict(constraint) for constraint in dct
-        ]
-    )
-
-
-TableConstraintList.from_dict = from_dict
-
-
-class DatabricksUnityCatalogSource(DatabaseServiceSource, MultiDBSource):
+class UnitycatalogSource(DatabaseServiceSource, MultiDBSource):
     """
     Implements the necessary methods to extract
     Database metadata from Databricks Source using
@@ -100,7 +89,7 @@ class DatabricksUnityCatalogSource(DatabaseServiceSource, MultiDBSource):
         )
         self.context.table_views = []
         self.metadata = metadata
-        self.service_connection: DatabricksConnection = (
+        self.service_connection: UnityCatalogConnection = (
             self.config.serviceConnection.__root__.config
         )
         self.client = get_connection(self.service_connection)
@@ -118,10 +107,10 @@ class DatabricksUnityCatalogSource(DatabaseServiceSource, MultiDBSource):
     @classmethod
     def create(cls, config_dict, metadata: OpenMetadata):
         config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
-        connection: DatabricksConnection = config.serviceConnection.__root__.config
-        if not isinstance(connection, DatabricksConnection):
+        connection: UnityCatalogConnection = config.serviceConnection.__root__.config
+        if not isinstance(connection, UnityCatalogConnection):
             raise InvalidSourceException(
-                f"Expected DatabricksConnection, but got {connection}"
+                f"Expected UnityCatalogConnection, but got {connection}"
             )
         return cls(config, metadata)
 
@@ -164,7 +153,7 @@ class DatabricksUnityCatalogSource(DatabaseServiceSource, MultiDBSource):
                         StackTraceError(
                             name=catalog_name,
                             error=f"Unexpected exception to get database name [{catalog_name}]: {exc}",
-                            stack_trace=traceback.format_exc(),
+                            stackTrace=traceback.format_exc(),
                         )
                     )
 
@@ -210,7 +199,7 @@ class DatabricksUnityCatalogSource(DatabaseServiceSource, MultiDBSource):
                     StackTraceError(
                         name=schema.name,
                         error=f"Unexpected exception to get database schema [{schema.name}]: {exc}",
-                        stack_trace=traceback.format_exc(),
+                        stackTrace=traceback.format_exc(),
                     )
                 )
 
@@ -233,7 +222,7 @@ class DatabricksUnityCatalogSource(DatabaseServiceSource, MultiDBSource):
             )
         )
 
-    def get_tables_name_and_type(self) -> Optional[Iterable[Tuple[str, str]]]:
+    def get_tables_name_and_type(self) -> Iterable[Tuple[str, str]]:
         """
         Handle table and views.
 
@@ -281,7 +270,7 @@ class DatabricksUnityCatalogSource(DatabaseServiceSource, MultiDBSource):
                     StackTraceError(
                         name=table.Name,
                         error=f"Unexpected exception to get table [{table.Name}]: {exc}",
-                        stack_trace=traceback.format_exc(),
+                        stackTrace=traceback.format_exc(),
                     )
                 )
 
@@ -343,12 +332,12 @@ class DatabricksUnityCatalogSource(DatabaseServiceSource, MultiDBSource):
                 left=StackTraceError(
                     name=table_name,
                     error=f"Unexpected exception to yield table [{table_name}]: {exc}",
-                    stack_trace=traceback.format_exc(),
+                    stackTrace=traceback.format_exc(),
                 )
             )
 
     def get_table_constraints(
-        self, constraints: TableConstraintList
+        self, constraints: List[DBTableConstraint]
     ) -> Tuple[List[TableConstraint], List[ForeignConstrains]]:
         """
         Function to handle table constraint for the current table and add it to context
@@ -356,23 +345,22 @@ class DatabricksUnityCatalogSource(DatabaseServiceSource, MultiDBSource):
 
         primary_constraints = []
         foreign_constraints = []
-        if constraints and constraints.table_constraints:
-            for constraint in constraints.table_constraints:
-                if constraint.primary_key_constraint:
-                    primary_constraints.append(
-                        TableConstraint(
-                            constraintType=ConstraintType.PRIMARY_KEY,
-                            columns=constraint.primary_key_constraint.child_columns,
-                        )
+        for constraint in constraints:
+            if constraint.primary_key_constraint:
+                primary_constraints.append(
+                    TableConstraint(
+                        constraintType=ConstraintType.PRIMARY_KEY,
+                        columns=constraint.primary_key_constraint.child_columns,
                     )
-                if constraint.foreign_key_constraint:
-                    foreign_constraints.append(
-                        ForeignConstrains(
-                            child_columns=constraint.foreign_key_constraint.child_columns,
-                            parent_columns=constraint.foreign_key_constraint.parent_columns,
-                            parent_table=constraint.foreign_key_constraint.parent_table,
-                        )
+                )
+            if constraint.foreign_key_constraint:
+                foreign_constraints.append(
+                    ForeignConstrains(
+                        child_columns=constraint.foreign_key_constraint.child_columns,
+                        parent_columns=constraint.foreign_key_constraint.parent_columns,
+                        parent_table=constraint.foreign_key_constraint.parent_table,
                     )
+                )
         return primary_constraints, foreign_constraints
 
     def _get_foreign_constraints(self, foreign_columns) -> List[TableConstraint]:
@@ -471,11 +459,16 @@ class DatabricksUnityCatalogSource(DatabaseServiceSource, MultiDBSource):
             )
 
     def get_columns(self, column_data: List[ColumnInfo]) -> Iterable[Column]:
-        # process table regular columns info
+        """
+        process table regular columns info
+        """
 
         for column in column_data:
             if column.type_text.lower().startswith("union"):
                 column.type_text = column.Type.replace(" ", "")
+            if column.type_text.lower() == "struct":
+                column.type_text = "struct<>"
+
             parsed_string = ColumnTypeParser._parse_datatype_string(  # pylint: disable=protected-access
                 column.type_text.lower()
             )

metadata/ingestion/source/database/unitycatalog/query_parser.py

@@ -0,0 +1,60 @@
+#  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.
+"""
+UnityCatalog Query parser module
+"""
+from abc import ABC
+
+from metadata.generated.schema.entity.services.connections.database.unityCatalogConnection import (
+    UnityCatalogConnection,
+)
+from metadata.generated.schema.metadataIngestion.workflow import (
+    Source as WorkflowSource,
+)
+from metadata.ingestion.api.steps import InvalidSourceException
+from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.ingestion.source.database.databricks.query_parser import (
+    DatabricksQueryParserSource,
+)
+from metadata.ingestion.source.database.query_parser_source import QueryParserSource
+from metadata.ingestion.source.database.unitycatalog.client import UnityCatalogClient
+from metadata.utils.logger import ingestion_logger
+
+logger = ingestion_logger()
+
+
+class UnityCatalogQueryParserSource(
+    DatabricksQueryParserSource, QueryParserSource, ABC
+):
+    """
+    UnityCatalog Query Parser Source
+
+    This class would be inheriting all the methods
+    from DatabricksQueryParserSource
+    """
+
+    filters: str
+
+    # pylint: disable=super-init-not-called
+    def __init__(self, config: WorkflowSource, metadata: OpenMetadata):
+        self._init_super(config=config, metadata=metadata)
+        self.client = UnityCatalogClient(self.service_connection)
+
+    @classmethod
+    def create(cls, config_dict, metadata: OpenMetadata):
+        """Create class instance"""
+        config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
+        connection: UnityCatalogConnection = config.serviceConnection.__root__.config
+        if not isinstance(connection, UnityCatalogConnection):
+            raise InvalidSourceException(
+                f"Expected UnityCatalogConnection, but got {connection}"
+            )
+        return cls(config, metadata)

metadata/ingestion/source/database/unitycatalog/usage.py

@@ -0,0 +1,31 @@
+#  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.
+"""
+unity catalog usage module
+"""
+
+from metadata.ingestion.source.database.databricks.usage import DatabricksUsageSource
+from metadata.ingestion.source.database.unitycatalog.query_parser import (
+    UnityCatalogQueryParserSource,
+)
+from metadata.utils.logger import ingestion_logger
+
+logger = ingestion_logger()
+
+
+class UnitycatalogUsageSource(UnityCatalogQueryParserSource, DatabricksUsageSource):
+    """
+    UnityCatalog Usage Source
+
+    This class would be inheriting all the methods from
+    DatabricksUsageSource as both the sources would call
+    the same API for fetching Usage Queries
+    """

metadata/ingestion/source/database/usage_source.py

@@ -107,12 +107,13 @@ class UsageSource(QueryParserSource, ABC):
                             row = dict(row)
                             try:
                                 query_type = row.get("query_type")
+                                query = self.format_query(row["query_text"])
                                 queries.append(
                                     TableQuery(
-                                        query=self.format_query(row["query_text"]),
+                                        query=query,
                                         query_type=query_type,
                                         exclude_usage=self.check_life_cycle_query(
-                                            query_type=query_type
+                                            query_type=query_type, query_text=query
                                         ),
                                         userName=row["user_name"],
                                         startTime=str(row["start_time"]),

metadata/ingestion/source/messaging/common_broker_source.py

@@ -27,11 +27,14 @@ from confluent_kafka.schema_registry.schema_registry_client import Schema
 from metadata.generated.schema.api.data.createTopic import CreateTopicRequest
 from metadata.generated.schema.entity.data.topic import Topic as TopicEntity
 from metadata.generated.schema.entity.data.topic import TopicSampleData
+from metadata.generated.schema.entity.services.ingestionPipelines.status import (
+    StackTraceError,
+)
 from metadata.generated.schema.metadataIngestion.workflow import (
     Source as WorkflowSource,
 )
 from metadata.generated.schema.type.schema import SchemaType, Topic
-from metadata.ingestion.api.models import Either, StackTraceError
+from metadata.ingestion.api.models import Either
 from metadata.ingestion.models.ometa_topic_data import OMetaTopicSampleData
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.messaging.messaging_service import (
@@ -145,7 +148,7 @@ class CommonBrokerSource(MessagingServiceSource, ABC):
                 left=StackTraceError(
                     name=topic_details.topic_name,
                     error=f"Unexpected exception to yield topic [{topic_details}]: {exc}",
-                    stack_trace=traceback.format_exc(),
+                    stackTrace=traceback.format_exc(),
                 )
             )
 
@@ -223,27 +226,28 @@ class CommonBrokerSource(MessagingServiceSource, ABC):
         if topic_entity and self.generate_sample_data:
             topic_name = topic_details.topic_name
             sample_data = []
+            messages = None
             try:
-                self.consumer_client.subscribe(
-                    [topic_name], on_assign=on_partitions_assignment_to_consumer
-                )
-                logger.info(
-                    f"Broker consumer polling for sample messages in topic {topic_name}"
-                )
-                messages = self.consumer_client.consume(num_messages=10, timeout=10)
+                if self.consumer_client:
+                    self.consumer_client.subscribe(
+                        [topic_name], on_assign=on_partitions_assignment_to_consumer
+                    )
+                    logger.info(
+                        f"Broker consumer polling for sample messages in topic {topic_name}"
+                    )
+                    messages = self.consumer_client.consume(num_messages=10, timeout=10)
             except Exception as exc:
                 yield Either(
                     left=StackTraceError(
                         name=topic_details.topic_name,
                         error=f"Failed to fetch sample data from topic {topic_name}: {exc}",
-                        stack_trace=traceback.format_exc(),
+                        stackTrace=traceback.format_exc(),
                     )
                 )
             else:
                 if messages:
                     for message in messages:
                         try:
-
                             value = message.value()
                             sample_data.append(
                                 self.decode_message(

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

@@ -15,7 +15,7 @@ Source connection handler
 from dataclasses import dataclass
 from typing import Optional, Union
 
-from confluent_kafka.admin import AdminClient
+from confluent_kafka.admin import AdminClient, KafkaException
 from confluent_kafka.avro import AvroConsumer
 from confluent_kafka.schema_registry.schema_registry_client import SchemaRegistryClient
 
@@ -35,6 +35,21 @@ from metadata.utils.logger import ingestion_logger
 logger = ingestion_logger()
 
 
+class InvalidKafkaCreds(Exception):
+    """
+    Class to indicate invalid kafka credentials exception
+    """
+
+
+class SchemaRegistryException(Exception):
+    """
+    Class to indicate invalid schema registry not initialized
+    """
+
+
+TIMEOUT_SECONDS = 10
+
+
 @dataclass
 class KafkaClient:
     def __init__(self, admin_client, schema_registry_client, consumer_client) -> None:
@@ -60,6 +75,14 @@ def get_connection(
         if connection.saslMechanism:
             connection.consumerConfig["sasl.mechanism"] = connection.saslMechanism.value
 
+        if (
+            connection.consumerConfig.get("security.protocol") is None
+            and connection.securityProtocol
+        ):
+            connection.consumerConfig[
+                "security.protocol"
+            ] = connection.securityProtocol.value
+
     if connection.basicAuthUserInfo:
         connection.schemaRegistryConfig = connection.schemaRegistryConfig or {}
         connection.schemaRegistryConfig[
@@ -109,9 +132,27 @@ def test_connection(
     """
 
     def custom_executor():
-        _ = client.admin_client.list_topics().topics
-
-    test_fn = {"GetTopics": custom_executor}
+        try:
+            client.admin_client.list_topics(timeout=TIMEOUT_SECONDS).topics
+        except KafkaException as err:
+            raise InvalidKafkaCreds(
+                f"Failed to fetch topics due to: {err}. "
+                "Please validate credentials and check if you are using correct security protocol"
+            )
+
+    def schema_registry_test():
+        if client.schema_registry_client:
+            client.schema_registry_client.get_subjects()
+        else:
+            raise SchemaRegistryException(
+                "Schema Registry not initialized, please provide schema registry "
+                "credentials in case you want topic schema and sample data to be ingested"
+            )
+
+    test_fn = {
+        "GetTopics": custom_executor,
+        "CheckSchemaRegistry": schema_registry_test,
+    }
 
     test_connection_steps(
         metadata=metadata,

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

@@ -21,11 +21,14 @@ from metadata.generated.schema.entity.data.topic import TopicSampleData
 from metadata.generated.schema.entity.services.connections.messaging.kinesisConnection import (
     KinesisConnection,
 )
+from metadata.generated.schema.entity.services.ingestionPipelines.status import (
+    StackTraceError,
+)
 from metadata.generated.schema.metadataIngestion.workflow import (
     Source as WorkflowSource,
 )
 from metadata.generated.schema.type.schema import Topic
-from metadata.ingestion.api.models import Either, StackTraceError
+from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import InvalidSourceException
 from metadata.ingestion.models.ometa_topic_data import OMetaTopicSampleData
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
@@ -137,7 +140,7 @@ class KinesisSource(MessagingServiceSource):
                 left=StackTraceError(
                     name=topic_details.topic_name,
                     error=f"Unexpected exception to yield topic [{topic_details}]: {exc}",
-                    stack_trace=traceback.format_exc(),
+                    stackTrace=traceback.format_exc(),
                 )
             )
 
@@ -215,7 +218,7 @@ class KinesisSource(MessagingServiceSource):
                 left=StackTraceError(
                     name=topic_details.topic_name,
                     error=f"Error while yielding topic sample data for topic: {topic_details.topic_name} - {err}",
-                    stack_trace=traceback.format_exc(),
+                    stackTrace=traceback.format_exc(),
                 )
             )
 

metadata/ingestion/source/messaging/messaging_service.py

@@ -37,8 +37,8 @@ from metadata.ingestion.models.delete_entity import DeleteEntity
 from metadata.ingestion.models.topology import (
     NodeStage,
     ServiceTopology,
+    TopologyContext,
     TopologyNode,
-    create_source_context,
 )
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.connections import get_connection, get_test_connection_fn
@@ -114,7 +114,7 @@ class MessagingServiceSource(TopologyRunnerMixin, Source, ABC):
     service_connection: MessagingConnection.__fields__["config"].type_
 
     topology = MessagingServiceTopology()
-    context = create_source_context(topology)
+    context = TopologyContext.create(topology)
     topic_source_state: Set = set()
 
     def __init__(
@@ -135,6 +135,10 @@ class MessagingServiceSource(TopologyRunnerMixin, Source, ABC):
         self.connection_obj = self.connection
         self.test_connection()
 
+    @property
+    def name(self) -> str:
+        return self.service_connection.type.name
+
     @abstractmethod
     def yield_topic(self, topic_details: Any) -> Iterable[Either[CreateTopicRequest]]:
         """

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

@@ -43,13 +43,16 @@ from metadata.generated.schema.entity.services.databaseService import (
     DatabaseService,
     DatabaseServiceType,
 )
+from metadata.generated.schema.entity.services.ingestionPipelines.status import (
+    StackTraceError,
+)
 from metadata.generated.schema.entity.teams import team
 from metadata.generated.schema.entity.teams.user import User
 from metadata.generated.schema.metadataIngestion.workflow import (
     Source as WorkflowSource,
 )
 from metadata.ingestion.api.common import Entity
-from metadata.ingestion.api.models import Either, StackTraceError
+from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import InvalidSourceException, Source
 from metadata.ingestion.models.user import OMetaUserProfile
 from metadata.ingestion.ometa.client_utils import get_chart_entities_from_id
@@ -174,7 +177,7 @@ class AmundsenSource(Source):
                 left=StackTraceError(
                     name=user.name,
                     error=f"Failed to create user entity [{user}]: {exc}",
-                    stack_trace=traceback.format_exc(),
+                    stackTrace=traceback.format_exc(),
                 )
             )
 
@@ -222,7 +225,7 @@ class AmundsenSource(Source):
                     left=StackTraceError(
                         name=user.get("full_name") or "User",
                         error=f"Failed to add table from user [{user}]: {exc}",
-                        stack_trace=traceback.format_exc(),
+                        stackTrace=traceback.format_exc(),
                     )
                 )
 
@@ -257,7 +260,7 @@ class AmundsenSource(Source):
                 left=StackTraceError(
                     name="Database",
                     error=f"Failed to Ingest database due to - {err}",
-                    stack_trace=traceback.format_exc(),
+                    stackTrace=traceback.format_exc(),
                 )
             )
 
@@ -284,7 +287,7 @@ class AmundsenSource(Source):
                 left=StackTraceError(
                     name="Database Schema",
                     error=f"Failed to Ingest database schema due to - {err}",
-                    stack_trace=traceback.format_exc(),
+                    stackTrace=traceback.format_exc(),
                 )
             )
 
@@ -354,7 +357,7 @@ class AmundsenSource(Source):
                 left=StackTraceError(
                     name=table.get("name") or "Table",
                     error=f"Failed to create table entity [{table}]: {exc}",
-                    stack_trace=traceback.format_exc(),
+                    stackTrace=traceback.format_exc(),
                 )
             )
 
@@ -397,7 +400,7 @@ class AmundsenSource(Source):
                 left=StackTraceError(
                     name=dashboard.get("name") or "Dashboard",
                     error=f"Failed to create dashboard entity [{dashboard}]: {exc}",
-                    stack_trace=traceback.format_exc(),
+                    stackTrace=traceback.format_exc(),
                 )
             )
 

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

@@ -122,7 +122,7 @@ class AtlasSource(Source):
                     left=StackTraceError(
                         name=service,
                         error=f"Cannot find service for {service} - type DatabaseService",
-                        stack_trace=traceback.format_exc(),
+                        stackTrace=traceback.format_exc(),
                     )
                 )
 
@@ -142,7 +142,7 @@ class AtlasSource(Source):
                     left=StackTraceError(
                         name=service,
                         error=f"Cannot find service for {service} - type MessagingService",
-                        stack_trace=traceback.format_exc(),
+                        stackTrace=traceback.format_exc(),
                     )
                 )
 
@@ -184,7 +184,7 @@ class AtlasSource(Source):
                         left=StackTraceError(
                             name="Topic",
                             error=f"Failed to parse topi entry [{topic_entity}]: {exc}",
-                            stack_trace=traceback.format_exc(),
+                            stackTrace=traceback.format_exc(),
                         )
                     )
 
@@ -279,7 +279,7 @@ class AtlasSource(Source):
                         left=StackTraceError(
                             name="Database",
                             error=f"Failed to parse for database : {db_entity} - table {table}: {exc}",
-                            stack_trace=traceback.format_exc(),
+                            stackTrace=traceback.format_exc(),
                         )
                     )
 
@@ -420,7 +420,7 @@ class AtlasSource(Source):
                 left=StackTraceError(
                     name="Lineage",
                     error=f"failed to parse lineage due to {exc}",
-                    stack_trace=traceback.format_exc(),
+                    stackTrace=traceback.format_exc(),
                 )
             )
 

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

@@ -29,10 +29,13 @@ from metadata.generated.schema.entity.data.mlmodel import (
 from metadata.generated.schema.entity.services.connections.mlmodel.mlflowConnection import (
     MlflowConnection,
 )
+from metadata.generated.schema.entity.services.ingestionPipelines.status import (
+    StackTraceError,
+)
 from metadata.generated.schema.metadataIngestion.workflow import (
     Source as WorkflowSource,
 )
-from metadata.ingestion.api.models import Either, StackTraceError
+from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import InvalidSourceException
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.mlmodel.mlmodel_service import MlModelServiceSource
@@ -91,7 +94,7 @@ class MlflowSource(MlModelServiceSource):
                     StackTraceError(
                         name=model.name,
                         error="Version not found",
-                        stack_trace=f"Unable to ingest model {model.name} due to missing version from version list {model.latest_versions}",  # pylint: disable=line-too-long
+                        stackTrace=f"Unable to ingest model {model.name} due to missing version from version list {model.latest_versions}",  # pylint: disable=line-too-long
                     )
                 )
                 continue