metabase-test-connection-completed (#4671)

* metabase-test-connection-completed

* debug-traceback-added

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

@@ -10,7 +10,6 @@
 #  limitations under the License.
 """Metabase source module"""
 
-import json
 import traceback
 import uuid
 from typing import Iterable
@@ -28,6 +27,9 @@ from metadata.generated.schema.entity.services.connections.metadata.openMetadata
     OpenMetadataConnection,
 )
 from metadata.generated.schema.entity.services.dashboardService import DashboardService
+from metadata.generated.schema.metadataIngestion.dashboardServiceMetadataPipeline import (
+    DashboardServiceMetadataPipeline,
+)
 from metadata.generated.schema.metadataIngestion.workflow import (
     Source as WorkflowSource,
 )
@@ -38,6 +40,7 @@ from metadata.ingestion.api.source import InvalidSourceException, Source, Source
 from metadata.ingestion.models.table_metadata import Chart, Dashboard
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.sql_source import SQLSourceStatus
+from metadata.utils.connections import get_connection, test_connection
 from metadata.utils.filters import filter_by_chart, filter_by_dashboard
 from metadata.utils.logger import ingestion_logger
 
@@ -74,7 +77,9 @@ class MetabaseSource(Source[Entity]):
         super().__init__()
         self.config = config
         self.service_connection = self.config.serviceConnection.__root__.config
-        self.source_config = self.config.sourceConfig.config
+        self.source_config: DashboardServiceMetadataPipeline = (
+            self.config.sourceConfig.config
+        )
         self.metadata_config = metadata_config
         self.metadata = OpenMetadata(metadata_config)
 
@@ -82,16 +87,9 @@ class MetabaseSource(Source[Entity]):
         params = dict()
         params["username"] = self.service_connection.username
         params["password"] = self.service_connection.password.get_secret_value()
-        try:
-            resp = requests.post(
-                self.service_connection.hostPort + "/api/session/",
-                data=json.dumps(params),
-                headers=HEADERS,
-            )
-        except Exception as err:
-            raise ConnectionError() from err
-        session_id = resp.json()["id"]
-        self.metabase_session = {"X-Metabase-Session": session_id}
+
+        self.connection = get_connection(self.service_connection)
+        self.metabase_session = self.connection.client["metabase_session"]
 
         self.dashboard_service = self.metadata.get_service_or_create(
             entity=DashboardService, config=config

ingestion/src/metadata/utils/connection_clients.py

@@ -40,3 +40,9 @@ class DeltaLakeClient:
 class KafkaClient:
     def __init__(self, client) -> None:
         self.client = client
+
+
+@dataclass
+class MetabaseClient:
+    def __init__(self, client) -> None:
+        self.client = client

ingestion/src/metadata/utils/connections.py

@@ -12,10 +12,13 @@
 """
 Build and document all supported Engines
 """
+import json
 import logging
+import traceback
 from functools import singledispatch
 from typing import Union
 
+import requests
 from sqlalchemy import create_engine
 from sqlalchemy.engine.base import Engine
 from sqlalchemy.exc import OperationalError
@@ -25,6 +28,9 @@ from sqlalchemy.orm.session import Session
 from metadata.generated.schema.entity.services.connections.connectionBasicType import (
     ConnectionOptions,
 )
+from metadata.generated.schema.entity.services.connections.dashboard.metabaseConnection import (
+    MetabaseConnection,
+)
 from metadata.generated.schema.entity.services.connections.database.bigQueryConnection import (
     BigQueryConnection,
 )
@@ -51,6 +57,7 @@ from metadata.utils.connection_clients import (
     DynamoClient,
     GlueClient,
     KafkaClient,
+    MetabaseClient,
     SalesforceClient,
 )
 from metadata.utils.credentials import set_google_credentials
@@ -323,3 +330,42 @@ def _(connection: DeltaLakeClient) -> None:
         raise SourceConnectionException(
             f"Unknown error connecting with {connection} - {err}."
         )
+
+
+@get_connection.register
+def _(connection: MetabaseConnection, verbose: bool = False):
+    try:
+        params = dict()
+        params["username"] = connection.username
+        params["password"] = connection.password.get_secret_value()
+
+        HEADERS = {"Content-Type": "application/json", "Accept": "*/*"}
+
+        resp = requests.post(
+            connection.hostPort + "/api/session/",
+            data=json.dumps(params),
+            headers=HEADERS,
+        )
+
+        session_id = resp.json()["id"]
+        metabase_session = {"X-Metabase-Session": session_id}
+        conn = {"connection": connection, "metabase_session": metabase_session}
+        return MetabaseClient(conn)
+
+    except Exception as err:
+        logger.error(f"Failed to connect with error :  {err}")
+        logger.debug(traceback.format_exc())
+
+
+@test_connection.register
+def _(connection: MetabaseClient) -> None:
+    try:
+        requests.get(
+            connection.client["connection"].hostPort + "/api/dashboard",
+            headers=connection.client["metabase_session"],
+        )
+
+    except Exception as err:
+        raise SourceConnectionException(
+            f"Unknown error connecting with {connection} - {err}."
+        )