Improve QueryFailedError

Author: aiguofer

dbtsl/api/adbc/client/base.py

@@ -57,7 +57,7 @@ def _handle_error(self, err: Exception) -> None:
                 raise AuthError(err.args) from err
 
             if err.status_code == AdbcStatusCode.INVALID_ARGUMENT:
-                raise QueryFailedError(err.args) from err
+                raise QueryFailedError(err.details, err.status_code) from err
 
             # TODO: timeouts are not implemented for ADBC
             # See: https://arrow.apache.org/adbc/current/driver/flight_sql.html#timeouts

dbtsl/api/graphql/client/asyncio.py

@@ -157,7 +157,7 @@ async def query(self, **params: Unpack[QueryParameters]) -> "pa.Table":
             variables={"query_id": query_id, "page_num": 1},
         )
         if first_page_results.status != QueryStatus.SUCCESSFUL:
-            raise QueryFailedError()
+            raise QueryFailedError(first_page_results.error, first_page_results.status, query_id)
 
         assert first_page_results.total_pages is not None
 

dbtsl/api/graphql/client/sync.py

@@ -148,7 +148,7 @@ def query(self, **params: Unpack[QueryParameters]) -> "pa.Table":
             },
         )
         if first_page_results.status != QueryStatus.SUCCESSFUL:
-            raise QueryFailedError()
+            raise QueryFailedError(first_page_results.error, first_page_results.status, query_id)
 
         assert first_page_results.total_pages is not None
 

dbtsl/error.py

@@ -1,4 +1,5 @@
 import json
+from typing import Any
 
 
 class SemanticLayerError(RuntimeError):
@@ -46,6 +47,24 @@ class RetryTimeoutError(TimeoutError):
 class QueryFailedError(SemanticLayerError):
     """Raise whenever a query has failed."""
 
+    def __init__(self, message: Any, status: Any, query_id: str | None = None) -> None:
+        """Initialize the query failed error.
+
+        Args:
+            message: The message or error details
+            status: The stringified status or the response
+            query_id: The query ID for GQL requests
+        """
+        self.message = str(message)
+        self.status_code = str(status)
+        self.query_id = query_id
+
+    def __str__(self) -> str:  # noqa: D105
+        content = f"message={self.message}), status_code={self.status}"
+        if self.query_id:
+            content += f", query_id={self.query_id}"
+        return f"{self.__class__.__name__}({content})"
+
 
 class AuthError(SemanticLayerError):
     """Raise whenever there was a problem authenticating to the API."""