From 34cbdf9c3fd38453a06df0cba5b51d0bfdcec3f0 Mon Sep 17 00:00:00 2001
From: Yongjie Zhao <yongjie.zhao@gmail.com>
Date: Wed, 1 Jun 2022 10:06:54 +0800
Subject: [PATCH] fix: failed samples should throw exception

---
 .../DataTablesPane/components/SamplesPane.tsx |  8 ++++----
 superset/datasets/api.py                      |  8 +-------
 superset/datasets/commands/samples.py         |  7 ++++++-
 tests/integration_tests/datasets/api_tests.py | 20 +++++++++++++++++++
 4 files changed, 31 insertions(+), 12 deletions(-)

diff --git a/superset-frontend/src/explore/components/DataTablesPane/components/SamplesPane.tsx b/superset-frontend/src/explore/components/DataTablesPane/components/SamplesPane.tsx
index de3ef919f6b6c..769e217692771 100644
--- a/superset-frontend/src/explore/components/DataTablesPane/components/SamplesPane.tsx
+++ b/superset-frontend/src/explore/components/DataTablesPane/components/SamplesPane.tsx
@@ -17,7 +17,7 @@
  * under the License.
  */
 import React, { useState, useEffect, useMemo } from 'react';
-import { GenericDataType, styled, t } from '@superset-ui/core';
+import { ensureIsArray, GenericDataType, styled, t } from '@superset-ui/core';
 import Loading from 'src/components/Loading';
 import { EmptyStateMedium } from 'src/components/EmptyState';
 import TableView, { EmptyWrapperType } from 'src/components/TableView';
@@ -63,9 +63,9 @@ export const SamplesPane = ({
       setIsLoading(true);
       getDatasetSamples(datasource.id, queryForce)
         .then(response => {
-          setData(response.data);
-          setColnames(response.colnames);
-          setColtypes(response.coltypes);
+          setData(ensureIsArray(response.data));
+          setColnames(ensureIsArray(response.colnames));
+          setColtypes(ensureIsArray(response.coltypes));
           setResponseError('');
           cache.add(datasource);
           if (queryForce && actions) {
diff --git a/superset/datasets/api.py b/superset/datasets/api.py
index 313634766c001..17e99959e9675 100644
--- a/superset/datasets/api.py
+++ b/superset/datasets/api.py
@@ -825,10 +825,4 @@ def samples(self, pk: int) -> Response:
         except DatasetForbiddenError:
             return self.response_403()
         except DatasetSamplesFailedError as ex:
-            logger.error(
-                "Error get dataset samples %s: %s",
-                self.__class__.__name__,
-                str(ex),
-                exc_info=True,
-            )
-            return self.response_422(message=str(ex))
+            return self.response_400(message=str(ex))
diff --git a/superset/datasets/commands/samples.py b/superset/datasets/commands/samples.py
index 4be2c6e90f850..2d49a3b3daf64 100644
--- a/superset/datasets/commands/samples.py
+++ b/superset/datasets/commands/samples.py
@@ -30,6 +30,7 @@
 )
 from superset.datasets.dao import DatasetDAO
 from superset.exceptions import SupersetSecurityException
+from superset.utils.core import QueryStatus
 from superset.views.base import check_ownership
 
 logger = logging.getLogger(__name__)
@@ -58,7 +59,11 @@ def run(self) -> Dict[str, Any]:
         )
         results = qc_instance.get_payload()
         try:
-            return results["queries"][0]
+            sample_data = results["queries"][0]
+            error_msg = sample_data.get("error")
+            if sample_data.get("status") == QueryStatus.FAILED and error_msg:
+                raise DatasetSamplesFailedError(error_msg)
+            return sample_data
         except (IndexError, KeyError) as exc:
             raise DatasetSamplesFailedError from exc
 
diff --git a/tests/integration_tests/datasets/api_tests.py b/tests/integration_tests/datasets/api_tests.py
index b426ddad71523..88458c8740b76 100644
--- a/tests/integration_tests/datasets/api_tests.py
+++ b/tests/integration_tests/datasets/api_tests.py
@@ -1903,3 +1903,23 @@ def test_get_dataset_samples(self):
             f' limit {self.app.config["SAMPLES_ROW_LIMIT"]}'
         ).to_dict(orient="records")
         assert eager_samples == rv_data2["result"]["data"]
+
+    @pytest.mark.usefixtures("create_datasets")
+    def test_get_dataset_samples_with_failed_cc(self):
+        dataset = self.get_fixture_datasets()[0]
+
+        self.login(username="admin")
+        uri = f"api/v1/dataset/{dataset.id}/samples"
+        failed_column = TableColumn(
+            column_name="DUMMY CC",
+            type="VARCHAR(255)",
+            table=dataset,
+            expression="INCORRECT SQL",
+        )
+        dataset.columns.append(failed_column)
+        rv = self.client.get(uri)
+        assert rv.status_code == 400
+        rv_data = json.loads(rv.data)
+        assert "message" in rv_data
+        if dataset.database.db_engine_spec.engine_name == "PostgreSQL":
+            assert "INCORRECT SQL" in rv_data.get("message")