Merge branch 'main' into GEN-658

ingestion/src/metadata/ingestion/connections/test_connections.py

@@ -103,7 +103,7 @@ def _test_connection_steps(
             metadata=metadata, steps=steps, automation_workflow=automation_workflow
         )
 
-    return _test_connection_steps_and_raise(steps=steps)
+    return _test_connection_steps_during_ingestion(steps=steps)
 
 
 def _test_connection_steps_automation_workflow(
@@ -231,16 +231,9 @@ def _test_connection_steps_during_ingestion(
     return test_connection_result
 
 
-def _test_connection_steps_and_raise(
-    steps: List[TestConnectionStep],
-) -> TestConnectionResult:
-    """
-    Run the test connection as part of the ingestion workflow
-    Raise an exception if something fails
-    """
-    test_connection_result = _test_connection_steps_during_ingestion(steps)
-
-    for step in test_connection_result.steps:
+def raise_test_connection_exception(result: TestConnectionResult) -> None:
+    """Raise if needed an exception for the test connection"""
+    for step in result.steps:
         if not step.passed and step.mandatory:
             raise SourceConnectionException(
                 f"Failed to run the test connection step: {step.name}"
@@ -250,8 +243,6 @@ def _test_connection_steps_and_raise(
                 f"You might be missing metadata in: {step.name} due to {step.message}"
             )
 
-    return test_connection_result
-
 
 def test_connection_steps(
     metadata: OpenMetadata,
@@ -319,7 +310,7 @@ def test_connection_db_common(
     automation_workflow: Optional[AutomationWorkflow] = None,
     queries: dict = None,
     timeout_seconds: Optional[int] = THREE_MIN,
-) -> None:
+) -> TestConnectionResult:
     """
     Test connection. This can be executed either as part
     of a metadata workflow or during an Automation Workflow
@@ -351,7 +342,7 @@ def test_connection_db_common(
     for key, query in queries.items():
         test_fn[key] = partial(test_query, statement=query, engine=engine)
 
-    test_connection_steps(
+    result = test_connection_steps(
         metadata=metadata,
         test_fn=test_fn,
         service_type=service_connection.type.value,
@@ -361,6 +352,8 @@ def test_connection_db_common(
 
     kill_active_connections(engine)
 
+    return result
+
 
 def test_connection_db_schema_sources(
     metadata: OpenMetadata,
@@ -369,7 +362,7 @@ def test_connection_db_schema_sources(
     automation_workflow: Optional[AutomationWorkflow] = None,
     queries: dict = None,
     timeout_seconds: Optional[int] = THREE_MIN,
-) -> None:
+) -> TestConnectionResult:
     """
     Test connection. This can be executed either as part
     of a metadata workflow or during an Automation Workflow
@@ -418,7 +411,7 @@ def custom_executor(engine_: Engine, inspector_fn_str: str):
     for key, query in queries.items():
         test_fn[key] = partial(test_query, statement=query, engine=engine)
 
-    test_connection_steps(
+    result = test_connection_steps(
         metadata=metadata,
         test_fn=test_fn,
         service_type=service_connection.type.value,
@@ -428,6 +421,8 @@ def custom_executor(engine_: Engine, inspector_fn_str: str):
 
     kill_active_connections(engine)
 
+    return result
+
 
 def test_query(engine: Engine, statement: str):
     """

ingestion/src/metadata/ingestion/source/api/api_service.py

@@ -39,6 +39,9 @@
 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,
@@ -175,7 +178,10 @@ def close(self):
 
     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_api_collections_as_deleted(self) -> Iterable[Either[DeleteEntity]]:
         """Method to mark the api collection as deleted"""

ingestion/src/metadata/ingestion/source/api/rest/connection.py

@@ -23,6 +23,9 @@
 from metadata.generated.schema.entity.services.connections.api.restConnection import (
     RestConnection,
 )
+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
@@ -56,7 +59,7 @@ def test_connection(
     service_connection: RestConnection,
     automation_workflow: Optional[AutomationWorkflow] = None,
     timeout_seconds: Optional[int] = THREE_MIN,
-) -> None:
+) -> TestConnectionResult:
     """
     Test connection. This can be executed either as part
     of a metadata workflow or during an Automation Workflow
@@ -81,7 +84,7 @@ def custom_schema_exec():
 
     test_fn = {"CheckURL": custom_url_exec, "CheckSchema": custom_schema_exec}
 
-    test_connection_steps(
+    return test_connection_steps(
         metadata=metadata,
         test_fn=test_fn,
         service_type=service_connection.type.value,

ingestion/src/metadata/ingestion/source/dashboard/dashboard_service.py

@@ -55,6 +55,9 @@
 from metadata.ingestion.api.models import Either, Entity
 from metadata.ingestion.api.steps import Source
 from metadata.ingestion.api.topology_runner import C, TopologyRunnerMixin
+from metadata.ingestion.connections.test_connections import (
+    raise_test_connection_exception,
+)
 from metadata.ingestion.lineage.sql_lineage import get_column_fqn
 from metadata.ingestion.models.delete_entity import DeleteEntity
 from metadata.ingestion.models.ometa_classification import OMetaTagAndClassification
@@ -554,7 +557,10 @@ def get_dashboard(self) -> Any:
 
     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 prepare(self):
         """By default, nothing to prepare"""

ingestion/src/metadata/ingestion/source/dashboard/domodashboard/connection.py

@@ -23,6 +23,9 @@
 from metadata.generated.schema.entity.services.connections.dashboard.domoDashboardConnection import (
     DomoDashboardConnection,
 )
+from metadata.generated.schema.entity.services.connections.testConnectionResult import (
+    TestConnectionResult,
+)
 from metadata.ingestion.connections.test_connections import (
     SourceConnectionException,
     test_connection_steps,
@@ -58,7 +61,7 @@ def test_connection(
     service_connection: DomoDashboardConnection,
     automation_workflow: Optional[AutomationWorkflow] = None,
     timeout_seconds: Optional[int] = THREE_MIN,
-) -> None:
+) -> TestConnectionResult:
     """
     Test connection. This can be executed either as part
     of a metadata workflow or during an Automation Workflow
@@ -73,7 +76,7 @@ def custom_test_page_list():
         "GetCharts": client.custom.test_list_cards,
     }
 
-    test_connection_steps(
+    return test_connection_steps(
         metadata=metadata,
         test_fn=test_fn,
         service_type=service_connection.type.value,

ingestion/src/metadata/ingestion/source/dashboard/lightdash/connection.py

@@ -20,6 +20,9 @@
 from metadata.generated.schema.entity.services.connections.dashboard.lightdashConnection import (
     LightdashConnection,
 )
+from metadata.generated.schema.entity.services.connections.testConnectionResult import (
+    TestConnectionResult,
+)
 from metadata.ingestion.connections.test_connections import (
     SourceConnectionException,
     test_connection_steps,
@@ -50,7 +53,7 @@ def test_connection(
     service_connection: LightdashConnection,
     automation_workflow: Optional[AutomationWorkflow] = None,
     timeout_seconds: Optional[int] = THREE_MIN,
-) -> None:
+) -> TestConnectionResult:
     """
     Test connection. This can be executed either as part
     of a metadata workflow or during an Automation Workflow
@@ -61,7 +64,7 @@ def custom_executor():
 
     test_fn = {"GetDashboards": custom_executor}
 
-    test_connection_steps(
+    return test_connection_steps(
         metadata=metadata,
         test_fn=test_fn,
         service_type=service_connection.type.value,

ingestion/src/metadata/ingestion/source/dashboard/looker/connection.py

@@ -24,6 +24,9 @@
 from metadata.generated.schema.entity.services.connections.dashboard.lookerConnection import (
     LookerConnection,
 )
+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
@@ -51,7 +54,7 @@ def test_connection(
     service_connection: LookerConnection,
     automation_workflow: Optional[AutomationWorkflow] = None,
     timeout_seconds: Optional[int] = THREE_MIN,
-) -> None:
+) -> TestConnectionResult:
     """
     Test connection. This can be executed either as part
     of a metadata workflow or during an Automation Workflow
@@ -78,7 +81,7 @@ def validate_api_version():
         "ListLookMLModels": list_datamodels_test,
     }
 
-    test_connection_steps(
+    return test_connection_steps(
         metadata=metadata,
         test_fn=test_fn,
         service_type=service_connection.type.value,

ingestion/src/metadata/ingestion/source/dashboard/metabase/connection.py

@@ -20,6 +20,9 @@
 from metadata.generated.schema.entity.services.connections.dashboard.metabaseConnection import (
     MetabaseConnection,
 )
+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.dashboard.metabase.client import MetabaseClient
@@ -39,7 +42,7 @@ def test_connection(
     service_connection: MetabaseConnection,
     automation_workflow: Optional[AutomationWorkflow] = None,
     timeout_seconds: Optional[int] = THREE_MIN,
-) -> None:
+) -> TestConnectionResult:
     """
     Test connection. This can be executed either as part
     of a metadata workflow or during an Automation Workflow
@@ -51,7 +54,7 @@ def custom_executor():
 
     test_fn = {"GetDashboards": custom_executor}
 
-    test_connection_steps(
+    return test_connection_steps(
         metadata=metadata,
         test_fn=test_fn,
         service_type=service_connection.type.value,

ingestion/src/metadata/ingestion/source/dashboard/mode/connection.py

@@ -21,6 +21,9 @@
 from metadata.generated.schema.entity.services.connections.dashboard.modeConnection import (
     ModeConnection,
 )
+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.dashboard.mode.client import ModeApiClient
@@ -40,7 +43,7 @@ def test_connection(
     service_connection: ModeConnection,
     automation_workflow: Optional[AutomationWorkflow] = None,
     timeout_seconds: Optional[int] = THREE_MIN,
-) -> None:
+) -> TestConnectionResult:
     """
     Test connection. This can be executed either as part
     of a metadata workflow or during an Automation Workflow
@@ -52,7 +55,7 @@ def test_connection(
         )
     }
 
-    test_connection_steps(
+    return test_connection_steps(
         metadata=metadata,
         test_fn=test_fn,
         service_type=service_connection.type.value,

ingestion/src/metadata/ingestion/source/dashboard/mstr/connection.py

@@ -20,6 +20,9 @@
 from metadata.generated.schema.entity.services.connections.dashboard.mstrConnection import (
     MstrConnection,
 )
+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.dashboard.mstr.client import MSTRClient
@@ -39,15 +42,15 @@ def test_connection(
     service_connection: MstrConnection,
     automation_workflow: Optional[AutomationWorkflow] = None,
     timeout_seconds: Optional[int] = THREE_MIN,
-) -> None:
+) -> TestConnectionResult:
     """
     Test connection. This can be executed either as part
     of a metadata workflow or during an Automation Workflow
     """
 
     test_fn = {"GetProjects": client.get_projects_list}
 
-    test_connection_steps(
+    return test_connection_steps(
         metadata=metadata,
         test_fn=test_fn,
         service_type=service_connection.type.value,

ingestion/src/metadata/ingestion/source/dashboard/powerbi/connection.py

@@ -20,6 +20,9 @@
 from metadata.generated.schema.entity.services.connections.dashboard.powerBIConnection import (
     PowerBIConnection,
 )
+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.dashboard.powerbi.client import (
@@ -48,14 +51,14 @@ def test_connection(
     service_connection: PowerBIConnection,
     automation_workflow: Optional[AutomationWorkflow] = None,
     timeout_seconds: Optional[int] = THREE_MIN,
-) -> None:
+) -> TestConnectionResult:
     """
     Test connection. This can be executed either as part
     of a metadata workflow or during an Automation Workflow
     """
     test_fn = {"GetDashboards": client.api_client.fetch_dashboards}
 
-    test_connection_steps(
+    return test_connection_steps(
         metadata=metadata,
         test_fn=test_fn,
         service_type=service_connection.type.value,

ingestion/src/metadata/ingestion/source/dashboard/qlikcloud/connection.py

@@ -20,6 +20,9 @@
 from metadata.generated.schema.entity.services.connections.dashboard.qlikCloudConnection import (
     QlikCloudConnection,
 )
+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.dashboard.qlikcloud.client import QlikCloudClient
@@ -39,15 +42,15 @@ def test_connection(
     service_connection: QlikCloudConnection,
     automation_workflow: Optional[AutomationWorkflow] = None,
     timeout_seconds: Optional[int] = THREE_MIN,
-) -> None:
+) -> TestConnectionResult:
     """
     Test connection. This can be executed either as part
     of a metadata workflow or during an Automation Workflow
     """
 
     test_fn = {"GetDashboards": client.get_dashboards_list_test_conn}
 
-    test_connection_steps(
+    return test_connection_steps(
         metadata=metadata,
         test_fn=test_fn,
         service_type=service_connection.type.value,