Issue 4406 superset (#4708)

* superset-test-connection-fixed

* optimized-test-connection

* code-formatted

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

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
         )

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

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}."
+        )