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)
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")

View File

@ -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)