openmetadata-ingestion 1.7.5.3__py3-none-any.whl → 1.8.1.0__py3-none-any.whl

metadata/ingestion/source/mlmodel/mlmodel_service.py

@@ -38,9 +38,6 @@ from metadata.ingestion.api.delete import delete_entity_from_source
 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.connections.test_connections import (
-    raise_test_connection_exception,
-)
 from metadata.ingestion.models.delete_entity import DeleteEntity
 from metadata.ingestion.models.topology import (
     NodeStage,
@@ -49,7 +46,7 @@ from metadata.ingestion.models.topology import (
     TopologyNode,
 )
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
-from metadata.ingestion.source.connections import get_connection, get_test_connection_fn
+from metadata.ingestion.source.connections import get_connection, test_connection_common
 from metadata.utils import fqn
 from metadata.utils.helpers import retry_with_docker_host
 from metadata.utils.logger import ingestion_logger
@@ -180,11 +177,9 @@ class MlModelServiceSource(TopologyRunnerMixin, Source, ABC):
         """By default, nothing to close"""
 
     def test_connection(self) -> None:
-        test_connection_fn = get_test_connection_fn(self.service_connection)
-        result = test_connection_fn(
+        test_connection_common(
             self.metadata, self.connection_obj, self.service_connection
         )
-        raise_test_connection_exception(result)
 
     def mark_mlmodels_as_deleted(self) -> Iterable[Either[DeleteEntity]]:
         """Method to mark the mlmodels as deleted"""

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

@@ -24,7 +24,7 @@ from metadata.generated.schema.entity.automations.workflow import (
     Workflow as AutomationWorkflow,
 )
 from metadata.generated.schema.entity.services.connections.database.mysqlConnection import (
-    MysqlConnection,
+    MysqlConnection as MysqlConnectionConfig,
 )
 from metadata.generated.schema.entity.services.connections.database.postgresConnection import (
     PostgresConnection,
@@ -67,12 +67,10 @@ def _(_: BackendConnection) -> Engine:
 
 
 @_get_connection.register
-def _(airflow_connection: MysqlConnection) -> Engine:
-    from metadata.ingestion.source.database.mysql.connection import (
-        get_connection as get_mysql_connection,
-    )
+def _(airflow_connection: MysqlConnectionConfig) -> Engine:
+    from metadata.ingestion.source.database.mysql.connection import MySQLConnection
 
-    return get_mysql_connection(airflow_connection)
+    return MySQLConnection(airflow_connection).get_client()
 
 
 @_get_connection.register

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

@@ -24,17 +24,38 @@ from metadata.generated.schema.entity.services.connections.pipeline.databricksPi
 from metadata.generated.schema.entity.services.connections.testConnectionResult import (
     TestConnectionResult,
 )
+from metadata.ingestion.connections.builders import (
+    create_generic_db_connection,
+    get_connection_args_common,
+    init_empty_connection_arguments,
+)
 from metadata.ingestion.connections.test_connections import test_connection_steps
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.database.databricks.client import DatabricksClient
 from metadata.utils.constants import THREE_MIN
 
 
+def get_connection_url(connection: DatabricksPipelineConnection) -> str:
+    url = f"databricks+connector://token:{connection.token.get_secret_value()}@{connection.hostPort}"
+    return url
+
+
 def get_connection(connection: DatabricksPipelineConnection) -> DatabricksClient:
     """
     Create connection
     """
-    return DatabricksClient(connection)
+
+    if connection.httpPath:
+        if not connection.connectionArguments:
+            connection.connectionArguments = init_empty_connection_arguments()
+        connection.connectionArguments.root["http_path"] = connection.httpPath
+
+    engine = create_generic_db_connection(
+        connection=connection,
+        get_connection_url_fn=get_connection_url,
+        get_connection_args_fn=get_connection_args_common,
+    )
+    return DatabricksClient(connection, engine)
 
 
 def test_connection(
@@ -49,7 +70,10 @@ def test_connection(
     of a metadata workflow or during an Automation Workflow
     """
 
-    test_fn = {"GetPipelines": client.list_jobs_test_connection}
+    test_fn = {
+        "GetPipelines": client.list_jobs_test_connection,
+        "GetLineage": client.test_lineage_query,
+    }
 
     return test_connection_steps(
         metadata=metadata,

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

@@ -14,7 +14,7 @@ Databricks pipeline source to extract metadata
 """
 
 import traceback
-from typing import Iterable, List, Optional
+from typing import Iterable, List, Optional, Tuple
 
 from pydantic import ValidationError
 
@@ -27,6 +27,7 @@ from metadata.generated.schema.entity.data.pipeline import (
     Task,
     TaskStatus,
 )
+from metadata.generated.schema.entity.data.table import Table
 from metadata.generated.schema.entity.services.connections.pipeline.databricksPipelineConnection import (
     DatabricksPipelineConnection,
 )
@@ -43,8 +44,16 @@ from metadata.generated.schema.type.basic import (
     SourceUrl,
     Timestamp,
 )
+from metadata.generated.schema.type.entityLineage import (
+    ColumnLineage,
+    EntitiesEdge,
+    LineageDetails,
+)
+from metadata.generated.schema.type.entityLineage import Source as LineageSource
+from metadata.generated.schema.type.entityReference import EntityReference
 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.pipeline_status import OMetaPipelineStatus
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.pipeline.databrickspipeline.models import (
@@ -54,7 +63,6 @@ from metadata.ingestion.source.pipeline.databrickspipeline.models import (
 from metadata.ingestion.source.pipeline.pipeline_service import PipelineServiceSource
 from metadata.utils import fqn
 from metadata.utils.logger import ingestion_logger
-from metadata.utils.time_utils import convert_timestamp_to_milliseconds
 
 logger = ingestion_logger()
 
@@ -200,21 +208,15 @@ class DatabrickspipelineSource(PipelineServiceSource):
                         executionStatus=STATUS_MAP.get(
                             run.state.result_state, StatusType.Failed
                         ),
-                        startTime=Timestamp(
-                            convert_timestamp_to_milliseconds(run.start_time)
-                        ),
-                        endTime=Timestamp(
-                            convert_timestamp_to_milliseconds(run.end_time)
-                        ),
+                        startTime=Timestamp(run.start_time),
+                        endTime=Timestamp(run.end_time) if run.end_time else None,
                         logLink=run.run_page_url,
                     )
                     for task in run.tasks or []
                 ]
                 pipeline_status = PipelineStatus(
                     taskStatus=task_status,
-                    timestamp=Timestamp(
-                        convert_timestamp_to_milliseconds(run.start_time)
-                    ),
+                    timestamp=Timestamp(run.start_time),
                     executionStatus=STATUS_MAP.get(
                         run.state.result_state,
                         StatusType.Failed,
@@ -235,13 +237,173 @@ class DatabrickspipelineSource(PipelineServiceSource):
         except Exception as exc:
             yield Either(
                 left=StackTraceError(
-                    name=pipeline_fqn,
+                    name=pipeline_details.job_id,
                     error=f"Failed to yield pipeline status: {exc}",
                     stackTrace=traceback.format_exc(),
                 )
             )
 
+    def _process_and_validate_column_lineage(
+        self,
+        column_lineage: List[Tuple[str, str]],
+        from_entity: Table,
+        to_entity: Table,
+    ) -> List[ColumnLineage]:
+        """
+        Process and validate column lineage
+        """
+        processed_column_lineage = []
+        if column_lineage:
+            for column_tuple in column_lineage or []:
+                try:
+                    if len(column_tuple) < 2:
+                        logger.debug(f"Skipping invalid column tuple: {column_tuple}")
+                        continue
+
+                    source_col = column_tuple[0]
+                    target_col = column_tuple[-1]
+
+                    if not source_col or not target_col:
+                        logger.debug(
+                            f"Skipping column tuple with empty values: source={source_col}, "
+                            f"target={target_col}, to_entity={to_entity.name}"
+                        )
+                        continue
+
+                    from_column = get_column_fqn(
+                        table_entity=from_entity, column=str(source_col)
+                    )
+                    to_column = get_column_fqn(
+                        table_entity=to_entity,
+                        column=str(target_col),
+                    )
+                    if from_column and to_column:
+                        processed_column_lineage.append(
+                            ColumnLineage(
+                                fromColumns=[from_column],
+                                toColumn=to_column,
+                            )
+                        )
+                except Exception as err:
+                    logger.warning(
+                        f"Error processing column lineage {column_tuple}: {err}"
+                    )
+                    logger.debug(traceback.format_exc())
+                    continue
+        if not processed_column_lineage:
+            logger.warning(
+                f"No column lineage found for {from_entity.name} to {to_entity.name}"
+            )
+        return processed_column_lineage or []
+
     def yield_pipeline_lineage_details(
         self, pipeline_details: DataBrickPipelineDetails
     ) -> Iterable[Either[AddLineageRequest]]:
-        """Get lineage between pipeline and data sources. Not Implemented."""
+        try:
+            pipeline_fqn = fqn.build(
+                metadata=self.metadata,
+                entity_type=Pipeline,
+                service_name=self.context.get().pipeline_service,
+                pipeline_name=self.context.get().pipeline,
+            )
+
+            pipeline_entity = self.metadata.get_by_name(
+                entity=Pipeline, fqn=pipeline_fqn
+            )
+
+            table_lineage_list = self.client.get_table_lineage(
+                job_id=pipeline_details.job_id
+            )
+            if table_lineage_list:
+                for table_lineage in table_lineage_list:
+                    source_table_full_name = table_lineage.get("source_table_full_name")
+                    target_table_full_name = table_lineage.get("target_table_full_name")
+                    if source_table_full_name and target_table_full_name:
+                        source = fqn.split_table_name(source_table_full_name)
+                        target = fqn.split_table_name(target_table_full_name)
+                        for dbservicename in self.get_db_service_names() or ["*"]:
+
+                            from_entity = self.metadata.get_by_name(
+                                entity=Table,
+                                fqn=fqn.build(
+                                    metadata=self.metadata,
+                                    entity_type=Table,
+                                    table_name=source.get("table"),
+                                    database_name=source.get("database"),
+                                    schema_name=source.get("database_schema"),
+                                    service_name=dbservicename,
+                                ),
+                            )
+
+                            if from_entity is None:
+                                continue
+
+                            to_entity = self.metadata.get_by_name(
+                                entity=Table,
+                                fqn=fqn.build(
+                                    metadata=self.metadata,
+                                    entity_type=Table,
+                                    table_name=target.get("table"),
+                                    database_name=target.get("database"),
+                                    schema_name=target.get("database_schema"),
+                                    service_name=dbservicename,
+                                ),
+                            )
+
+                            if to_entity is None:
+                                continue
+
+                            processed_column_lineage = (
+                                self._process_and_validate_column_lineage(
+                                    column_lineage=self.client.get_column_lineage(
+                                        job_id=pipeline_details.job_id,
+                                        TableKey=(
+                                            source_table_full_name,
+                                            target_table_full_name,
+                                        ),
+                                    ),
+                                    from_entity=from_entity,
+                                    to_entity=to_entity,
+                                )
+                            )
+
+                            lineage_details = LineageDetails(
+                                pipeline=EntityReference(
+                                    id=pipeline_entity.id.root, type="pipeline"
+                                ),
+                                source=LineageSource.PipelineLineage,
+                                columnsLineage=processed_column_lineage,
+                            )
+
+                            yield Either(
+                                right=AddLineageRequest(
+                                    edge=EntitiesEdge(
+                                        fromEntity=EntityReference(
+                                            id=from_entity.id,
+                                            type="table",
+                                        ),
+                                        toEntity=EntityReference(
+                                            id=to_entity.id,
+                                            type="table",
+                                        ),
+                                        lineageDetails=lineage_details,
+                                    )
+                                )
+                            )
+
+                else:
+                    logger.debug(
+                        f"No source or target table full name found for job {pipeline_details.job_id}"
+                    )
+            else:
+                logger.debug(
+                    f"No table lineage found for job {pipeline_details.job_id}"
+                )
+        except Exception as exc:
+            yield Either(
+                left=StackTraceError(
+                    name=pipeline_details.job_id,
+                    error=f"Wild error ingesting pipeline lineage {pipeline_details} - {exc}",
+                    stackTrace=traceback.format_exc(),
+                )
+            )

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

@@ -12,13 +12,14 @@
 Nifi source to extract metadata
 """
 import traceback
-from typing import Iterable, List, Optional
+from collections import defaultdict
+from typing import Dict, Iterable, List, Optional
 
 from pydantic import BaseModel, ValidationError
 
 from metadata.generated.schema.api.data.createPipeline import CreatePipelineRequest
 from metadata.generated.schema.api.lineage.addLineage import AddLineageRequest
-from metadata.generated.schema.entity.data.pipeline import Task
+from metadata.generated.schema.entity.data.pipeline import Pipeline, Task
 from metadata.generated.schema.entity.services.connections.pipeline.nifiConnection import (
     NifiConnection,
 )
@@ -30,6 +31,9 @@ from metadata.generated.schema.type.basic import (
     FullyQualifiedEntityName,
     SourceUrl,
 )
+from metadata.generated.schema.type.entityLineage import EntitiesEdge, LineageDetails
+from metadata.generated.schema.type.entityLineage import Source as LineageSource
+from metadata.generated.schema.type.entityReference import EntityReference
 from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import InvalidSourceException
 from metadata.ingestion.models.pipeline_status import OMetaPipelineStatus
@@ -43,6 +47,7 @@ logger = ingestion_logger()
 
 PROCESS_GROUP_FLOW = "processGroupFlow"
 BREADCRUMB = "breadcrumb"
+PARENT_BREADCRUMB = "parentBreadcrumb"
 
 
 class NifiProcessor(BaseModel):
@@ -77,6 +82,7 @@ class NifiPipelineDetails(BaseModel):
     uri: str
     processors: List[NifiProcessor]
     connections: List[NifiProcessorConnections]
+    parent_pipeline_id: Optional[str] = None
 
 
 class NifiSource(PipelineServiceSource):
@@ -85,6 +91,11 @@ class NifiSource(PipelineServiceSource):
     Pipeline metadata from Airflow's metadata db
     """
 
+    def __init__(self, config: WorkflowSource, metadata: OpenMetadata):
+        super().__init__(config, metadata)
+        self.pipeline_parents_mapping: Dict[str, List[str]] = defaultdict(list)
+        self.process_group_connections: List[NifiProcessorConnections] = []
+
     @classmethod
     def create(
         cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
@@ -220,7 +231,7 @@ class NifiSource(PipelineServiceSource):
         """Get List of all pipelines"""
         for process_group in self.connection.list_process_groups():
             try:
-                yield NifiPipelineDetails(
+                nifi_pipeline_details = NifiPipelineDetails(
                     id_=process_group[PROCESS_GROUP_FLOW].get("id"),
                     name=process_group[PROCESS_GROUP_FLOW][BREADCRUMB][BREADCRUMB].get(
                         "name"
@@ -232,7 +243,18 @@ class NifiSource(PipelineServiceSource):
                     connections=self._get_connections_from_process_group(
                         process_group=process_group
                     ),
+                    parent_pipeline_id=process_group[PROCESS_GROUP_FLOW][BREADCRUMB]
+                    .get(PARENT_BREADCRUMB, {})
+                    .get("id"),
+                )
+                if nifi_pipeline_details.parent_pipeline_id:
+                    self.pipeline_parents_mapping[nifi_pipeline_details.id_].append(
+                        nifi_pipeline_details.parent_pipeline_id
+                    )
+                self.process_group_connections.extend(
+                    self.get_process_group_connections(process_group)
                 )
+                yield nifi_pipeline_details
             except (ValueError, KeyError, ValidationError) as err:
                 logger.debug(traceback.format_exc())
                 logger.warning(
@@ -244,5 +266,113 @@ class NifiSource(PipelineServiceSource):
                     f"Wild error encountered when trying to get pipelines from Process Group {process_group} - {err}."
                 )
 
+    def get_process_group_connections(
+        self, process_group: dict
+    ) -> List[NifiProcessorConnections]:
+        """Get all connections for a process group"""
+        connections_list = (
+            process_group.get(PROCESS_GROUP_FLOW).get("flow").get("connections")
+        )
+        connections = []
+
+        for connection in connections_list:
+            try:
+                source = connection.get("component", {}).get("source", {})
+                destination = connection.get("component", {}).get("destination", {})
+                if (
+                    source.get("type") == "OUTPUT_PORT"
+                    and destination.get("type") == "INPUT_PORT"
+                    and source.get("groupId") != destination.get("groupId")
+                ):
+                    connections.append(
+                        NifiProcessorConnections(
+                            id_=connection.get("id"),
+                            source_id=source.get("groupId"),
+                            destination_id=destination.get("groupId"),
+                        )
+                    )
+            except Exception as err:
+                logger.debug(traceback.format_exc())
+                logger.warning(
+                    f"Wild error encountered when trying to get process group connections from \
+                    {process_group[PROCESS_GROUP_FLOW][BREADCRUMB][BREADCRUMB].get('name')} - {err}."
+                )
+        return connections
+
+    def yield_pipeline_bulk_lineage_details(
+        self,
+    ) -> Iterable[Either[AddLineageRequest]]:
+        """
+        Process the pipeline bulk lineage details
+        """
+        # Process the lineage between parent and child pipelines
+        for pipeline_id, parent_pipeline_ids in self.pipeline_parents_mapping.items():
+            to_entity = self.metadata.get_by_name(
+                entity=Pipeline,
+                fqn=f"{self.context.get().pipeline_service}.{pipeline_id}",
+            )
+            if not to_entity:
+                logger.warning(
+                    f"Pipeline {pipeline_id} not found in metadata, skipping lineage"
+                )
+                continue
+            for parent_pipeline_id in parent_pipeline_ids:
+                from_entity = self.metadata.get_by_name(
+                    entity=Pipeline,
+                    fqn=f"{self.context.get().pipeline_service}.{parent_pipeline_id}",
+                )
+                if not from_entity:
+                    logger.warning(
+                        f"Parent Pipeline {parent_pipeline_id} not found in metadata, skipping lineage"
+                    )
+                    continue
+                yield Either(
+                    right=AddLineageRequest(
+                        edge=EntitiesEdge(
+                            fromEntity=EntityReference(
+                                id=from_entity.id, type="pipeline"
+                            ),
+                            toEntity=EntityReference(id=to_entity.id, type="pipeline"),
+                            lineageDetails=LineageDetails(
+                                source=LineageSource.PipelineLineage
+                            ),
+                        )
+                    )
+                )
+
+        # Process the lineage between connected pipelines
+        for connection in self.process_group_connections:
+            from_entity = self.metadata.get_by_name(
+                entity=Pipeline,
+                fqn=f"{self.context.get().pipeline_service}.{connection.source_id}",
+            )
+            if not from_entity:
+                logger.warning(
+                    f"Pipeline {connection.source_id} not found in metadata, skipping lineage"
+                )
+                continue
+
+            to_entity = self.metadata.get_by_name(
+                entity=Pipeline,
+                fqn=f"{self.context.get().pipeline_service}.{connection.destination_id}",
+            )
+            if not to_entity:
+                logger.warning(
+                    f"Pipeline {connection.destination_id} not found in metadata, skipping lineage"
+                )
+                continue
+
+            yield Either(
+                right=AddLineageRequest(
+                    edge=EntitiesEdge(
+                        fromEntity=EntityReference(id=from_entity.id, type="pipeline"),
+                        toEntity=EntityReference(id=to_entity.id, type="pipeline"),
+                        lineageDetails=LineageDetails(
+                            source=LineageSource.PipelineLineage
+                        ),
+                    )
+                )
+            )
+
     def get_pipeline_name(self, pipeline_details: NifiPipelineDetails) -> str:
         return pipeline_details.name

metadata/ingestion/source/pipeline/pipeline_service.py

@@ -41,9 +41,6 @@ from metadata.ingestion.api.delete import delete_entity_from_source
 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.connections.test_connections import (
-    raise_test_connection_exception,
-)
 from metadata.ingestion.models.delete_entity import DeleteEntity
 from metadata.ingestion.models.ometa_classification import OMetaTagAndClassification
 from metadata.ingestion.models.ometa_lineage import OMetaLineageRequest
@@ -55,7 +52,7 @@ from metadata.ingestion.models.topology import (
     TopologyNode,
 )
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
-from metadata.ingestion.source.connections import get_connection, get_test_connection_fn
+from metadata.ingestion.source.connections import get_connection, test_connection_common
 from metadata.ingestion.source.pipeline.openlineage.models import TableDetails
 from metadata.ingestion.source.pipeline.openlineage.utils import FQNNotFoundException
 from metadata.utils import fqn
@@ -137,6 +134,7 @@ class PipelineServiceTopology(ServiceTopology):
                 nullable=True,
             ),
         ],
+        post_process=["process_pipeline_bulk_lineage"],
     )
 
 
@@ -316,6 +314,23 @@ class PipelineServiceSource(TopologyRunnerMixin, Source, ABC):
                 else:
                     yield lineage
 
+    def yield_pipeline_bulk_lineage_details(self) -> Iterable[AddLineageRequest]:
+        """Method to yield the bulk pipeline lineage details"""
+
+    def process_pipeline_bulk_lineage(self) -> Iterable[AddLineageRequest]:
+        """Method to process the bulk pipeline lineage"""
+        if self.source_config.includeLineage:
+            for lineage in self.yield_pipeline_bulk_lineage_details() or []:
+                if lineage.right is not None:
+                    yield Either(
+                        right=OMetaLineageRequest(
+                            lineage_request=lineage.right,
+                            override_lineage=self.source_config.overrideLineage,
+                        )
+                    )
+                else:
+                    yield lineage
+
     def _get_table_fqn_from_om(self, table_details: TableDetails) -> Optional[str]:
         """
         Based on partial schema and table names look for matching table object in open metadata.
@@ -373,11 +388,9 @@ class PipelineServiceSource(TopologyRunnerMixin, Source, ABC):
             yield pipeline_detail
 
     def test_connection(self) -> None:
-        test_connection_fn = get_test_connection_fn(self.service_connection)
-        result = test_connection_fn(
+        test_connection_common(
             self.metadata, self.connection_obj, self.service_connection
         )
-        raise_test_connection_exception(result)
 
     def register_record(self, pipeline_request: CreatePipelineRequest) -> None:
         """Mark the pipeline record as scanned and update the pipeline_source_state"""

metadata/ingestion/source/search/search_service.py

@@ -41,9 +41,6 @@ from metadata.ingestion.api.delete import delete_entity_from_source
 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.connections.test_connections import (
-    raise_test_connection_exception,
-)
 from metadata.ingestion.models.delete_entity import DeleteEntity
 from metadata.ingestion.models.search_index_data import OMetaIndexSampleData
 from metadata.ingestion.models.topology import (
@@ -53,7 +50,7 @@ from metadata.ingestion.models.topology import (
     TopologyNode,
 )
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
-from metadata.ingestion.source.connections import get_connection, get_test_connection_fn
+from metadata.ingestion.source.connections import get_connection, test_connection_common
 from metadata.utils import fqn
 from metadata.utils.filters import filter_by_search_index
 from metadata.utils.helpers import retry_with_docker_host
@@ -239,11 +236,9 @@ class SearchServiceSource(TopologyRunnerMixin, Source, ABC):
         """Nothing to prepare by default"""
 
     def test_connection(self) -> None:
-        test_connection_fn = get_test_connection_fn(self.service_connection)
-        result = test_connection_fn(
+        test_connection_common(
             self.metadata, self.connection_obj, self.service_connection
         )
-        raise_test_connection_exception(result)
 
     def mark_search_indexes_as_deleted(self) -> Iterable[Either[DeleteEntity]]:
         """Method to mark the search index as deleted"""

metadata/ingestion/source/storage/storage_service.py

@@ -40,9 +40,6 @@ from metadata.ingestion.api.delete import delete_entity_from_source
 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.connections.test_connections import (
-    raise_test_connection_exception,
-)
 from metadata.ingestion.models.delete_entity import DeleteEntity
 from metadata.ingestion.models.ometa_classification import OMetaTagAndClassification
 from metadata.ingestion.models.topology import (
@@ -52,7 +49,7 @@ from metadata.ingestion.models.topology import (
     TopologyNode,
 )
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
-from metadata.ingestion.source.connections import get_connection, get_test_connection_fn
+from metadata.ingestion.source.connections import get_connection, test_connection_common
 from metadata.ingestion.source.database.glue.models import Column
 from metadata.readers.dataframe.models import DatalakeTableSchemaWrapper
 from metadata.readers.dataframe.reader_factory import SupportedTypes
@@ -240,11 +237,9 @@ class StorageServiceSource(TopologyRunnerMixin, Source, ABC):
         self.container_source_state.add(container_fqn)
 
     def test_connection(self) -> None:
-        test_connection_fn = get_test_connection_fn(self.service_connection)
-        result = test_connection_fn(
+        test_connection_common(
             self.metadata, self.connection_obj, self.service_connection
         )
-        raise_test_connection_exception(result)
 
     def mark_containers_as_deleted(self) -> Iterable[Either[DeleteEntity]]:
         """Method to mark the containers as deleted"""