fix bugs; add alert if query can't start

server/app/alert.tsx

@@ -1,9 +1,12 @@
 import Link from "next/link";
-import { CheckCircleIcon, XMarkIcon } from "@heroicons/react/20/solid";
+import {
+  CheckCircleIcon,
+  ExclamationTriangleIcon,
+} from "@heroicons/react/20/solid";
 
-export default function QueryStartedAlert({ queryId }: { queryId: string }) {
+export function QueryStartedAlert({ queryId }: { queryId: string }) {
   return (
-    <div className="rounded-md bg-green-50 p-4">
+    <div className="-mt-16 mb-4 rounded-md bg-green-50 p-4">
       <div className="flex">
         <div className="shrink-0">
           <CheckCircleIcon
@@ -19,16 +22,25 @@ export default function QueryStartedAlert({ queryId }: { queryId: string }) {
             </Link>.{" "}
           </p>
         </div>
-        <div className="ml-auto pl-3">
-          <div className="-m-1.5">
-            <button
-              type="button"
-              className="inline-flex rounded-md bg-green-50 p-1.5 text-green-500 hover:bg-green-100 focus:outline-none focus:ring-2 focus:ring-green-600 focus:ring-offset-2 focus:ring-offset-green-50"
-            >
-              <span className="sr-only">Dismiss</span>
-              <XMarkIcon className="size-5" aria-hidden="true" />
-            </button>
-          </div>
+      </div>
+    </div>
+  );
+}
+
+export function QueryFailedToStartAlert({ queryId }: { queryId: string }) {
+  return (
+    <div className="-mt-16 mb-4 rounded-md bg-red-50 p-4">
+      <div className="flex">
+        <div className="shrink-0">
+          <ExclamationTriangleIcon
+            className="size-5 text-red-400"
+            aria-hidden="true"
+          />
+        </div>
+        <div className="ml-3">
+          <p className="text-sm font-medium text-red-800">
+            Failed to started Query: {queryId}.
+          </p>
         </div>
       </div>
     </div>

server/app/query/create/page.tsx

@@ -9,7 +9,7 @@ import {
 } from "@heroicons/react/20/solid";
 import { useRouter } from "next/navigation";
 import { ExclamationCircleIcon } from "@heroicons/react/20/solid";
-import QueryStartedAlert from "@/app/alert";
+import { QueryStartedAlert, QueryFailedToStartAlert } from "@/app/alert";
 import {
   DemoLoggerRemoteServers,
   IPARemoteServers,
@@ -23,6 +23,9 @@ type QueryType = Database["public"]["Enums"]["query_type"];
 
 export default function Page() {
   const [queryId, setQueryId] = useState<string | null>(null);
+  const [querySubmitSuccess, setQuerySubmitSuccess] = useState<boolean | null>(
+    null,
+  );
   const router = useRouter();
 
   const handleFormSubmit = async (
@@ -43,9 +46,14 @@ export default function Page() {
           method: "POST",
           body: params,
         });
-        const _data = await response.json();
+        if (!response.ok) {
+          setQuerySubmitSuccess(false);
+          const error_message = await response.text();
+          throw new Error(error_message);
+        }
       }
 
+      setQuerySubmitSuccess(true);
       await new Promise((f) => setTimeout(f, 1000));
 
       // Redirect to /query/view/<newQueryId>
@@ -67,7 +75,12 @@ export default function Page() {
 
   return (
     <>
-      {queryId && <QueryStartedAlert queryId={queryId} />}
+      {queryId && querySubmitSuccess === true && (
+        <QueryStartedAlert queryId={queryId} />
+      )}
+      {queryId && querySubmitSuccess === false && (
+        <QueryFailedToStartAlert queryId={queryId} />
+      )}
       <div className="md:flex md:items-start md:justify-between">
         <DemoLogsForm handleDemoLogsFormSubmit={handleDemoLogsFormSubmit} />
         <IPAForm handleIPAFormSubmit={handleIPAFormSubmit} />

server/app/query/layout.tsx

@@ -7,7 +7,7 @@ export default function RootLayout({
   children: React.ReactNode;
 }) {
   return (
-    <div className="mx-auto min-h-full max-w-7xl bg-slate-100 py-10 sm:px-6 lg:px-8 dark:bg-slate-900">
+    <div className="mx-auto min-h-full max-w-7xl bg-slate-100 py-20 sm:px-6 lg:px-8 dark:bg-slate-900">
       {children}
     </div>
   );

sidecar/app/query/base.py

@@ -169,9 +169,13 @@ class MaxQueriesRunningError(Exception):
 @dataclass
 class QueryManager:
     max_parallel_queries: int = field(init=True, repr=False, default=1)
-    running_queries: set[str] = field(init=False, repr=True, default_factory=set)
+    running_queries: dict[str, Query] = field(
+        init=False, repr=True, default_factory=dict
+    )
 
     def get_from_query_id(self, cls, query_id: str) -> Optional[Query]:
+        if query_id in self.running_queries:
+            return self.running_queries[query_id]
         if status_file_path(query_id).exists():
             return cls(query_id)
         return None
@@ -183,9 +187,9 @@ def run_query(self, query: Query):
                 f"Only {self.max_parallel_queries} allowed. Currently running {self}"
             )
 
-        self.running_queries.add(query.query_id)
+        self.running_queries[query.query_id] = query
         query.start()
-        self.running_queries.remove(query.query_id)
+        del self.running_queries[query.query_id]
 
     @property
     def capacity_available(self):

sidecar/app/routes/start.py

@@ -37,7 +37,7 @@ def running_queries(
     request: Request,
 ):
     query_manager = request.app.state.QUERY_MANAGER
-    return {"running_queries": query_manager.running_queries}
+    return {"running_queries": list(query_manager.running_queries.keys())}
 
 
 @router.post("/demo-logger/{query_id}", status_code=status.HTTP_201_CREATED)

sidecar/app/routes/websockets.py

@@ -2,7 +2,7 @@
 import time
 from contextlib import asynccontextmanager
 
-from fastapi import APIRouter, HTTPException, Request, WebSocket, WebSocketDisconnect
+from fastapi import APIRouter, HTTPException, WebSocket, WebSocketDisconnect
 from websockets import ConnectionClosedError, ConnectionClosedOK
 
 from ..query.base import Query
@@ -30,9 +30,8 @@ async def use_websocket(websocket):
 async def status_websocket(
     websocket: WebSocket,
     query_id: str,
-    request: Request,
 ):
-    query_manager = request.app.state.QUERY_MANAGER
+    query_manager = websocket.app.state.QUERY_MANAGER
     query = query_manager.get_from_query_id(Query, query_id)
     if query is None:
         raise HTTPException(status_code=404, detail="Query not found")
@@ -51,9 +50,8 @@ async def status_websocket(
 async def logs_websocket(
     websocket: WebSocket,
     query_id: str,
-    request: Request,
 ):
-    query_manager = request.app.state.QUERY_MANAGER
+    query_manager = websocket.app.state.QUERY_MANAGER
     query = query_manager.get_from_query_id(Query, query_id)
     if query is None:
         raise HTTPException(status_code=404, detail="Query not found")
@@ -80,9 +78,8 @@ async def logs_websocket(
 async def stats_websocket(
     websocket: WebSocket,
     query_id: str,
-    request: Request,
 ):
-    query_manager = request.app.state.QUERY_MANAGER
+    query_manager = websocket.app.state.QUERY_MANAGER
     query = query_manager.get_from_query_id(Query, query_id)
     if query is None:
         raise HTTPException(status_code=404, detail="Query not found")

sidecar/tests/app/query/test_base.py

@@ -83,9 +83,9 @@ def test_query_manager_capacity_available():
     query_manager = QueryManager(max_parallel_queries=1)
     assert query_manager.capacity_available
     query = Query(str(uuid4()))
-    query_manager.running_queries.add(query.query_id)
+    query_manager.running_queries[query.query_id] = query
     assert not query_manager.capacity_available
-    query_manager.running_queries.remove(query.query_id)
+    del query_manager.running_queries[query.query_id]
     assert query_manager.capacity_available
 
 

sidecar/tests/app/routes/test_start.py

@@ -27,10 +27,10 @@ def __mock_role(role: Role):
 def _running_query():
     query = Query(str(uuid4()))
     query_manager = app.state.QUERY_MANAGER
-    query_manager.running_queries.add(query.query_id)
+    query_manager.running_queries[query.query_id] = query
     query.status = Status.STARTING
     yield query
-    query_manager.running_queries.remove(query.query_id)
+    del query_manager.running_queries[query.query_id]
 
 
 def test_capacity_available():