diff --git a/ingestion/src/metadata/ingestion/source/dashboard/superset/client.py b/ingestion/src/metadata/ingestion/source/dashboard/superset/client.py index f8f654296a9..dd963a00bce 100644 --- a/ingestion/src/metadata/ingestion/source/dashboard/superset/client.py +++ b/ingestion/src/metadata/ingestion/source/dashboard/superset/client.py @@ -99,6 +99,13 @@ class SupersetAPIClient: ) 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 @@ class SupersetAPIClient: 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,6 +146,13 @@ class SupersetAPIClient: 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 @@ -149,12 +160,8 @@ class SupersetAPIClient: 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 index 545369dcf81..d1a3a5fa3b5 100644 --- 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)