websocket

main.go

@@ -14,6 +14,7 @@ import (
 	"main/models"
 
 	"github.com/gorilla/mux"
+	"github.com/gorilla/websocket"
 )
 
 func main() {
@@ -75,7 +76,45 @@ func models_post(w http.ResponseWriter, r *http.Request) {
 	w.Write(ToJSON(model))
 }
 
+func models_item_get_ws(w http.ResponseWriter, r *http.Request) {
+	vars := mux.Vars(r)
+	id, _ := strconv.Atoi(vars["id"])
+	model := models.QueryModel(id)
+	if model.ID == 0 {
+		w.WriteHeader(http.StatusNotFound)
+		return
+	}
+	upgrader := websocket.Upgrader{}
+	conn, err := upgrader.Upgrade(w, r, nil)
+	if err != nil {
+		log.Println(err)
+		return
+	}
+	defer conn.Close()
+	for {
+		time.Sleep(1 * time.Second)
+		model = models.QueryModel(id)
+		if model.ID == 0 {
+			w.WriteHeader(http.StatusNotFound)
+			return
+		}
+		if model.Status == "success" || model.Status == "error" {
+			break
+		}
+		err = conn.WriteJSON(model)
+		if err != nil {
+			log.Println(err)
+			return
+		}
+	}
+}
+
 func models_item_get(w http.ResponseWriter, r *http.Request) {
+	if r.Header.Get("Upgrade") == "websocket" {
+		models_item_get_ws(w, r)
+		return
+	}
+
 	model := models.Model{ID: ParamInt(mux.Vars(r)["id"], 0)}
 	model.Get()
 	w.Header().Set("Content-Type", "application/json; charset=utf-8")
@@ -195,7 +234,38 @@ func tasks_post(w http.ResponseWriter, r *http.Request) {
 	w.Write(ToJSON(task))
 }
 
+func tasks_item_get_ws(w http.ResponseWriter, r *http.Request) {
+	vars := mux.Vars(r)
+	id, _ := strconv.Atoi(vars["id"])
+	task := models.QueryTask(id)
+	if task.ID == 0 {
+		w.WriteHeader(http.StatusNotFound)
+		return
+	}
+	upgrader := websocket.Upgrader{}
+	ws, err := upgrader.Upgrade(w, r, nil)
+	if err != nil {
+		log.Println(err)
+		return
+	}
+	defer ws.Close()
+	for {
+		_, message, err := ws.ReadMessage()
+		if err != nil {
+			log.Println(err)
+			break
+		}
+		task.Status = string(message)
+		task.Update()
+	}
+}
+
 func tasks_item_get(w http.ResponseWriter, r *http.Request) {
+	if r.Header.Get("Upgrade") == "websocket" {
+		tasks_item_get_ws(w, r)
+		return
+	}
+
 	task := models.Task{ID: ParamInt(mux.Vars(r)["id"], 0)}
 	task.Get()
 	w.Header().Set("Content-Type", "application/json; charset=utf-8")