Minor: Fix Superset API Test Connection (#14696)

This commit is contained in:
Mayur Singal 2024-01-15 15:25:48 +05:30 committed by GitHub
parent 21a701f360
commit 355f933ef0
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 19 additions and 12 deletions

View File

@ -99,6 +99,13 @@ class SupersetAPIClient:
) )
self.client = REST(client_config) 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: def fetch_total_dashboards(self) -> int:
""" """
Fetch total dashboard Fetch total dashboard
@ -107,10 +114,7 @@ class SupersetAPIClient:
int int
""" """
try: try:
resp_dashboards = self.client.get("/dashboard/?q=(page:0,page_size:1)") return self.get_dashboard_count()
if resp_dashboards:
dashboard_count = SupersetDashboardCount(**resp_dashboards)
return dashboard_count.count
except Exception: except Exception:
logger.debug(traceback.format_exc()) logger.debug(traceback.format_exc())
logger.warning("Failed to fetch the dashboard count") logger.warning("Failed to fetch the dashboard count")
@ -142,6 +146,13 @@ class SupersetAPIClient:
logger.warning("Failed to fetch the dashboard list") logger.warning("Failed to fetch the dashboard list")
return SupersetDashboardCount() 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: def fetch_total_charts(self) -> int:
""" """
Fetch the total number of charts Fetch the total number of charts
@ -149,12 +160,8 @@ class SupersetAPIClient:
Returns: Returns:
int int
""" """
try: try:
resp_chart = self.client.get("/chart/?q=(page:0,page_size:1)") return self.get_chart_count()
if resp_chart:
chart_count = SupersetChart(**resp_chart)
return chart_count.count
except Exception: except Exception:
logger.debug(traceback.format_exc()) logger.debug(traceback.format_exc())
logger.warning("Failed to fetch the chart count") logger.warning("Failed to fetch the chart count")

View File

@ -78,9 +78,9 @@ def test_connection(
test_fn = {} test_fn = {}
if isinstance(client, SupersetAPIClient): if isinstance(client, SupersetAPIClient):
test_fn["CheckAccess"] = client.fetch_total_dashboards test_fn["CheckAccess"] = client.get_dashboard_count
test_fn["GetDashboards"] = client.fetch_total_dashboards test_fn["GetDashboards"] = client.get_dashboard_count
test_fn["GetCharts"] = client.fetch_total_charts test_fn["GetCharts"] = client.get_chart_count
else: else:
test_fn["CheckAccess"] = partial(test_connection_engine_step, client) test_fn["CheckAccess"] = partial(test_connection_engine_step, client)
test_fn["GetDashboards"] = partial(test_query, client, FETCH_DASHBOARDS_TEST) test_fn["GetDashboards"] = partial(test_query, client, FETCH_DASHBOARDS_TEST)