diff --git a/ingestion/src/metadata/ingestion/ometa/superset_rest.py b/ingestion/src/metadata/ingestion/ometa/superset_rest.py index 0a0e3e29367..83339612fae 100644 --- a/ingestion/src/metadata/ingestion/ometa/superset_rest.py +++ b/ingestion/src/metadata/ingestion/ometa/superset_rest.py @@ -32,9 +32,9 @@ class SupersetAuthenticationProvider(AuthenticationProvider): def __init__(self, config: WorkflowSource): self.config = config - self.service_connection = config.serviceConnection.__root__.config + self.service_connection = self.config client_config = ClientConfig( - base_url=config.serviceConnection.__root__.config.hostPort, + base_url=config.hostPort, api_version="api/v1", auth_token=lambda: ("no_token", 0), auth_header="Authorization", @@ -79,7 +79,7 @@ class SupersetAPIClient: self.config = config self._auth_provider = SupersetAuthenticationProvider.create(config) client_config = ClientConfig( - base_url=config.serviceConnection.__root__.config.hostPort, + base_url=config.hostPort, api_version="api/v1", auth_token=lambda: self._auth_provider.get_access_token(), auth_header="Authorization", @@ -166,3 +166,16 @@ class SupersetAPIClient: """ response = self.client.get(f"/database/{database_id}") return response + + def fetch_menu(self): + """ + Check Current User + + Args: + No Arguments + + Returns: + requests.Response + """ + response = self.client.get(f"/menu/") + return response diff --git a/ingestion/src/metadata/ingestion/source/superset.py b/ingestion/src/metadata/ingestion/source/superset.py index 4ea2a77f313..6a9ff56cd93 100644 --- a/ingestion/src/metadata/ingestion/source/superset.py +++ b/ingestion/src/metadata/ingestion/source/superset.py @@ -44,6 +44,7 @@ from metadata.ingestion.api.source import InvalidSourceException, Source, Source from metadata.ingestion.models.table_metadata import Chart, Dashboard, DashboardOwner from metadata.ingestion.ometa.ometa_api import OpenMetadata from metadata.ingestion.ometa.superset_rest import SupersetAPIClient +from metadata.utils.connections import get_connection, test_connection from metadata.utils.logger import ingestion_logger logger = ingestion_logger() @@ -184,7 +185,8 @@ class SupersetSource(Source[Entity]): self.metadata = OpenMetadata(self.metadata_config) self.status = SourceStatus() - self.client = SupersetAPIClient(self.config) + self.connection = get_connection(self.service_connection) + self.client = self.connection.client self.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 ed109c5a01f..0ba8183b93b 100644 --- a/ingestion/src/metadata/utils/connection_clients.py +++ b/ingestion/src/metadata/utils/connection_clients.py @@ -58,3 +58,9 @@ class MetabaseClient: class RedashClient: def __init__(self, client) -> None: self.client = client + + +@dataclass +class SupersetClient: + 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 a194ee328c0..c5fac5c9ff3 100644 --- a/ingestion/src/metadata/utils/connections.py +++ b/ingestion/src/metadata/utils/connections.py @@ -34,6 +34,9 @@ from metadata.generated.schema.entity.services.connections.dashboard.metabaseCon from metadata.generated.schema.entity.services.connections.dashboard.redashConnection import ( RedashConnection, ) +from metadata.generated.schema.entity.services.connections.dashboard.supersetConnection import ( + SupersetConnection, +) from metadata.generated.schema.entity.services.connections.database.bigQueryConnection import ( BigQueryConnection, ) @@ -63,6 +66,7 @@ from metadata.utils.connection_clients import ( MetabaseClient, RedashClient, SalesforceClient, + SupersetClient, ) from metadata.utils.credentials import set_google_credentials from metadata.utils.source_connections import get_connection_args, get_connection_url @@ -397,3 +401,22 @@ def _(connection: RedashClient) -> None: raise SourceConnectionException( f"Unknown error connecting with {connection} - {err}." ) + + +@get_connection.register +def _(connection: SupersetConnection, verbose: bool = False): + from metadata.ingestion.ometa.superset_rest import SupersetAPIClient + + superset_connection = SupersetAPIClient(connection) + superset_client = SupersetClient(superset_connection) + return superset_client + + +@test_connection.register +def _(connection: SupersetClient) -> None: + try: + connection.client.fetch_menu() + except Exception as err: + raise SourceConnectionException( + f"Unknown error connecting with {connection} - {err}." + )