fix(querying): Use appropriate HTTP status for async queries (#21175)

* fix(querying): Use appropriate HTTP status for async queries

* Make `error_message` nullable

* Update test_query.py

* Satisfy mypy

* Satisfy mypy more

* Update UI snapshots for `chromium` (2)

* Update UI snapshots for `chromium` (2)

---------

Co-authored-by: github-actions <41898282+github-actions[bot]@users.noreply.github.com>

frontend/src/queries/schema.json

@@ -4469,8 +4469,8 @@
                     "type": "boolean"
                 },
                 "error_message": {
-                    "default": "",
-                    "type": "string"
+                    "default": null,
+                    "type": ["string", "null"]
                 },
                 "expiration_time": {
                     "format": "date-time",

frontend/src/queries/schema.ts

@@ -893,8 +893,8 @@ export type QueryStatus = {
     error: boolean
     /**  @default false */
     complete: boolean
-    /**  @default "" */
-    error_message: string
+    /**  @default null */
+    error_message: string | null
     results?: any
     /**  @format date-time */
     start_time?: string

posthog/api/query.py

@@ -92,8 +92,18 @@ def create(self, request, *args, **kwargs) -> Response:
         },
     )
     def retrieve(self, request: Request, pk=None, *args, **kwargs) -> JsonResponse:
-        status = get_query_status(team_id=self.team.pk, query_id=pk)
-        return JsonResponse(status.__dict__, safe=False)
+        query_status = get_query_status(team_id=self.team.pk, query_id=pk)
+
+        http_code: int = status.HTTP_202_ACCEPTED
+        if query_status.error:
+            if query_status.error_message:
+                http_code = status.HTTP_400_BAD_REQUEST  # An error where a user can likely take an action to resolve it
+            else:
+                http_code = status.HTTP_500_INTERNAL_SERVER_ERROR  # An internal surprise
+        elif query_status.complete:
+            http_code = status.HTTP_200_OK
+
+        return JsonResponse(query_status.model_dump(), safe=False, status=http_code)
 
     @extend_schema(
         description="(Experimental)",

posthog/api/test/test_query.py

@@ -833,7 +833,7 @@ def test_full_hogql_query_async(self):
                     "complete": False,
                     "end_time": None,
                     "error": False,
-                    "error_message": "",
+                    "error_message": None,
                     "expiration_time": None,
                     "id": mock.ANY,
                     "query_async": True,
@@ -923,20 +923,33 @@ def test_running_query(self):
             }
         ).encode()
         response = self.client.get(f"/api/projects/{self.team.id}/query/{self.valid_query_id}/")
-        self.assertEqual(response.status_code, 200)
+        self.assertEqual(response.status_code, 202)
         self.assertFalse(response.json()["complete"])
 
-    def test_failed_query(self):
+    def test_failed_query_with_internal_error(self):
         self.redis_client_mock.get.return_value = json.dumps(
             {
                 "id": self.valid_query_id,
                 "team_id": self.team_id,
                 "error": True,
-                "error_message": "Query failed",
+                "error_message": None,
             }
         ).encode()
         response = self.client.get(f"/api/projects/{self.team.id}/query/{self.valid_query_id}/")
-        self.assertEqual(response.status_code, 200)
+        self.assertEqual(response.status_code, 500)
+        self.assertTrue(response.json()["error"])
+
+    def test_failed_query_with_exposed_error(self):
+        self.redis_client_mock.get.return_value = json.dumps(
+            {
+                "id": self.valid_query_id,
+                "team_id": self.team_id,
+                "error": True,
+                "error_message": "Try changing the time range",
+            }
+        ).encode()
+        response = self.client.get(f"/api/projects/{self.team.id}/query/{self.valid_query_id}/")
+        self.assertEqual(response.status_code, 400)
         self.assertTrue(response.json()["error"])
 
     def test_destroy(self):

posthog/clickhouse/client/execute_async.py

@@ -9,7 +9,9 @@
 
 from posthog import celery, redis
 from posthog.clickhouse.query_tagging import tag_queries
+from posthog.errors import ExposedCHQueryError
 from posthog.hogql.constants import LimitContext
+from posthog.hogql.errors import HogQLException
 from posthog.renderers import SafeJSONRenderer
 from posthog.schema import QueryStatus
 from posthog.tasks.tasks import process_query_task
@@ -105,11 +107,15 @@ def execute_process_query(
         query_status.expiration_time = query_status.end_time + datetime.timedelta(seconds=manager.STATUS_TTL_SECONDS)
         process_duration = (query_status.end_time - pickup_time) / datetime.timedelta(seconds=1)
         QUERY_PROCESS_TIME.observe(process_duration)
-    except Exception as err:
+    except (HogQLException, ExposedCHQueryError) as err:  # We can expose the error to the user
         query_status.results = None  # Clear results in case they are faulty
         query_status.error_message = str(err)
         logger.error("Error processing query for team %s query %s: %s", team_id, query_id, err)
         raise err
+    except Exception as err:  # We cannot reveal anything about the error
+        query_status.results = None  # Clear results in case they are faulty
+        logger.error("Error processing query for team %s query %s: %s", team_id, query_id, err)
+        raise err
     finally:
         manager.store_query_status(query_status)
 
@@ -163,7 +169,7 @@ def enqueue_process_query_task(
     return query_status
 
 
-def get_query_status(team_id, query_id):
+def get_query_status(team_id, query_id) -> QueryStatus:
     """
     Abstracts away the manager for any caller and returns a QueryStatus object
     """

posthog/clickhouse/client/test/test_execute_async.py

@@ -69,8 +69,9 @@ def test_async_query_client(self):
         team_id = self.team_id
         query_id = client.enqueue_process_query_task(team_id, self.user_id, query, _test_only_bypass_celery=True).id
         result = client.get_query_status(team_id, query_id)
-        self.assertFalse(result.error, result.error_message)
+        self.assertFalse(result.error, result.error_message or "<no error message>")
         self.assertTrue(result.complete)
+        assert result.results is not None
         self.assertEqual(result.results["results"], [[2]])
 
     def test_async_query_client_errors(self):
@@ -90,15 +91,17 @@ def test_async_query_client_errors(self):
 
         result = client.get_query_status(self.team_id, query_id)
         self.assertTrue(result.error)
+        assert result.error_message
         self.assertRegex(result.error_message, "Unknown table")
 
     def test_async_query_client_uuid(self):
         query = build_query("SELECT toUUID('00000000-0000-0000-0000-000000000000')")
         team_id = self.team_id
         query_id = client.enqueue_process_query_task(team_id, self.user_id, query, _test_only_bypass_celery=True).id
         result = client.get_query_status(team_id, query_id)
-        self.assertFalse(result.error, result.error_message)
+        self.assertFalse(result.error, result.error_message or "<no error message>")
         self.assertTrue(result.complete)
+        assert result.results is not None
         self.assertEqual(result.results["results"], [["00000000-0000-0000-0000-000000000000"]])
 
     def test_async_query_client_does_not_leak(self):

posthog/schema.py

@@ -733,7 +733,7 @@ class QueryStatus(BaseModel):
     complete: Optional[bool] = False
     end_time: Optional[AwareDatetime] = None
     error: Optional[bool] = False
-    error_message: Optional[str] = ""
+    error_message: Optional[str] = None
     expiration_time: Optional[AwareDatetime] = None
     id: str
     query_async: Optional[bool] = True