From ebf9a301eb8945c5e9fbbe6859595f7838bd6ad2 Mon Sep 17 00:00:00 2001
From: Josh Temple <8672171+joshtemple@users.noreply.github.com>
Date: Thu, 25 May 2023 11:18:20 -0400
Subject: [PATCH] Allow Looker API errors with a single, string-typed error

---
 spectacles/types.py              | 25 ++++++++++++++++++++-----
 spectacles/validators/sql.py     | 26 ++++++++++++++++++--------
 tests/unit/test_sql_validator.py |  2 +-
 tests/unit/test_types.py         | 18 ++++++++++++++++++
 4 files changed, 57 insertions(+), 14 deletions(-)

diff --git a/spectacles/types.py b/spectacles/types.py
index d490a178..c4576369 100644
--- a/spectacles/types.py
+++ b/spectacles/types.py
@@ -52,14 +52,20 @@ def runtime(self) -> float:
 
 
 class ErrorQueryResult(BaseModel):
-    class QueryResultData(BaseModel):
+    class ErrorData(BaseModel):
+        id: str
+        error: str
+        runtime = 0.0
+        sql = ""
+
+    class MultiErrorData(BaseModel):
         id: str
         runtime: float
         sql: str
         errors: Optional[Tuple[QueryError, ...]]
 
     status: Literal["error"]
-    data: QueryResultData
+    data: Union[ErrorData, MultiErrorData]
 
     @property
     def runtime(self) -> float:
@@ -71,9 +77,18 @@ def sql(self) -> str:
 
     @property
     def errors(self) -> Tuple[QueryError, ...]:
-        if self.data.errors is None:
-            raise TypeError("No errors contained in this query result")
-        return self.data.errors
+        if isinstance(self.data, self.ErrorData):
+            return (
+                QueryError(
+                    message=self.data.error, message_details=None, sql_error_loc=None
+                ),
+            )
+        elif isinstance(self.data, self.MultiErrorData):
+            if self.data.errors is None:
+                raise TypeError("No errors contained in this query result")
+            return self.data.errors
+        else:
+            raise TypeError("Unexpected type for ErrorQueryResult.data")
 
     def get_valid_errors(self) -> Tuple[QueryError, ...]:
         WARNINGS = (
diff --git a/spectacles/validators/sql.py b/spectacles/validators/sql.py
index fdc1556b..f4501efd 100644
--- a/spectacles/validators/sql.py
+++ b/spectacles/validators/sql.py
@@ -3,7 +3,7 @@
 from dataclasses import dataclass
 import time
 from tabulate import tabulate
-from typing import List, Optional, Tuple, Iterator, Union, cast
+from typing import List, Optional, Tuple, Iterator
 import pydantic
 from spectacles.client import LookerClient
 from spectacles.lookml import CompiledSql, Dimension, Explore
@@ -11,7 +11,12 @@
 from spectacles.logger import GLOBAL_LOGGER as logger
 from spectacles.printer import print_header
 from spectacles.utils import consume_queue, halt_queue
-from spectacles.types import QueryResult, CompletedQueryResult, ErrorQueryResult
+from spectacles.types import (
+    QueryResult,
+    CompletedQueryResult,
+    ErrorQueryResult,
+    InterruptedQueryResult,
+)
 
 QUERY_TASK_LIMIT = 250
 DEFAULT_CHUNK_SIZE = 500
@@ -316,10 +321,9 @@ async def _get_query_results(
                     )
 
                     # Append long-running queries for the profiler
-                    if query_result.status in ("complete", "error"):
-                        query_result = cast(
-                            Union[CompletedQueryResult, ErrorQueryResult], query_result
-                        )
+                    if isinstance(
+                        query_result, (CompletedQueryResult, ErrorQueryResult)
+                    ):
                         query = self._task_to_query[task_id]
                         query.runtime = query_result.runtime
                         if query_result.runtime > self.runtime_threshold:
@@ -393,7 +397,10 @@ async def _get_query_results(
                             # Indicate there are no more queries or subqueries to run
                             queries_to_run.task_done()
 
-                    elif query_result.status == "killed":
+                    elif (
+                        isinstance(query_result, InterruptedQueryResult)
+                        and query_result.status == "killed"
+                    ):
                         query = self._task_to_query[task_id]
                         query.errored = True
                         explore = query.explore
@@ -416,7 +423,10 @@ async def _get_query_results(
                         query_slot.release()
                         queries_to_run.task_done()
 
-                    elif query_result.status == "expired":
+                    elif (
+                        isinstance(query_result, InterruptedQueryResult)
+                        and query_result.status == "expired"
+                    ):
                         query = self._task_to_query[task_id]
                         query.expired_at = query.expired_at or time.time()
                         expired_for = time.time() - query.expired_at
diff --git a/tests/unit/test_sql_validator.py b/tests/unit/test_sql_validator.py
index 7afec30f..c98eb5cc 100644
--- a/tests/unit/test_sql_validator.py
+++ b/tests/unit/test_sql_validator.py
@@ -439,7 +439,7 @@ async def test_get_query_results_handles_bogus_expired_queries(
             json={
                 query_task_id: ErrorQueryResult(
                     status="error",
-                    data=ErrorQueryResult.QueryResultData(
+                    data=ErrorQueryResult.MultiErrorData(
                         id="",
                         runtime=1.0,
                         sql="",
diff --git a/tests/unit/test_types.py b/tests/unit/test_types.py
index 06501d24..346072d9 100644
--- a/tests/unit/test_types.py
+++ b/tests/unit/test_types.py
@@ -129,3 +129,21 @@ def test_get_valid_errors_should_ignore_warnings():
     query_result = cast(ErrorQueryResult, QueryResult.parse_obj(response_json).__root__)
     valid_errors = query_result.get_valid_errors()
     assert not valid_errors
+
+
+def test_can_parse_string_errors():
+    response = {
+        "status": "error",
+        "result_source": None,
+        "data": {
+            "from_cache": True,
+            "id": "67120bd3c7d23eb81f72692be9581c4a",
+            "error": "View Not Found",
+        },
+    }
+
+    result = QueryResult.parse_obj(response).__root__
+    assert isinstance(result, ErrorQueryResult)
+    assert result.errors[0].message == "View Not Found"
+    assert result.runtime == 0.0
+    assert result.sql == ""