Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
merged 9 commits into from
Mar 28, 2024
Merged
4 changes: 2 additions & 2 deletions frontend/src/queries/schema.json
Original file line number Diff line number Diff line change
Expand Up @@ -4469,8 +4469,8 @@
"type": "boolean"
},
"error_message": {
"default": "",
"type": "string"
"default": null,
"type": ["string", "null"]
},
"expiration_time": {
"format": "date-time",
Expand Down
4 changes: 2 additions & 2 deletions frontend/src/queries/schema.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
14 changes: 12 additions & 2 deletions posthog/api/query.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)",
Expand Down
23 changes: 18 additions & 5 deletions posthog/api/test/test_query.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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):
Expand Down
10 changes: 8 additions & 2 deletions posthog/clickhouse/client/execute_async.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)

Expand Down Expand Up @@ -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
"""
Expand Down
7 changes: 5 additions & 2 deletions posthog/clickhouse/client/test/test_execute_async.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand All @@ -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):
Expand Down
2 changes: 1 addition & 1 deletion posthog/schema.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Loading