From 2423213e9a4a028119a139f27409b6e020a4eff3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=A1=9C=E8=8F=AF?= Date: Sun, 14 May 2023 04:45:52 +0800 Subject: [PATCH] gorm --- configs/sqlite3.go | 11 ++++ go.mod | 7 +++ go.sum | 10 ++++ models/ListView.go | 4 -- models/Model.go | 145 ++------------------------------------------- routers/models.go | 56 ++++++++++++----- 6 files changed, 75 insertions(+), 158 deletions(-) diff --git a/configs/sqlite3.go b/configs/sqlite3.go index 8e043b5..3915353 100644 --- a/configs/sqlite3.go +++ b/configs/sqlite3.go @@ -6,6 +6,8 @@ import ( "os" _ "github.com/mattn/go-sqlite3" + "gorm.io/driver/sqlite" + "gorm.io/gorm" ) // 使用SQLite3初始化數據庫 @@ -103,3 +105,12 @@ func GetDB() (*sql.DB, error) { } return db, nil } + +// ORMDB 使用 GORM +func ORMDB() (db *gorm.DB) { + db, err := gorm.Open(sqlite.Open("data/sqlite3.db"), &gorm.Config{}) + if err != nil { + log.Println(err) + } + return db +} diff --git a/go.mod b/go.mod index 538582e..2911b1b 100644 --- a/go.mod +++ b/go.mod @@ -9,4 +9,11 @@ require ( github.com/gorilla/websocket v1.5.0 github.com/mattn/go-sqlite3 v1.14.16 github.com/russross/blackfriday v1.6.0 + gorm.io/gorm v1.25.1 +) + +require ( + github.com/jinzhu/inflection v1.0.0 // indirect + github.com/jinzhu/now v1.1.5 // indirect + gorm.io/driver/sqlite v1.5.0 ) diff --git a/go.sum b/go.sum index 76cf8d1..7601b57 100644 --- a/go.sum +++ b/go.sum @@ -6,7 +6,17 @@ github.com/gorilla/mux v1.8.0 h1:i40aqfkR1h2SlN9hojwV5ZA91wcXFOvkdNIeFDP5koI= github.com/gorilla/mux v1.8.0/go.mod h1:DVbg23sWSpFRCP0SfiEN6jmj59UnW/n46BH5rLB71So= github.com/gorilla/websocket v1.5.0 h1:PPwGk2jz7EePpoHN/+ClbZu8SPxiqlu12wZP/3sWmnc= github.com/gorilla/websocket v1.5.0/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE= +github.com/jinzhu/inflection v1.0.0 h1:K317FqzuhWc8YvSVlFMCCUb36O/S9MCKRDI7QkRKD/E= +github.com/jinzhu/inflection v1.0.0/go.mod h1:h+uFLlag+Qp1Va5pdKtLDYj+kHp5pxUVkryuEj+Srlc= +github.com/jinzhu/now v1.1.5 h1:/o9tlHleP7gOFmsnYNz3RGnqzefHA47wQpKrrdTIwXQ= +github.com/jinzhu/now v1.1.5/go.mod h1:d3SSVoowX0Lcu0IBviAWJpolVfI5UJVZZ7cO71lE/z8= +github.com/mattn/go-sqlite3 v1.14.15/go.mod h1:2eHXhiwb8IkHr+BDWZGa96P6+rkvnG63S2DGjv9HUNg= github.com/mattn/go-sqlite3 v1.14.16 h1:yOQRA0RpS5PFz/oikGwBEqvAWhWg5ufRz4ETLjwpU1Y= github.com/mattn/go-sqlite3 v1.14.16/go.mod h1:2eHXhiwb8IkHr+BDWZGa96P6+rkvnG63S2DGjv9HUNg= github.com/russross/blackfriday v1.6.0 h1:KqfZb0pUVN2lYqZUYRddxF4OR8ZMURnJIG5Y3VRLtww= github.com/russross/blackfriday v1.6.0/go.mod h1:ti0ldHuxg49ri4ksnFxlkCfN+hvslNlmVHqNRXXJNAY= +gorm.io/driver/sqlite v1.5.0 h1:zKYbzRCpBrT1bNijRnxLDJWPjVfImGEn0lSnUY5gZ+c= +gorm.io/driver/sqlite v1.5.0/go.mod h1:kDMDfntV9u/vuMmz8APHtHF0b4nyBB7sfCieC6G8k8I= +gorm.io/gorm v1.24.7-0.20230306060331-85eaf9eeda11/go.mod h1:L4uxeKpfBml98NYqVqwAdmV1a2nBtAec/cf3fpucW/k= +gorm.io/gorm v1.25.1 h1:nsSALe5Pr+cM3V1qwwQ7rOkw+6UeLrX5O4v3llhHa64= +gorm.io/gorm v1.25.1/go.mod h1:L4uxeKpfBml98NYqVqwAdmV1a2nBtAec/cf3fpucW/k= diff --git a/models/ListView.go b/models/ListView.go index d218b49..57437a1 100644 --- a/models/ListView.go +++ b/models/ListView.go @@ -16,13 +16,9 @@ type ListView struct { // 轉換爲JSON並返回 func (listview *ListView) ToJSON() []byte { - - // 即使list爲空，也要返回空的JSON數組 if listview.List == nil { listview.List = make([]interface{}, 0) } - - // 輸出格式化的JSON b, err := json.MarshalIndent(listview, "", " ") if err != nil { log.Println(err) diff --git a/models/Model.go b/models/Model.go index 0aa8bb0..68fb612 100644 --- a/models/Model.go +++ b/models/Model.go @@ -6,7 +6,7 @@ import ( ) type Model struct { - ID int `json:"id"` + ID int `json:"id" gorm:"primary_key"` Name string `json:"name"` Type string `json:"type"` // (lora|ckp|hyper|ti) TriggerWords string `json:"trigger_words"` // 觸發詞 @@ -21,6 +21,10 @@ type Model struct { UserID int `json:"user_id"` } +func init() { + configs.ORMDB().AutoMigrate(&Model{}) +} + func (model *Model) SendToTrain() error { db, err := configs.GetDB() if err != nil { @@ -43,142 +47,3 @@ func (model *Model) SendToTrain() error { // 將任務放入隊列中, 將自動回調更新任務狀態 return nil } - -func (model *Model) Create() error { - db, err := configs.GetDB() - if err != nil { - log.Println(err) - return err - } - defer db.Close() - stmt, err := db.Prepare("INSERT INTO models(name, type, trigger_words, base_model, model_path, status, progress, image, tags, created_at, updated_at, user_id) values(?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)") - if err != nil { - log.Println(err) - return err - } - defer stmt.Close() - result, err := stmt.Exec(model.Name, model.Type, model.TriggerWords, model.BaseModel, model.ModelPath, model.Status, model.Progress, model.Image, model.Tags, model.CreatedAt, model.UpdatedAt, model.UserID) - if err != nil { - log.Println(err) - return err - } - id, err := result.LastInsertId() - if err != nil { - return err - } - model.ID = int(id) - return nil -} - -func (model *Model) Delete() error { - db, err := configs.GetDB() - if err != nil { - log.Println(err) - return err - } - defer db.Close() - stmt, err := db.Prepare("DELETE FROM models WHERE id = ?") - if err != nil { - log.Println(err) - return err - } - defer stmt.Close() - _, err = stmt.Exec(model.ID) - if err != nil { - log.Println(err) - return err - } - return nil -} - -func (model *Model) Update() error { - db, err := configs.GetDB() - if err != nil { - log.Println(err) - return err - } - defer db.Close() - stmt, err := db.Prepare("UPDATE models SET name = ?, type = ?, updated_at = ? WHERE id = ?") - if err != nil { - log.Println(err) - return err - } - defer stmt.Close() - _, err = stmt.Exec(model.Name, model.Type, model.UpdatedAt, model.ID) - if err != nil { - log.Println(err) - return err - } - return nil -} - -func (model *Model) Get() error { - db, err := configs.GetDB() - if err != nil { - log.Println(err) - return err - } - defer db.Close() - err = db.QueryRow("SELECT id,name,type,trigger_words,base_model,model_path,status,progress,image,tags,creaated_at,updated_at,user_id FROM models WHERE id = ?", model.ID).Scan(&model.ID, &model.Name, &model.Type, &model.TriggerWords, &model.BaseModel, &model.ModelPath, &model.Status, &model.Progress, &model.Image, &model.Tags, &model.CreatedAt, &model.UpdatedAt, &model.UserID) - if err != nil { - log.Println(err) - return err - } - - return nil -} - -func QueryModel(id int) (model Model) { - db, err := configs.GetDB() - if err != nil { - log.Println(err) - return - } - defer db.Close() - err = db.QueryRow("SELECT id,name,type,trigger_words,base_model,model_path,status,progress,image,tags,created_at,updated_at,user_id FROM models WHERE id = ?", id).Scan(&model.ID, &model.Name, &model.Type, &model.TriggerWords, &model.BaseModel, &model.ModelPath, &model.Status, &model.Progress, &model.Image, &model.Tags, &model.CreatedAt, &model.UpdatedAt, &model.UserID) - if err != nil { - log.Println(err) - return - } - return -} - -func QueryModels(page int, pagesize int) (models []interface{}) { - db, err := configs.GetDB() - if err != nil { - log.Println(err) - return - } - defer db.Close() - rows, err := db.Query("SELECT id, name, type, trigger_words, base_model, model_path, status, progress, image, tags, created_at, updated_at, user_id FROM models LIMIT ?, ?", (page-1)*pagesize, pagesize) - if err != nil { - log.Println(err) - return - } - defer rows.Close() - for rows.Next() { - var model Model - err = rows.Scan(&model.ID, &model.Name, &model.Type, &model.TriggerWords, &model.BaseModel, &model.ModelPath, &model.Status, &model.Progress, &model.Image, &model.Tags, &model.CreatedAt, &model.UpdatedAt, &model.UserID) - if err != nil { - log.Println(err) - return - } - models = append(models, model) - } - return -} - -func CountModels() (count int) { - db, err := configs.GetDB() - if err != nil { - log.Println(err) - return - } - defer db.Close() - err = db.QueryRow("SELECT COUNT(*) FROM models").Scan(&count) - if err != nil { - log.Println(err) - return - } - return -} diff --git a/routers/models.go b/routers/models.go index f399031..bebee3d 100644 --- a/routers/models.go +++ b/routers/models.go @@ -5,6 +5,7 @@ import ( "fmt" "io/ioutil" "log" + "main/configs" "main/models" "main/utils" "net/http" @@ -17,11 +18,21 @@ import ( 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() + // 獲取模型列表 + var model_list []models.Model + db := configs.ORMDB() + db.Offset((listview.Page - 1) * listview.PageSize).Limit(listview.PageSize).Find(&model_list) + for _, model := range model_list { + listview.List = append(listview.List, model) + } + // 獲取總數 + var total int64 + db.Model(&models.Model{}).Count(&total) + listview.Total = int(total) listview.Next = listview.Page*listview.PageSize < listview.Total listview.WriteJSON(w) } @@ -43,7 +54,10 @@ func ModelsPost(w http.ResponseWriter, r *http.Request) { log.Println(err) return } - model.Create() + if err := configs.ORMDB().Create(&model).Error; err != nil { + log.Println(err) + return + } w.Header().Set("Content-Type", "application/json; charset=utf-8") w.Write(utils.ToJSON(model)) }) @@ -53,11 +67,13 @@ 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 { + + var model = models.Model{ID: id} + if err := configs.ORMDB().Take(&model, id).Error; err != nil { w.WriteHeader(http.StatusNotFound) return } + upgrader := websocket.Upgrader{} conn, err := upgrader.Upgrade(w, r, nil) if err != nil { @@ -80,16 +96,20 @@ func ModelItemGet(w http.ResponseWriter, r *http.Request) { } return } - model := models.Model{ID: utils.ParamInt(mux.Vars(r)["id"], 0)} - model.Get() + + var model = models.Model{ID: utils.ParamInt(mux.Vars(r)["id"], 0)} + if err := configs.ORMDB().Take(&model, utils.ParamInt(mux.Vars(r)["id"], 0)); err != nil { + w.WriteHeader(http.StatusNotFound) + return + } + w.Header().Set("Content-Type", "application/json; charset=utf-8") w.Write(utils.ToJSON(model)) } func ModelItemPatch(w http.ResponseWriter, r *http.Request) { - // 取出原始数据 - model := models.Model{ID: utils.ParamInt(mux.Vars(r)["id"], 0)} - if err := model.Get(); err != nil { + var model = models.Model{ID: utils.ParamInt(mux.Vars(r)["id"], 0)} + if err := configs.ORMDB().Take(&model, utils.ParamInt(mux.Vars(r)["id"], 0)); err != nil { w.WriteHeader(http.StatusNotFound) return } @@ -125,13 +145,21 @@ func ModelItemPatch(w http.ResponseWriter, r *http.Request) { model.Image = model_new.Image } - model.Update() + if err := configs.ORMDB().Save(&model).Error; err != nil { + log.Println(err) + return + } + 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) + var model = models.Model{ID: utils.ParamInt(mux.Vars(r)["id"], 0)} + if err := configs.ORMDB().Take(&model, utils.ParamInt(mux.Vars(r)["id"], 0)); err != nil { + w.WriteHeader(http.StatusNotFound) + return + } + w.Header().Set("Content-Type", "application/json; charset=utf-8") + w.Write(utils.ToJSON(model)) }