openmetadata-ingestion 1.3.4.0__py3-none-any.whl → 1.4.0.0rc2__py3-none-any.whl

metadata/ingestion/source/database/trino/metadata.py

@@ -16,7 +16,7 @@ import traceback
 from copy import deepcopy
 from typing import Any, Dict, Iterable, List, Optional, Tuple
 
-from sqlalchemy import exc, inspect, sql, util
+from sqlalchemy import exc, sql, util
 from sqlalchemy.engine.base import Connection
 from sqlalchemy.sql import sqltypes
 from trino.sqlalchemy import datatype, error
@@ -183,7 +183,9 @@ class TrinoSource(CommonDbSourceService):
     ] = "JSON"
 
     @classmethod
-    def create(cls, config_dict, metadata: OpenMetadata):
+    def create(
+        cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
+    ):
         config = WorkflowSource.parse_obj(config_dict)
         connection: TrinoConnection = config.serviceConnection.__root__.config
         if not isinstance(connection, TrinoConnection):
@@ -203,7 +205,8 @@ class TrinoSource(CommonDbSourceService):
         new_service_connection = deepcopy(self.service_connection)
         new_service_connection.catalog = database_name
         self.engine = get_connection(new_service_connection)
-        self.inspector = inspect(self.engine)
+        self._connection_map = {}  # Lazy init as well
+        self._inspector_map = {}
 
     def get_database_names(self) -> Iterable[str]:
         configured_catalog = self.service_connection.catalog
