diff --git a/server/worker.go b/server/worker.go index e79c2ea..fa5fd18 100644 --- a/server/worker.go +++ b/server/worker.go @@ -36,6 +36,21 @@ func handleWorkerWebsocket(w http.ResponseWriter, r *http.Request) { w.WriteHeader(http.StatusNotImplemented) return } + // Connection ID + uuid := r.Header.Get(constants.UUID_HEADER) + if uuid == "" { + w.WriteHeader(http.StatusBadRequest) + return + } + name := r.Header.Get(constants.NAME_HEADER) + + _, err := db.Exec(r.Context(), "INSERT INTO workers (id, name) VALUES ($1, $2) ON CONFLICT (id) DO UPDATE SET name = $2", uuid, name) + if err != nil { + w.WriteHeader(http.StatusInternalServerError) + slog.ErrorContext(r.Context(), "Error Upserting Worker", "err", err) + return + } + c, err := websocket.Accept(w, r, &websocket.AcceptOptions{ InsecureSkipVerify: true, }) @@ -45,9 +60,6 @@ func handleWorkerWebsocket(w http.ResponseWriter, r *http.Request) { } defer c.CloseNow() - // Connection ID - uuid := r.Header.Get(constants.UUID_HEADER) - // Track Connection func() { WorkersMutex.Lock() @@ -70,7 +82,7 @@ func handleWorkerWebsocket(w http.ResponseWriter, r *http.Request) { } Workers[uuid] = &Worker{ - Name: r.Header.Get(constants.NAME_HEADER), + Name: name, Address: r.RemoteAddr, Conn: c, Connected: true,