路由分拆

2023-04-28 16:53:23 +08:00
parent bb58879595
commit 95107e7bdc
10 changed files with 457 additions and 345 deletions
--- a/routers/images.go
+++ b/routers/images.go
@@ -0,0 +1,72 @@
+package routers
+
+import (
+	"encoding/json"
+	"io/ioutil"
+	"log"
+	"main/models"
+	"main/utils"
+	"net/http"
+
+	"github.com/gorilla/mux"
+)
+
+func ImagesGet(w http.ResponseWriter, r *http.Request) {
+	var listview models.ListView
+	listview.Page = utils.ParamInt(r.URL.Query().Get("page"), 1)
+	listview.PageSize = utils.ParamInt(r.URL.Query().Get("pageSize"), 10)
+	listview.List = models.QueryImages(listview.Page, listview.PageSize)
+	listview.Total = models.CountImages()
+	listview.Next = listview.Page*listview.PageSize < listview.Total
+	w.Header().Set("Content-Type", "application/json; charset=utf-8")
+	w.Write(listview.ToJSON())
+}
+
+func ImagesPost(w http.ResponseWriter, r *http.Request) {
+	var image models.Image
+	body, err := ioutil.ReadAll(r.Body)
+	if err != nil {
+		log.Println(err)
+		return
+	}
+	defer r.Body.Close()
+	if err = json.Unmarshal(body, &image); err != nil {
+		log.Println(err)
+		return
+	}
+	image.Create()
+	w.Header().Set("Content-Type", "application/json; charset=utf-8")
+	w.Write(utils.ToJSON(image))
+}
+
+func ImagesItemGet(w http.ResponseWriter, r *http.Request) {
+	image := models.Image{ID: utils.ParamInt(mux.Vars(r)["id"], 0)}
+	image.Get()
+	w.Header().Set("Content-Type", "application/json; charset=utf-8")
+	w.Write(utils.ToJSON(image))
+}
+
+func ImagesItemPatch(w http.ResponseWriter, r *http.Request) {
+	image := models.Image{}
+	body, err := ioutil.ReadAll(r.Body)
+	if err != nil {
+		log.Println(err)
+		return
+	}
+	defer r.Body.Close()
+	if err = json.Unmarshal(body, &image); err != nil {
+		log.Println(err)
+		return
+	}
+	image.ID = utils.ParamInt(mux.Vars(r)["id"], 0)
+	image.Update()
+	w.Header().Set("Content-Type", "application/json; charset=utf-8")
+	w.Write(utils.ToJSON(image))
+}
+
+func ImagesItemDelete(w http.ResponseWriter, r *http.Request) {
+	image := models.Image{ID: utils.ParamInt(mux.Vars(r)["id"], 0)}
+	image.Delete()
+	w.Header().Set("Content-Type", "application/json; charset=utf-8")
+	w.Write(utils.ToJSON(image))
+}
--- a/routers/models.go
+++ b/routers/models.go
@@ -0,0 +1,105 @@
+package routers
+
+import (
+	"encoding/json"
+	"io/ioutil"
+	"log"
+	"main/models"
+	"main/utils"
+	"net/http"
+	"strconv"
+
+	"github.com/gorilla/mux"
+	"github.com/gorilla/websocket"
+)
+
+var manager = models.NewWebSocketManager()
+
+func ModelsGet(w http.ResponseWriter, r *http.Request) {
+	var listview models.ListView
+	listview.Page = utils.ParamInt(r.URL.Query().Get("page"), 1)
+	listview.PageSize = utils.ParamInt(r.URL.Query().Get("pageSize"), 10)
+	listview.List = models.QueryModels(listview.Page, listview.PageSize)
+	listview.Total = models.CountModels()
+	listview.Next = listview.Page*listview.PageSize < listview.Total
+	w.Header().Set("Content-Type", "application/json; charset=utf-8")
+	w.Write(listview.ToJSON())
+}
+
+func ModelsPost(w http.ResponseWriter, r *http.Request) {
+	var model models.Model
+	body, err := ioutil.ReadAll(r.Body)
+	if err != nil {
+		log.Println(err)
+		return
+	}
+	defer r.Body.Close()
+	if err = json.Unmarshal(body, &model); err != nil {
+		log.Println(err)
+		return
+	}
+	model.Create()
+	w.Header().Set("Content-Type", "application/json; charset=utf-8")
+	w.Write(utils.ToJSON(model))
+}
+
+func ModelItemGet(w http.ResponseWriter, r *http.Request) {
+	if r.Header.Get("Upgrade") == "websocket" {
+		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()
+		wsid := manager.AddConnection(conn)
+		defer manager.RemoveConnection(wsid)
+		for {
+			_, msg, err := conn.ReadMessage()
+			if err != nil {
+				log.Println(err)
+				return
+			}
+			log.Println(string(msg))
+			if string(msg) == "close" {
+				break
+			}
+		}
+		return
+	}
+	model := models.Model{ID: utils.ParamInt(mux.Vars(r)["id"], 0)}
+	model.Get()
+	w.Header().Set("Content-Type", "application/json; charset=utf-8")
+	w.Write(utils.ToJSON(model))
+}
+
+func ModelItemPatch(w http.ResponseWriter, r *http.Request) {
+	var model models.Model
+	body, err := ioutil.ReadAll(r.Body)
+	if err != nil {
+		log.Println(err)
+		return
+	}
+	defer r.Body.Close()
+	if err = json.Unmarshal(body, &model); err != nil {
+		log.Println(err)
+		return
+	}
+	model.ID = utils.ParamInt(mux.Vars(r)["id"], 0)
+	model.Update()
+	w.Header().Set("Content-Type", "application/json; charset=utf-8")
+	w.Write(utils.ToJSON(model))
+}
+
+func ModelItemDelete(w http.ResponseWriter, r *http.Request) {
+	model := models.Model{ID: utils.ParamInt(mux.Vars(r)["id"], 0)}
+	model.Delete()
+	w.WriteHeader(http.StatusNoContent)
+}
--- a/routers/params.go
+++ b/routers/params.go
@@ -0,0 +1,15 @@
+package routers
+
+import (
+	"main/utils"
+	"net/http"
+)
+
+func ParamsModelsGet(w http.ResponseWriter, r *http.Request) {
+	params := make(map[string]interface{})
+	params["type"] = []string{"lora", "ckp", "hyper", "ti"}
+	params["status"] = []string{"pending", "running", "finished", "failed"}
+	params["base_model"] = []string{"SD1.5", "SD2"}
+	w.Header().Set("Content-Type", "application/json; charset=utf-8")
+	w.Write(utils.ToJSON(params))
+}
--- a/routers/tasks.go
+++ b/routers/tasks.go
@@ -0,0 +1,104 @@
+package routers
+
+import (
+	"encoding/json"
+	"io/ioutil"
+	"log"
+	"main/models"
+	"main/utils"
+	"net/http"
+	"strconv"
+
+	"github.com/gorilla/mux"
+	"github.com/gorilla/websocket"
+)
+
+func TasksGet(w http.ResponseWriter, r *http.Request) {
+	var listview models.ListView
+	listview.Page = utils.ParamInt(r.URL.Query().Get("page"), 1)
+	listview.PageSize = utils.ParamInt(r.URL.Query().Get("pageSize"), 10)
+	listview.List = models.QueryTasks(listview.Page, listview.PageSize)
+	listview.Total = models.CountTasks()
+	listview.Next = listview.Page*listview.PageSize < listview.Total
+	w.Header().Set("Content-Type", "application/json; charset=utf-8")
+	w.Write(listview.ToJSON())
+}
+
+func TasksPost(w http.ResponseWriter, r *http.Request) {
+	var task models.Task
+	body, err := ioutil.ReadAll(r.Body)
+	if err != nil {
+		log.Println(err)
+		return
+	}
+	defer r.Body.Close()
+	if err = json.Unmarshal(body, &task); err != nil {
+		log.Println(err)
+		return
+	}
+	task.Create()
+	w.Header().Set("Content-Type", "application/json; charset=utf-8")
+	w.Write(utils.ToJSON(task))
+}
+
+func TasksItemGet(w http.ResponseWriter, r *http.Request) {
+	if r.Header.Get("Upgrade") == "websocket" {
+		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()
+		}
+		return
+	}
+	task := models.Task{ID: utils.ParamInt(mux.Vars(r)["id"], 0)}
+	task.Get()
+	w.Header().Set("Content-Type", "application/json; charset=utf-8")
+	w.Write(utils.ToJSON(task))
+}
+
+func TasksItemPatch(w http.ResponseWriter, r *http.Request) {
+	var task models.Task
+	body, err := ioutil.ReadAll(r.Body)
+	if err != nil {
+		log.Println(err)
+		return
+	}
+	defer r.Body.Close()
+	if err = json.Unmarshal(body, &task); err != nil {
+		log.Println(err)
+		return
+	}
+	task.ID = utils.ParamInt(mux.Vars(r)["id"], 0)
+	task.Update()
+	w.Header().Set("Content-Type", "application/json; charset=utf-8")
+	w.Write(utils.ToJSON(task))
+}
+
+func TasksItemDelete(w http.ResponseWriter, r *http.Request) {
+	task := models.Task{ID: utils.ParamInt(mux.Vars(r)["id"], 0)}
+	task.Delete()
+	if task.ID == 0 {
+		w.WriteHeader(http.StatusNotFound)
+		return
+	}
+	task.Delete()
+	w.WriteHeader(http.StatusNoContent)
+}