@@ -218,7 +221,7 @@ class TrinoSource(CommonDbSourceService):
                     database_fqn = fqn.build(
                         self.metadata,
                         entity_type=Database,
-                        service_name=self.context.database_service,
+                        service_name=self.context.get().database_service,
                         database_name=new_catalog,
                     )
                     if filter_by_database(

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

@@ -12,6 +12,7 @@
 Trino usage module
 """
 from abc import ABC
+from typing import Optional
 
 from metadata.generated.schema.entity.services.connections.database.trinoConnection import (
     TrinoConnection,
@@ -32,7 +33,9 @@ class TrinoQueryParserSource(QueryParserSource, ABC):
     filters: str
 
     @classmethod
-    def create(cls, config_dict, metadata: OpenMetadata):
+    def create(
+        cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
+    ):
         """Create class instance"""
         config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
         connection: TrinoConnection = config.serviceConnection.__root__.config

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

@@ -72,7 +72,9 @@ class UnitycatalogLineageSource(Source):
         """
 
     @classmethod
-    def create(cls, config_dict, metadata: OpenMetadata):
+    def create(
+        cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
+    ):
         """Create class instance"""
         config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
         connection: UnityCatalogConnection = config.serviceConnection.__root__.config

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

@@ -51,11 +51,13 @@ from metadata.generated.schema.metadataIngestion.workflow import (
 )
 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.external_table_lineage_mixin import (
+    ExternalTableLineageMixin,
+)
 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
@@ -74,7 +76,9 @@ from metadata.utils.logger import ingestion_logger
 logger = ingestion_logger()
 
 
-class UnitycatalogSource(DatabaseServiceSource, MultiDBSource):
+class UnitycatalogSource(
+    ExternalTableLineageMixin, DatabaseServiceSource, MultiDBSource
+):
     """
     Implements the necessary methods to extract
     Database metadata from Databricks Source using
@@ -87,14 +91,16 @@ class UnitycatalogSource(DatabaseServiceSource, MultiDBSource):
         self.source_config: DatabaseServiceMetadataPipeline = (
             self.config.sourceConfig.config
         )
-        self.context.table_views = []
+        self.context.get_global().table_views = []
         self.metadata = metadata
         self.service_connection: UnityCatalogConnection = (
             self.config.serviceConnection.__root__.config
         )
+        self.external_location_map = {}
         self.client = get_connection(self.service_connection)
         self.connection_obj = self.client
         self.table_constraints = []
+        self.context.storage_location = None
         self.test_connection()
 
     def get_configured_database(self) -> Optional[str]:
@@ -105,7 +111,9 @@ class UnitycatalogSource(DatabaseServiceSource, MultiDBSource):
             yield catalog.name
 
     @classmethod
-    def create(cls, config_dict, metadata: OpenMetadata):
+    def create(
+        cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
+    ):
         config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
         connection: UnityCatalogConnection = config.serviceConnection.__root__.config
         if not isinstance(connection, UnityCatalogConnection):
@@ -133,7 +141,7 @@ class UnitycatalogSource(DatabaseServiceSource, MultiDBSource):
                     database_fqn = fqn.build(
                         self.metadata,
                         entity_type=Database,
-                        service_name=self.context.database_service,
+                        service_name=self.context.get().database_service,
                         database_name=catalog_name,
                     )
                     if filter_by_database(
@@ -167,7 +175,7 @@ class UnitycatalogSource(DatabaseServiceSource, MultiDBSource):
         yield Either(
             right=CreateDatabaseRequest(
                 name=database_name,
-                service=self.context.database_service,
+                service=self.context.get().database_service,
             )
         )
 
@@ -175,14 +183,14 @@ class UnitycatalogSource(DatabaseServiceSource, MultiDBSource):
         """
         return schema names
         """
-        catalog_name = self.context.database
+        catalog_name = self.context.get().database
         for schema in self.client.schemas.list(catalog_name=catalog_name):
             try:
                 schema_fqn = fqn.build(
                     self.metadata,
                     entity_type=DatabaseSchema,
-                    service_name=self.context.database_service,
-                    database_name=self.context.database,
+                    service_name=self.context.get().database_service,
+                    database_name=self.context.get().database,
                     schema_name=schema.name,
                 )
                 if filter_by_schema(
@@ -216,8 +224,8 @@ class UnitycatalogSource(DatabaseServiceSource, MultiDBSource):
                 database=fqn.build(
                     metadata=self.metadata,
                     entity_type=Database,
-                    service_name=self.context.database_service,
-                    database_name=self.context.database,
+                    service_name=self.context.get().database_service,
+                    database_name=self.context.get().database,
                 ),
             )
         )
@@ -231,8 +239,8 @@ class UnitycatalogSource(DatabaseServiceSource, MultiDBSource):
 
         :return: tables or views, depending on config
         """
-        schema_name = self.context.database_schema
-        catalog_name = self.context.database
+        schema_name = self.context.get().database_schema
+        catalog_name = self.context.get().database
         for table in self.client.tables.list(
             catalog_name=catalog_name,
             schema_name=schema_name,
@@ -242,9 +250,9 @@ class UnitycatalogSource(DatabaseServiceSource, MultiDBSource):
                 table_fqn = fqn.build(
                     self.metadata,
                     entity_type=Table,
-                    service_name=self.context.database_service,
-                    database_name=self.context.database,
-                    schema_name=self.context.database_schema,
+                    service_name=self.context.get().database_service,
+                    database_name=self.context.get().database,
+                    schema_name=self.context.get().database_schema,
                     table_name=table_name,
                 )
                 if filter_by_table(
@@ -263,7 +271,7 @@ class UnitycatalogSource(DatabaseServiceSource, MultiDBSource):
                     table_type: TableType = TableType.View
                 if table.table_type.value.lower() == TableType.External.value.lower():
                     table_type: TableType = TableType.External
-                self.context.table_data = table
+                self.context.get().table_data = table
                 yield table_name, table_type
             except Exception as exc:
                 self.status.failed(
@@ -282,9 +290,13 @@ class UnitycatalogSource(DatabaseServiceSource, MultiDBSource):
         Prepare a table request and pass it to the sink
         """
         table_name, table_type = table_name_and_type
-        table = self.client.tables.get(self.context.table_data.full_name)
-        schema_name = self.context.database_schema
-        db_name = self.context.database
+        table = self.client.tables.get(self.context.get().table_data.full_name)
+        schema_name = self.context.get().database_schema
+        db_name = self.context.get().database
+        if table.storage_location and not table.storage_location.startswith("dbfs"):
+            self.external_location_map[
+                (db_name, schema_name, table_name)
+            ] = table.storage_location
         table_constraints = None
         try:
             columns = self.get_columns(table.columns)
@@ -306,15 +318,15 @@ class UnitycatalogSource(DatabaseServiceSource, MultiDBSource):
                 databaseSchema=fqn.build(
                     metadata=self.metadata,
                     entity_type=DatabaseSchema,
-                    service_name=self.context.database_service,
-                    database_name=self.context.database,
+                    service_name=self.context.get().database_service,
+                    database_name=self.context.get().database,
                     schema_name=schema_name,
                 ),
             )
             yield Either(right=table_request)
 
             if table_type == TableType.View or table.view_definition:
-                self.context.table_views.append(
+                self.context.get_global().table_views.append(
                     TableView(
                         table_name=table_name,
                         schema_name=schema_name,
@@ -375,18 +387,17 @@ class UnitycatalogSource(DatabaseServiceSource, MultiDBSource):
             ref_table_fqn = column.parent_table
             table_fqn_list = fqn.split(ref_table_fqn)
 
-            referred_table = fqn.search_table_from_es(
-                metadata=self.metadata,
+            referred_table_fqn = fqn.build(
+                self.metadata,
+                entity_type=Table,
                 table_name=table_fqn_list[2],
                 schema_name=table_fqn_list[1],
                 database_name=table_fqn_list[0],
-                service_name=self.context.database_service,
+                service_name=self.context.get().database_service,
             )
-            if referred_table:
+            if referred_table_fqn:
                 for parent_column in column.parent_columns:
-                    col_fqn = get_column_fqn(
-                        table_entity=referred_table, column=parent_column
-                    )
+                    col_fqn = fqn._build(referred_table_fqn, parent_column, quote=False)
                     if col_fqn:
                         referred_column_fqns.append(col_fqn)
             else:
@@ -485,12 +496,12 @@ class UnitycatalogSource(DatabaseServiceSource, MultiDBSource):
     def yield_view_lineage(self) -> Iterable[Either[AddLineageRequest]]:
         logger.info("Processing Lineage for Views")
         for view in [
-            v for v in self.context.table_views if v.view_definition is not None
+            v for v in self.context.get().table_views if v.view_definition is not None
         ]:
             yield from get_view_lineage(
                 view=view,
                 metadata=self.metadata,
-                service_name=self.context.database_service,
+                service_name=self.context.get().database_service,
                 connection_type=self.service_connection.type.value,
             )
 

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

@@ -12,6 +12,7 @@
 UnityCatalog Query parser module
 """
 from abc import ABC
+from typing import Optional
 
 from metadata.generated.schema.entity.services.connections.database.unityCatalogConnection import (
     UnityCatalogConnection,
@@ -49,7 +50,9 @@ class UnityCatalogQueryParserSource(
         self.client = UnityCatalogClient(self.service_connection)
 
     @classmethod
-    def create(cls, config_dict, metadata: OpenMetadata):
+    def create(
+        cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
+    ):
         """Create class instance"""
         config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
         connection: UnityCatalogConnection = config.serviceConnection.__root__.config

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

@@ -274,7 +274,9 @@ class VerticaSource(CommonDbSourceService, MultiDBSource):
         self.schema_desc_map = {}
 
     @classmethod
-    def create(cls, config_dict, metadata: OpenMetadata):
+    def create(
+        cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
+    ):
         config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
         connection: VerticaConnection = config.serviceConnection.__root__.config
         if not isinstance(connection, VerticaConnection):
@@ -311,7 +313,7 @@ class VerticaSource(CommonDbSourceService, MultiDBSource):
                 database_fqn = fqn.build(
                     self.metadata,
                     entity_type=Database,
-                    service_name=self.context.database_service,
+                    service_name=self.context.get().database_service,
                     database_name=new_database,
                 )
 

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

@@ -12,7 +12,7 @@
 Vertica usage module
 """
 from abc import ABC
-from typing import Iterable
+from typing import Iterable, Optional
 
 from metadata.generated.schema.entity.services.connections.database.verticaConnection import (
     VerticaConnection,
@@ -43,7 +43,9 @@ class VerticaQueryParserSource(QueryParserSource, ABC):
     filters: str
 
     @classmethod
-    def create(cls, config_dict, metadata: OpenMetadata):
+    def create(
+        cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
+    ):
         """Create class instance"""
         config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
         connection: VerticaConnection = config.serviceConnection.__root__.config

metadata/ingestion/source/messaging/common_broker_source.py

@@ -47,6 +47,7 @@ from metadata.parsers.schema_parsers import (
 )
 from metadata.utils import fqn
 from metadata.utils.logger import ingestion_logger
+from metadata.utils.messaging_utils import merge_and_clean_protobuf_schema
 
 logger = ingestion_logger()
 
@@ -75,6 +76,7 @@ class CommonBrokerSource(MessagingServiceSource, ABC):
         self.generate_sample_data = self.config.sourceConfig.config.generateSampleData
         self.admin_client = self.connection.admin_client
         self.schema_registry_client = self.connection.schema_registry_client
+        self.context.processed_schemas = {}
         if self.generate_sample_data:
             self.consumer_client = self.connection.consumer_client
 
@@ -104,7 +106,7 @@ class CommonBrokerSource(MessagingServiceSource, ABC):
             logger.info(f"Fetching topic config {topic_details.topic_name}")
             topic = CreateTopicRequest(
                 name=topic_details.topic_name,
-                service=self.context.messaging_service,
+                service=self.context.get().messaging_service,
                 partitions=len(topic_details.topic_metadata.partitions),
                 replicationFactor=len(
                     topic_details.topic_metadata.partitions.get(0).replicas
@@ -125,9 +127,14 @@ class CommonBrokerSource(MessagingServiceSource, ABC):
                     raise InvalidSchemaTypeException(
                         f"Cannot find {schema_type} in parser providers registry."
                     )
-                schema_fields = load_parser_fn(
-                    topic_details.topic_name, topic_schema.schema_str
-                )
+                schema_text = topic_schema.schema_str
+
+                # In protobuf schema, we need to merge all the schema text with references
+                if schema_type == SchemaType.Protobuf.value.lower():
+                    schema_text = merge_and_clean_protobuf_schema(
+                        self._get_schema_text_with_references(schema=topic_schema)
+                    )
+                schema_fields = load_parser_fn(topic_details.topic_name, schema_text)
 
                 topic.messageSchema = Topic(
                     schemaText=topic_schema.schema_str,
@@ -195,6 +202,38 @@ class CommonBrokerSource(MessagingServiceSource, ABC):
                 f"Exception adding properties to topic [{topic.name}]: {exc}"
             )
 
+    def _get_schema_text_with_references(self, schema) -> Optional[str]:
+        """
+        Returns the schema text with references resolved using recursive calls
+        """
+        try:
+            if schema:
+                schema_text = schema.schema_str
+                for reference in schema.references or []:
+                    if not self.context.processed_schemas.get(reference.name):
+                        self.context.processed_schemas[reference.name] = True
+                        reference_schema = (
+                            self.schema_registry_client.get_latest_version(
+                                reference.name
+                            )
+                        )
+                        if reference_schema.schema.references:
+                            schema_text = (
+                                schema_text
+                                + self._get_schema_text_with_references(
+                                    reference_schema.schema
+                                )
+                            )
+                        else:
+                            schema_text = (
+                                schema_text + reference_schema.schema.schema_str
+                            )
+                return schema_text
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.warning(f"Failed to get schema with references: {exc}")
+        return None
+
     def _parse_topic_metadata(self, topic_name: str) -> Optional[Schema]:
         try:
             if self.schema_registry_client:
@@ -219,8 +258,8 @@ class CommonBrokerSource(MessagingServiceSource, ABC):
         topic_fqn = fqn.build(
             metadata=self.metadata,
             entity_type=Topic,
-            service_name=self.context.messaging_service,
-            topic_name=self.context.topic,
+            service_name=self.context.get().messaging_service,
+            topic_name=self.context.get().topic,
         )
         topic_entity = self.metadata.get_by_name(entity=TopicEntity, fqn=topic_fqn)
         if topic_entity and self.generate_sample_data:

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

@@ -11,6 +11,8 @@
 """
 Kafka source ingestion
 """
+from typing import Optional
+
 from metadata.generated.schema.entity.services.connections.messaging.kafkaConnection import (
     KafkaConnection,
 )
@@ -20,11 +22,29 @@ from metadata.generated.schema.metadataIngestion.workflow import (
 from metadata.ingestion.api.steps import InvalidSourceException
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.messaging.common_broker_source import CommonBrokerSource
+from metadata.utils.ssl_manager import SSLManager
 
 
 class KafkaSource(CommonBrokerSource):
+    def __init__(self, config: WorkflowSource, metadata: OpenMetadata):
+        self.ssl_manager = None
+        service_connection = config.serviceConnection.__root__.config
+        if service_connection.schemaRegistrySSL:
+
+            self.ssl_manager = SSLManager(
+                ca=service_connection.schemaRegistrySSL.__root__.caCertificate,
+                key=service_connection.schemaRegistrySSL.__root__.sslKey,
+                cert=service_connection.schemaRegistrySSL.__root__.sslCertificate,
+            )
+            service_connection = self.ssl_manager.setup_ssl(
+                config.serviceConnection.__root__.config.sslConfig
+            )
+        super().__init__(config, metadata)
+
     @classmethod
-    def create(cls, config_dict, metadata: OpenMetadata):
+    def create(
+        cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
+    ):
         config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
         connection: KafkaConnection = config.serviceConnection.__root__.config
         if not isinstance(connection, KafkaConnection):

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

@@ -62,13 +62,19 @@ class KinesisSource(MessagingServiceSource):
     topics metadata from Kinesis Source
     """
 
-    def __init__(self, config: WorkflowSource, metadata: OpenMetadata):
+    def __init__(
+        self,
+        config: WorkflowSource,
+        metadata: OpenMetadata,
+    ):
         super().__init__(config, metadata)
         self.generate_sample_data = self.config.sourceConfig.config.generateSampleData
         self.kinesis = self.connection
 
     @classmethod
-    def create(cls, config_dict, metadata: OpenMetadata):
+    def create(
+        cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
+    ):
         config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
         connection: KinesisConnection = config.serviceConnection.__root__.config
         if not isinstance(connection, KinesisConnection):
@@ -124,7 +130,7 @@ class KinesisSource(MessagingServiceSource):
 
             topic = CreateTopicRequest(
                 name=topic_details.topic_name,
-                service=self.context.messaging_service,
+                service=self.context.get().messaging_service,
                 partitions=len(topic_details.topic_metadata.partitions),
                 retentionTime=self._compute_retention_time(
                     topic_details.topic_metadata.summary
@@ -199,8 +205,8 @@ class KinesisSource(MessagingServiceSource):
             topic_fqn = fqn.build(
                 metadata=self.metadata,
                 entity_type=Topic,
-                service_name=self.context.messaging_service,
-                topic_name=self.context.topic,
+                service_name=self.context.get().messaging_service,
+                topic_name=self.context.get().topic,
             )
             topic_entity = self.metadata.get_by_name(entity=Topic, fqn=topic_fqn)
             if topic_entity and self.generate_sample_data:

metadata/ingestion/source/messaging/messaging_service.py

@@ -37,7 +37,7 @@ from metadata.ingestion.models.delete_entity import DeleteEntity
 from metadata.ingestion.models.topology import (
     NodeStage,
     ServiceTopology,
-    TopologyContext,
+    TopologyContextManager,
     TopologyNode,
 )
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
@@ -114,7 +114,7 @@ class MessagingServiceSource(TopologyRunnerMixin, Source, ABC):
     service_connection: MessagingConnection.__fields__["config"].type_
 
     topology = MessagingServiceTopology()
-    context = TopologyContext.create(topology)
+    context = TopologyContextManager(topology)
     topic_source_state: Set = set()
 
     def __init__(
@@ -203,7 +203,7 @@ class MessagingServiceSource(TopologyRunnerMixin, Source, ABC):
                 entity_type=Topic,
                 entity_source_state=self.topic_source_state,
                 mark_deleted_entity=self.source_config.markDeletedTopics,
-                params={"service": self.context.messaging_service},
+                params={"service": self.context.get().messaging_service},
             )
 
     def register_record(self, topic_request: CreateTopicRequest) -> None:

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

@@ -11,6 +11,8 @@
 """
 RedPanda source ingestion
 """
+from typing import Optional
+
 from metadata.generated.schema.entity.services.connections.messaging.redpandaConnection import (
     RedpandaConnection,
 )
@@ -24,7 +26,9 @@ from metadata.ingestion.source.messaging.common_broker_source import CommonBroke
 
 class RedpandaSource(CommonBrokerSource):
     @classmethod
-    def create(cls, config_dict, metadata: OpenMetadata):
+    def create(
+        cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
+    ):
         config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
         connection: RedpandaConnection = config.serviceConnection.__root__.config
         if not isinstance(connection, RedpandaConnection):

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

@@ -126,7 +126,9 @@ class AmundsenSource(Source):
         self.test_connection()
 
     @classmethod
-    def create(cls, config_dict, metadata: OpenMetadata):
+    def create(
+        cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
+    ):
         """Create class instance"""
         config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
         connection: AmundsenConnection = config.serviceConnection.__root__.config

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

@@ -93,7 +93,9 @@ class AtlasSource(Source):
         self.test_connection()
 
     @classmethod
-    def create(cls, config_dict, metadata: OpenMetadata):
+    def create(
+        cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
+    ):
         config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
         connection: AtlasConnection = config.serviceConnection.__root__.config
         if not isinstance(connection, AtlasConnection):

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

@@ -55,7 +55,9 @@ class MlflowSource(MlModelServiceSource):
     """
 
     @classmethod
-    def create(cls, config_dict, metadata: OpenMetadata):
+    def create(
+        cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
+    ):
         config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
         connection: MlflowConnection = config.serviceConnection.__root__.config
         if not isinstance(connection, MlflowConnection):
@@ -126,7 +128,7 @@ class MlflowSource(MlModelServiceSource):
                 run.data, latest_version.run_id, model.name
             ),
             mlStore=self._get_ml_store(latest_version),
-            service=self.context.mlmodel_service,
+            service=self.context.get().mlmodel_service,
             sourceUrl=source_url,
         )
         yield Either(right=mlmodel_request)

metadata/ingestion/source/mlmodel/mlmodel_service.py

@@ -39,7 +39,7 @@ from metadata.ingestion.models.delete_entity import DeleteEntity
 from metadata.ingestion.models.topology import (
     NodeStage,
     ServiceTopology,
-    TopologyContext,
+    TopologyContextManager,
     TopologyNode,
 )
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
@@ -100,7 +100,7 @@ class MlModelServiceSource(TopologyRunnerMixin, Source, ABC):
     service_connection: MlModelConnection.__fields__["config"].type_
 
     topology = MlModelServiceTopology()
-    context = TopologyContext.create(topology)
+    context = TopologyContextManager(topology)
     mlmodel_source_state: Set = set()
 
     def __init__(
@@ -179,7 +179,7 @@ class MlModelServiceSource(TopologyRunnerMixin, Source, ABC):
                 entity_type=MlModel,
                 entity_source_state=self.mlmodel_source_state,
                 mark_deleted_entity=self.source_config.markDeletedMlModels,
-                params={"service": self.context.mlmodel_service},
+                params={"service": self.context.get().mlmodel_service},
             )
 
     def register_record(self, mlmodel_request: CreateMlModelRequest) -> None:

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

@@ -67,7 +67,9 @@ class SagemakerSource(MlModelServiceSource):
         self.sagemaker = self.client
 
     @classmethod
-    def create(cls, config_dict, metadata: OpenMetadata):
+    def create(
+        cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
+    ):
         config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
         connection: SageMakerConnection = config.serviceConnection.__root__.config
         if not isinstance(connection, SageMakerConnection):
@@ -136,7 +138,7 @@ class SagemakerSource(MlModelServiceSource):
                 name=model.name,
                 algorithm=self._get_algorithm(),  # Setting this to a constant
                 mlStore=self._get_ml_store(model.name),
-                service=self.context.mlmodel_service,
+                service=self.context.get().mlmodel_service,
             )
             yield Either(right=mlmodel_request)
             self.register_record(mlmodel_request=mlmodel_request)
@@ -158,7 +160,10 @@ class SagemakerSource(MlModelServiceSource):
         """
         try:
             model_info = self.sagemaker.describe_model(ModelName=model_name)
-            return MlStore(imageRepository=model_info["PrimaryContainer"]["Image"])
+            storage = model_info.get("PrimaryContainer", {}).get("ModelDataUrl")
+            image_repository = model_info.get("PrimaryContainer", {}).get("Image")
+            if image_repository or storage:
+                return MlStore(storage=storage, imageRepository=image_repository)
         except ValidationError as err:
             logger.debug(traceback.format_exc())
             logger.warning(