From 6cff9620b2e05f826a0ce4755d2f81072c63cd3a Mon Sep 17 00:00:00 2001 From: codingwithabhi <63392662+codingwithabhi@users.noreply.github.com> Date: Wed, 4 May 2022 16:26:33 +0530 Subject: [PATCH] metabase-test-connection-completed (#4671) * metabase-test-connection-completed * debug-traceback-added --- .../src/metadata/ingestion/source/metabase.py | 22 ++++----- .../src/metadata/utils/connection_clients.py | 6 +++ ingestion/src/metadata/utils/connections.py | 46 +++++++++++++++++++ 3 files changed, 62 insertions(+), 12 deletions(-) diff --git a/ingestion/src/metadata/ingestion/source/metabase.py b/ingestion/src/metadata/ingestion/source/metabase.py index 06e7228d762..3d2494e5a63 100644 --- a/ingestion/src/metadata/ingestion/source/metabase.py +++ b/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 diff --git a/ingestion/src/metadata/utils/connection_clients.py b/ingestion/src/metadata/utils/connection_clients.py index 62bb3ac29e8..74829a39cd4 100644 --- a/ingestion/src/metadata/utils/connection_clients.py +++ b/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 diff --git a/ingestion/src/metadata/utils/connections.py b/ingestion/src/metadata/utils/connections.py index 95091590214..f244c3d3146 100644 --- a/ingestion/src/metadata/utils/connections.py +++ b/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}." + )