fix(queries): Fix cancel_query to cancel on cluster (#18819)

posthog/clickhouse/client/execute_async.py

@@ -141,16 +141,21 @@ def get_query_status(team_id, query_id):
 
 
 def cancel_query(team_id, query_id):
-    query_status = get_query_status(team_id, query_id)
+    try:
+        query_status = get_query_status(team_id, query_id)
+
+        if query_status.task_id:
+            logger.info("Got task id %s, attempting to revoke", query_status.task_id)
+            celery.app.control.revoke(query_status.task_id, terminate=True)
 
-    if query_status.task_id:
-        logger.info("Got task id %s, attempting to revoke", query_status.task_id)
-        celery.app.control.revoke(query_status.task_id, terminate=True)
+            logger.info("Revoked task id %s", query_status.task_id)
+    except QueryNotFoundError:
+        # Continue, to attempt to cancel the query even if it's not a task
+        pass
 
-        from posthog.clickhouse.cancel import cancel_query_on_cluster
+    from posthog.clickhouse.cancel import cancel_query_on_cluster
 
-        logger.info("Revoked task id %s, attempting to cancel on cluster", query_status.task_id)
-        cancel_query_on_cluster(team_id, query_id)
+    cancel_query_on_cluster(team_id, query_id)
 
     redis_client = redis.get_client()
     key = generate_redis_results_key(query_id, team_id)