Minor: Fix Superset API Test Connection (#14696)

open-metadata · Jan 15, 2024 · 355f933 · 355f933
1 parent 21a701f
commit 355f933
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 12 deletions.
diff --git a/ingestion/src/metadata/ingestion/source/dashboard/superset/client.py b/ingestion/src/metadata/ingestion/source/dashboard/superset/client.py
@@ -99,6 +99,13 @@ def __init__(self, config: SupersetConnection):
         )
         self.client = REST(client_config)
 
+    def get_dashboard_count(self) -> int:
+        resp_dashboards = self.client.get("/dashboard/?q=(page:0,page_size:1)")
+        if resp_dashboards:
+            dashboard_count = SupersetDashboardCount(**resp_dashboards)
+            return dashboard_count.count
+        return 0
+
     def fetch_total_dashboards(self) -> int:
         """
         Fetch total dashboard
@@ -107,10 +114,7 @@ def fetch_total_dashboards(self) -> int:
             int
         """
         try:
-            resp_dashboards = self.client.get("/dashboard/?q=(page:0,page_size:1)")
-            if resp_dashboards:
-                dashboard_count = SupersetDashboardCount(**resp_dashboards)
-                return dashboard_count.count
+            return self.get_dashboard_count()
         except Exception:
             logger.debug(traceback.format_exc())
             logger.warning("Failed to fetch the dashboard count")
@@ -142,19 +146,22 @@ def fetch_dashboards(
             logger.warning("Failed to fetch the dashboard list")
         return SupersetDashboardCount()
 
+    def get_chart_count(self) -> int:
+        resp_chart = self.client.get("/chart/?q=(page:0,page_size:1)")
+        if resp_chart:
+            chart_count = SupersetChart(**resp_chart)
+            return chart_count.count
+        return 0
+
     def fetch_total_charts(self) -> int:
         """
         Fetch the total number of charts
 
         Returns:
              int
         """
-
         try:
-            resp_chart = self.client.get("/chart/?q=(page:0,page_size:1)")
-            if resp_chart:
-                chart_count = SupersetChart(**resp_chart)
-                return chart_count.count
+            return self.get_chart_count()
         except Exception:
             logger.debug(traceback.format_exc())
             logger.warning("Failed to fetch the chart count")

diff --git a/ingestion/src/metadata/ingestion/source/dashboard/superset/connection.py b/ingestion/src/metadata/ingestion/source/dashboard/superset/connection.py
@@ -78,9 +78,9 @@ def test_connection(
     test_fn = {}
 
     if isinstance(client, SupersetAPIClient):
-        test_fn["CheckAccess"] = client.fetch_total_dashboards
-        test_fn["GetDashboards"] = client.fetch_total_dashboards
-        test_fn["GetCharts"] = client.fetch_total_charts
+        test_fn["CheckAccess"] = client.get_dashboard_count
+        test_fn["GetDashboards"] = client.get_dashboard_count
+        test_fn["GetCharts"] = client.get_chart_count
     else:
         test_fn["CheckAccess"] = partial(test_connection_engine_step, client)
         test_fn["GetDashboards"] = partial(test_query, client, FETCH_DASHBOARDS_TEST)