gorm

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
+}

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
 )

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=

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)

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
-}

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) {
-	// 取出原始数据
+	var model = models.Model{ID: utils.ParamInt(mux.Vars(r)["id"], 0)}
-	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 {
-	if err := model.Get(); 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)}
+	var model = models.Model{ID: utils.ParamInt(mux.Vars(r)["id"], 0)}
-	model.Delete()
+	if err := configs.ORMDB().Take(&model, utils.ParamInt(mux.Vars(r)["id"], 0)); err != nil {
-	w.WriteHeader(http.StatusNoContent)
+		w.WriteHeader(http.StatusNotFound)
+		return
+	}
+	w.Header().Set("Content-Type", "application/json; charset=utf-8")
+	w.Write(utils.ToJSON(model))
 }