openmetadata-ingestion 1.5.7.0__py3-none-any.whl → 1.5.8.0__py3-none-any.whl

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

@@ -22,8 +22,12 @@ from metadata.generated.schema.entity.automations.workflow import (
 from metadata.generated.schema.entity.services.connections.mlmodel.mlflowConnection import (
     MlflowConnection,
 )
+from metadata.generated.schema.entity.services.connections.testConnectionResult import (
+    TestConnectionResult,
+)
 from metadata.ingestion.connections.test_connections import test_connection_steps
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.utils.constants import THREE_MIN
 
 
 def get_connection(connection: MlflowConnection) -> MlflowClient:
@@ -41,7 +45,8 @@ def test_connection(
     client: MlflowClient,
     service_connection: MlflowConnection,
     automation_workflow: Optional[AutomationWorkflow] = None,
-) -> None:
+    timeout_seconds: Optional[int] = THREE_MIN,
+) -> TestConnectionResult:
     """
     Test connection. This can be executed either as part
     of a metadata workflow or during an Automation Workflow
@@ -49,9 +54,10 @@ def test_connection(
 
     test_fn = {"GetModels": client.search_registered_models}
 
-    test_connection_steps(
+    return test_connection_steps(
         metadata=metadata,
         test_fn=test_fn,
         service_type=service_connection.type.value,
         automation_workflow=automation_workflow,
+        timeout_seconds=timeout_seconds,
     )

metadata/ingestion/source/mlmodel/mlmodel_service.py

@@ -38,6 +38,9 @@ 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,
@@ -176,7 +179,10 @@ class MlModelServiceSource(TopologyRunnerMixin, Source, ABC):
 
     def test_connection(self) -> None:
         test_connection_fn = get_test_connection_fn(self.service_connection)
-        test_connection_fn(self.metadata, self.connection_obj, self.service_connection)
+        result = test_connection_fn(
+            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/mlmodel/sagemaker/connection.py

@@ -21,8 +21,12 @@ from metadata.generated.schema.entity.automations.workflow import (
 from metadata.generated.schema.entity.services.connections.mlmodel.sageMakerConnection import (
     SageMakerConnection,
 )
+from metadata.generated.schema.entity.services.connections.testConnectionResult import (
+    TestConnectionResult,
+)
 from metadata.ingestion.connections.test_connections import test_connection_steps
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.utils.constants import THREE_MIN
 
 
 def get_connection(connection: SageMakerConnection):
@@ -37,7 +41,8 @@ def test_connection(
     client,
     service_connection: SageMakerConnection,
     automation_workflow: Optional[AutomationWorkflow] = None,
-) -> None:
+    timeout_seconds: Optional[int] = THREE_MIN,
+) -> TestConnectionResult:
     """
     Test connection. This can be executed either as part
     of a metadata workflow or during an Automation Workflow
