From da06206ea65ab4aad2e502a8902a2254bbdaa994 Mon Sep 17 00:00:00 2001 From: John Bodley <4567245+john-bodley@users.noreply.github.com> Date: Thu, 16 Nov 2023 15:42:39 -0800 Subject: [PATCH] chore: Optimize fetching samples logic (#25995) (cherry picked from commit 326ac4a6c49c49d60ac92b9722a2fd5379817c76) --- superset/views/datasource/utils.py | 23 ++++++++++------------- 1 file changed, 10 insertions(+), 13 deletions(-) diff --git a/superset/views/datasource/utils.py b/superset/views/datasource/utils.py index 65b19c34938..9baabdcc541 100644 --- a/superset/views/datasource/utils.py +++ b/superset/views/datasource/utils.py @@ -104,21 +104,18 @@ def get_samples( # pylint: disable=too-many-arguments,too-many-locals result_type=ChartDataResultType.FULL, force=force, ) - samples_results = samples_instance.get_payload() - count_star_results = count_star_instance.get_payload() try: - sample_data = samples_results["queries"][0] - count_star_data = count_star_results["queries"][0] - failed_status = ( - sample_data.get("status") == QueryStatus.FAILED - or count_star_data.get("status") == QueryStatus.FAILED - ) - error_msg = sample_data.get("error") or count_star_data.get("error") - if failed_status and error_msg: - cache_key = sample_data.get("cache_key") - QueryCacheManager.delete(cache_key, region=CacheRegion.DATA) - raise DatasetSamplesFailedError(error_msg) + count_star_data = count_star_instance.get_payload()["queries"][0] + + if count_star_data.get("status") == QueryStatus.FAILED: + raise DatasetSamplesFailedError(count_star_data.get("error")) + + sample_data = samples_instance.get_payload()["queries"][0] + + if sample_data.get("status") == QueryStatus.FAILED: + QueryCacheManager.delete(sample_data.get("cache_key"), CacheRegion.DATA) + raise DatasetSamplesFailedError(sample_data.get("error")) sample_data["page"] = page sample_data["per_page"] = per_page