@@ -45,9 +50,10 @@ def test_connection(
 
     test_fn = {"GetModels": client.list_models}
 
-    test_connection_steps(
+    return test_connection_steps(
         metadata=metadata,
         test_fn=test_fn,
         service_type=service_connection.type.value,
         automation_workflow=automation_workflow,
+        timeout_seconds=timeout_seconds,
     )

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

@@ -20,9 +20,13 @@ from metadata.generated.schema.entity.automations.workflow import (
 from metadata.generated.schema.entity.services.connections.pipeline.airbyteConnection import (
     AirbyteConnection,
 )
+from metadata.generated.schema.entity.services.connections.testConnectionResult import (
+    TestConnectionResult,
+)
 from metadata.ingestion.connections.test_connections import test_connection_steps
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.pipeline.airbyte.client import AirbyteClient
+from metadata.utils.constants import THREE_MIN
 
 
 def get_connection(connection: AirbyteConnection) -> AirbyteClient:
@@ -37,7 +41,8 @@ def test_connection(
     client: AirbyteClient,
     service_connection: AirbyteConnection,
     automation_workflow: Optional[AutomationWorkflow] = None,
-) -> None:
+    timeout_seconds: Optional[int] = THREE_MIN,
+) -> TestConnectionResult:
     """
     Test connection. This can be executed either as part
     of a metadata workflow or during an Automation Workflow
@@ -45,9 +50,10 @@ def test_connection(
 
     test_fn = {"GetPipelines": client.list_workspaces}
 
-    test_connection_steps(
+    return test_connection_steps(
         metadata=metadata,
         test_fn=test_fn,
         service_type=service_connection.type.value,
         automation_workflow=automation_workflow,
+        timeout_seconds=timeout_seconds,
     )

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

@@ -36,12 +36,16 @@ from metadata.generated.schema.entity.services.connections.pipeline.airflowConne
 from metadata.generated.schema.entity.services.connections.pipeline.backendConnection import (
     BackendConnection,
 )
+from metadata.generated.schema.entity.services.connections.testConnectionResult import (
+    TestConnectionResult,
+)
 from metadata.ingestion.connections.test_connections import (
     SourceConnectionException,
     test_connection_engine_step,
     test_connection_steps,
 )
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.utils.constants import THREE_MIN
 
 
 # Only import when needed
@@ -103,7 +107,8 @@ def test_connection(
     engine: Engine,
     service_connection: AirflowConnection,
     automation_workflow: Optional[AutomationWorkflow] = None,
-) -> None:
+    timeout_seconds: Optional[int] = THREE_MIN,
+) -> TestConnectionResult:
     """
     Test connection. This can be executed either as part
     of a metadata workflow or during an Automation Workflow
@@ -111,9 +116,10 @@ def test_connection(
 
     test_fn = {"CheckAccess": partial(test_connection_engine_step, engine)}
 
-    test_connection_steps(
+    return test_connection_steps(
         metadata=metadata,
         test_fn=test_fn,
         service_type=service_connection.type.value,
         automation_workflow=automation_workflow,
+        timeout_seconds=timeout_seconds,
     )

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

@@ -15,7 +15,7 @@ import traceback
 from collections import Counter
 from datetime import datetime
 from enum import Enum
-from typing import Iterable, List, Optional, cast
+from typing import Dict, Iterable, List, Optional, cast
 
 from airflow.models import BaseOperator, DagRun, TaskInstance
 from airflow.models.dag import DagModel
@@ -142,6 +142,16 @@ class AirflowSource(PipelineServiceSource):
 
         return self._session
 
+    @staticmethod
+    def _extract_serialized_task(task: Dict) -> Dict:
+        """
+        Given the serialization changes introduced in Airflow 2.10,
+        ensure compatibility with all versions.
+        """
+        if task.keys() == {"__var", "__type"}:
+            return task["__var"]
+        return task
+
     def get_pipeline_status(self, dag_id: str) -> List[DagRun]:
         """
         Return the DagRuns of given dag
@@ -328,7 +338,9 @@ class AirflowSource(PipelineServiceSource):
                     max_active_runs=data.get("max_active_runs", None),
                     description=data.get("_description", None),
                     start_date=data.get("start_date", None),
-                    tasks=data.get("tasks", []),
+                    tasks=list(
+                        map(self._extract_serialized_task, data.get("tasks", []))
+                    ),
                     schedule_interval=get_schedule_interval(data),
                     owner=self.fetch_dag_owners(data),
                 )

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

@@ -20,10 +20,14 @@ from metadata.generated.schema.entity.automations.workflow import (
 from metadata.generated.schema.entity.services.connections.pipeline.dagsterConnection import (
     DagsterConnection,
 )
+from metadata.generated.schema.entity.services.connections.testConnectionResult import (
+    TestConnectionResult,
+)
 from metadata.ingestion.connections.test_connections import test_connection_steps
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.pipeline.dagster.client import DagsterClient
 from metadata.ingestion.source.pipeline.dagster.queries import TEST_QUERY_GRAPHQL
+from metadata.utils.constants import THREE_MIN
 
 
 def get_connection(connection: DagsterConnection) -> DagsterClient:
@@ -38,7 +42,8 @@ def test_connection(
     client: DagsterClient,
     service_connection: DagsterConnection,
     automation_workflow: Optional[AutomationWorkflow] = None,
-) -> None:
+    timeout_seconds: Optional[int] = THREE_MIN,
+) -> TestConnectionResult:
     """
     Test connection. This can be executed either as part
     of a metadata workflow or during an Automation Workflow
@@ -49,9 +54,10 @@ def test_connection(
 
     test_fn = {"GetPipelines": custom_executor_for_pipeline}
 
-    test_connection_steps(
+    return test_connection_steps(
         metadata=metadata,
         test_fn=test_fn,
         service_type=service_connection.type.value,
         automation_workflow=automation_workflow,
+        timeout_seconds=timeout_seconds,
     )

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

@@ -21,9 +21,13 @@ from metadata.generated.schema.entity.automations.workflow import (
 from metadata.generated.schema.entity.services.connections.pipeline.databricksPipelineConnection import (
     DatabricksPipelineConnection,
 )
+from metadata.generated.schema.entity.services.connections.testConnectionResult import (
+    TestConnectionResult,
+)
 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(connection: DatabricksPipelineConnection) -> DatabricksClient:
@@ -38,7 +42,8 @@ def test_connection(
     client: DatabricksClient,
     service_connection: DatabricksPipelineConnection,
     automation_workflow: Optional[AutomationWorkflow] = None,
-) -> None:
+    timeout_seconds: Optional[int] = THREE_MIN,
+) -> TestConnectionResult:
     """
     Test connection. This can be executed either as part
     of a metadata workflow or during an Automation Workflow
@@ -46,9 +51,10 @@ def test_connection(
 
     test_fn = {"GetPipelines": client.list_jobs_test_connection}
 
-    test_connection_steps(
+    return test_connection_steps(
         metadata=metadata,
         test_fn=test_fn,
         service_type=service_connection.type.value,
         automation_workflow=automation_workflow,
+        timeout_seconds=timeout_seconds,
     )

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

@@ -22,9 +22,13 @@ from metadata.generated.schema.entity.automations.workflow import (
 from metadata.generated.schema.entity.services.connections.pipeline.dbtCloudConnection import (
     DBTCloudConnection,
 )
+from metadata.generated.schema.entity.services.connections.testConnectionResult import (
+    TestConnectionResult,
+)
 from metadata.ingestion.connections.test_connections import test_connection_steps
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.pipeline.dbtcloud.client import DBTCloudClient
+from metadata.utils.constants import THREE_MIN
 
 
 def get_connection(connection: DBTCloudConnection) -> DBTCloudClient:
@@ -39,7 +43,8 @@ def test_connection(
     client: DBTCloudClient,
     service_connection: DBTCloudConnection,
     automation_workflow: Optional[AutomationWorkflow] = None,
-) -> None:
+    timeout_seconds: Optional[int] = THREE_MIN,
+) -> TestConnectionResult:
     """
     Test connection. This can be executed either as part
     of a metadata workflow or during an Automation Workflow
@@ -52,9 +57,10 @@ def test_connection(
         "GetRuns": partial(client.get_runs, job_id=job_id),
     }
 
-    test_connection_steps(
+    return test_connection_steps(
         metadata=metadata,
         test_fn=test_fn,
         service_type=service_connection.type.value,
         automation_workflow=automation_workflow,
+        timeout_seconds=timeout_seconds,
     )

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

@@ -24,11 +24,15 @@ from metadata.generated.schema.entity.automations.workflow import (
 from metadata.generated.schema.entity.services.connections.pipeline.domoPipelineConnection import (
     DomoPipelineConnection,
 )
+from metadata.generated.schema.entity.services.connections.testConnectionResult import (
+    TestConnectionResult,
+)
 from metadata.ingestion.connections.test_connections import (
     SourceConnectionException,
     test_connection_steps,
 )
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.utils.constants import THREE_MIN
 
 
 def get_connection(connection: DomoPipelineConnection) -> Domo:
@@ -47,7 +51,8 @@ def test_connection(
     connection: Domo,
     service_connection: DomoPipelineConnection,
     automation_workflow: Optional[AutomationWorkflow] = None,
-) -> None:
+    timeout_seconds: Optional[int] = THREE_MIN,
+) -> TestConnectionResult:
     """
     Test connection. This can be executed either as part
     of a metadata workflow or during an Automation Workflow
@@ -59,9 +64,10 @@ def test_connection(
 
     test_fn = {"GetPipelines": custom_executor}
 
-    test_connection_steps(
+    return test_connection_steps(
         metadata=metadata,
         test_fn=test_fn,
         service_type=service_connection.type.value,
         automation_workflow=automation_workflow,
+        timeout_seconds=timeout_seconds,
     )

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

@@ -20,9 +20,13 @@ from metadata.generated.schema.entity.automations.workflow import (
 from metadata.generated.schema.entity.services.connections.pipeline.fivetranConnection import (
     FivetranConnection,
 )
+from metadata.generated.schema.entity.services.connections.testConnectionResult import (
+    TestConnectionResult,
+)
 from metadata.ingestion.connections.test_connections import test_connection_steps
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.pipeline.fivetran.client import FivetranClient
+from metadata.utils.constants import THREE_MIN
 
 
 def get_connection(connection: FivetranConnection) -> FivetranClient:
@@ -37,7 +41,8 @@ def test_connection(
     client: FivetranClient,
     service_connection: FivetranConnection,
     automation_workflow: Optional[AutomationWorkflow] = None,
-) -> None:
+    timeout_seconds: Optional[int] = THREE_MIN,
+) -> TestConnectionResult:
     """
     Test connection. This can be executed either as part
     of a metadata workflow or during an Automation Workflow
@@ -45,9 +50,10 @@ def test_connection(
 
     test_fn = {"GetPipelines": client.list_groups}
 
-    test_connection_steps(
+    return test_connection_steps(
         metadata=metadata,
         test_fn=test_fn,
         service_type=service_connection.type.value,
         automation_workflow=automation_workflow,
+        timeout_seconds=timeout_seconds,
     )

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

@@ -20,9 +20,13 @@ from metadata.generated.schema.entity.automations.workflow import (
 from metadata.generated.schema.entity.services.connections.pipeline.flinkConnection import (
     FlinkConnection,
 )
+from metadata.generated.schema.entity.services.connections.testConnectionResult import (
+    TestConnectionResult,
+)
 from metadata.ingestion.connections.test_connections import test_connection_steps
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.pipeline.flink.client import FlinkClient
+from metadata.utils.constants import THREE_MIN
 
 
 def get_connection(connection: FlinkConnection) -> FlinkClient:
@@ -37,16 +41,18 @@ def test_connection(
     client: FlinkClient,
     service_connection: FlinkConnection,
     automation_workflow: Optional[AutomationWorkflow] = None,
-) -> None:
+    timeout_seconds: Optional[int] = THREE_MIN,
+) -> TestConnectionResult:
     """
     Test connection. This can be executed either as part
     of a metadata workflow or during an Automation Workflow
     """
     test_fn = {"GetPipelines": client.get_jobs}
 
-    test_connection_steps(
+    return test_connection_steps(
         metadata=metadata,
         test_fn=test_fn,
         service_type=service_connection.type.value,
         automation_workflow=automation_workflow,
+        timeout_seconds=timeout_seconds,
     )

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

@@ -22,8 +22,12 @@ from metadata.generated.schema.entity.automations.workflow import (
 from metadata.generated.schema.entity.services.connections.pipeline.gluePipelineConnection import (
     GluePipelineConnection,
 )
+from metadata.generated.schema.entity.services.connections.testConnectionResult import (
+    TestConnectionResult,
+)
 from metadata.ingestion.connections.test_connections import test_connection_steps
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.utils.constants import THREE_MIN
 
 
 def get_connection(connection: GluePipelineConnection):
@@ -38,7 +42,8 @@ def test_connection(
     client,
     service_connection: GluePipelineConnection,
     automation_workflow: Optional[AutomationWorkflow] = None,
-) -> None:
+    timeout_seconds: Optional[int] = THREE_MIN,
+) -> TestConnectionResult:
     """
     Test connection. This can be executed either as part
     of a metadata workflow or during an Automation Workflow
@@ -46,9 +51,10 @@ def test_connection(
 
     test_fn = {"GetPipelines": client.list_workflows}
 
-    test_connection_steps(
+    return test_connection_steps(
         metadata=metadata,
         test_fn=test_fn,
         service_type=service_connection.type.value,
         automation_workflow=automation_workflow,
+        timeout_seconds=timeout_seconds,
     )

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

@@ -21,9 +21,13 @@ from metadata.generated.schema.entity.automations.workflow import (
 from metadata.generated.schema.entity.services.connections.pipeline.kafkaConnectConnection import (
     KafkaConnectConnection,
 )
+from metadata.generated.schema.entity.services.connections.testConnectionResult import (
+    TestConnectionResult,
+)
 from metadata.ingestion.connections.test_connections import test_connection_steps
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.pipeline.kafkaconnect.client import KafkaConnectClient
+from metadata.utils.constants import THREE_MIN
 
 
 def get_connection(connection: KafkaConnectConnection) -> KafkaConnectClient:
@@ -38,7 +42,8 @@ def test_connection(
     client: KafkaConnectClient,
     service_connection: KafkaConnectConnection,
     automation_workflow: Optional[AutomationWorkflow] = None,
-) -> None:
+    timeout_seconds: Optional[int] = THREE_MIN,
+) -> TestConnectionResult:
     """
     Test connection. This can be executed either as part
     of a metadata workflow or during an Automation Workflow
@@ -50,9 +55,10 @@ def test_connection(
         "GetPlugins": client.get_connector_plugins,
     }
 
-    test_connection_steps(
+    return test_connection_steps(
         metadata=metadata,
         test_fn=test_fn,
         service_type=service_connection.type.value,
         automation_workflow=automation_workflow,
+        timeout_seconds=timeout_seconds,
     )

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

@@ -21,9 +21,13 @@ from metadata.generated.schema.entity.services.connections.pipeline.nifiConnecti
     BasicAuthentication,
     NifiConnection,
 )
+from metadata.generated.schema.entity.services.connections.testConnectionResult import (
+    TestConnectionResult,
+)
 from metadata.ingestion.connections.test_connections import test_connection_steps
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.pipeline.nifi.client import NifiClient
+from metadata.utils.constants import THREE_MIN
 
 
 def get_connection(connection: NifiConnection) -> NifiClient:
@@ -53,7 +57,8 @@ def test_connection(
     client: NifiClient,
     service_connection: NifiConnection,
     automation_workflow: Optional[AutomationWorkflow] = None,
-) -> None:
+    timeout_seconds: Optional[int] = THREE_MIN,
+) -> TestConnectionResult:
     """
     Test connection. This can be executed either as part
     of a metadata workflow or during an Automation Workflow
@@ -64,9 +69,10 @@ def test_connection(
 
     test_fn = {"GetPipelines": custom_executor}
 
-    test_connection_steps(
+    return test_connection_steps(
         metadata=metadata,
         test_fn=test_fn,
         service_type=service_connection.type.value,
         automation_workflow=automation_workflow,
+        timeout_seconds=timeout_seconds,
     )

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

@@ -26,11 +26,15 @@ from metadata.generated.schema.entity.services.connections.pipeline.openLineageC
 from metadata.generated.schema.entity.services.connections.pipeline.openLineageConnection import (
     SecurityProtocol as KafkaSecProtocol,
 )
+from metadata.generated.schema.entity.services.connections.testConnectionResult import (
+    TestConnectionResult,
+)
 from metadata.ingestion.connections.test_connections import (
     SourceConnectionException,
     test_connection_steps,
 )
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.utils.constants import THREE_MIN
 
 
 def get_connection(connection: OpenLineageConnection) -> KafkaConsumer:
@@ -76,7 +80,8 @@ def test_connection(
     client: KafkaConsumer,
     service_connection: OpenLineageConnection,
     automation_workflow: Optional[AutomationWorkflow] = None,
-) -> None:
+    timeout_seconds: Optional[int] = THREE_MIN,
+) -> TestConnectionResult:
     """
     Test connection. This can be executed either as part
     of a metadata workflow or during an Automation Workflow
@@ -89,9 +94,10 @@ def test_connection(
 
     test_fn = {"GetWatermarkOffsets": custom_executor}
 
-    test_connection_steps(
+    return test_connection_steps(
         metadata=metadata,
         test_fn=test_fn,
         service_type=service_connection.type.value,
         automation_workflow=automation_workflow,
+        timeout_seconds=timeout_seconds,
     )

metadata/ingestion/source/pipeline/pipeline_service.py

@@ -35,6 +35,9 @@ 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
@@ -251,7 +254,10 @@ class PipelineServiceSource(TopologyRunnerMixin, Source, ABC):
 
     def test_connection(self) -> None:
         test_connection_fn = get_test_connection_fn(self.service_connection)
-        test_connection_fn(self.metadata, self.connection_obj, self.service_connection)
+        result = test_connection_fn(
+            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/pipeline/spline/connection.py

@@ -20,9 +20,13 @@ from metadata.generated.schema.entity.automations.workflow import (
 from metadata.generated.schema.entity.services.connections.pipeline.splineConnection import (
     SplineConnection,
 )
+from metadata.generated.schema.entity.services.connections.testConnectionResult import (
+    TestConnectionResult,
+)
 from metadata.ingestion.connections.test_connections import test_connection_steps
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.pipeline.spline.client import SplineClient
+from metadata.utils.constants import THREE_MIN
 
 
 def get_connection(connection: SplineConnection) -> SplineClient:
@@ -38,7 +42,8 @@ def test_connection(
     client: SplineClient,
     service_connection: SplineConnection,
     automation_workflow: Optional[AutomationWorkflow] = None,
-) -> None:
+    timeout_seconds: Optional[int] = THREE_MIN,
+) -> TestConnectionResult:
     """
     Test connection. This can be executed either as part
     of a metadata workflow or during an Automation Workflow
@@ -46,9 +51,10 @@ def test_connection(
 
     test_fn = {"GetPipelines": client.get_pipelines_test_connection}
 
-    test_connection_steps(
+    return test_connection_steps(
         metadata=metadata,
         test_fn=test_fn,
         service_type=service_connection.type.value,
         automation_workflow=automation_workflow,
+        timeout_seconds=timeout_seconds,
     )

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

@@ -40,10 +40,13 @@ from metadata.generated.schema.entity.services.connections.search.elasticSearch.
 from metadata.generated.schema.entity.services.connections.search.elasticSearchConnection import (
     ElasticsearchConnection,
 )
+from metadata.generated.schema.entity.services.connections.testConnectionResult import (
+    TestConnectionResult,
+)
 from metadata.ingestion.connections.builders import init_empty_connection_arguments
 from metadata.ingestion.connections.test_connections import test_connection_steps
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
-from metadata.utils.constants import UTF_8
+from metadata.utils.constants import THREE_MIN, UTF_8
 from metadata.utils.helpers import init_staging_dir
 
 CA_CERT_FILE_NAME = "root.pem"
@@ -185,7 +188,8 @@ def test_connection(
     client: Elasticsearch,
     service_connection: ElasticsearchConnection,
     automation_workflow: Optional[AutomationWorkflow] = None,
-) -> None:
+    timeout_seconds: Optional[int] = THREE_MIN,
+) -> TestConnectionResult:
     """
     Test connection. This can be executed either as part
     of a metadata workflow or during an Automation Workflow
@@ -199,9 +203,10 @@ def test_connection(
         "GetSearchIndexes": test_get_search_indexes,
     }
 
-    test_connection_steps(
+    return test_connection_steps(
         metadata=metadata,
         test_fn=test_fn,
         service_type=service_connection.type.value,
         automation_workflow=automation_workflow,
+        timeout_seconds=timeout_seconds,
     )