ai/routers/datasets.go

package routers

import (
	"database/sql/driver"
	"encoding/json"
	"io/ioutil"
	"main/configs"
	"main/models"
	"main/utils"
	"net/http"
	"time"

	"github.com/gorilla/mux"
)

type ImageList []string

func (list *ImageList) Scan(value interface{}) error {
	return json.Unmarshal(value.([]byte), list)
}
func (list ImageList) Value() (driver.Value, error) {
	return json.Marshal(list)
}

// 數據集模型
type Dataset struct {
	ID        int       `json:"id" gorm:"primary_key"`
	Name      string    `json:"name"`
	Info      string    `json:"info"`
	Images    ImageList `json:"images"`
	UserID    int       `json:"user_id"`
	CreatedAt time.Time `json:"created_at" gorm:"autoCreateTime"`
	UpdatedAt time.Time `json:"updated_at" gorm:"autoUpdateTime"`
}

func init() {
	configs.ORMDB().AutoMigrate(&Dataset{})
}

// 獲取數據集列表
func DatasetsGet(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)
	var dataset_list []Dataset
	db := configs.ORMDB()
	db.Offset((listview.Page - 1) * listview.PageSize).Limit(listview.PageSize).Find(&dataset_list)
	for _, dataset := range dataset_list {
		if dataset.Images == nil {
			dataset.Images = ImageList{}
		}
		listview.List = append(listview.List, dataset)
	}
	db.Model(&Dataset{}).Count(&listview.Total)
	listview.Next = listview.Page*listview.PageSize < int(listview.Total)
	listview.WriteJSON(w)
}

// 創建數據集
func DatasetsPost(w http.ResponseWriter, r *http.Request) {
	models.AccountRead(w, r, func(account *models.Account) {
		var dataset Dataset
		body, err := ioutil.ReadAll(r.Body)
		if err != nil {
			w.WriteHeader(http.StatusInternalServerError)
			w.Write([]byte("500 - Internal Server Error"))
			return
		}
		defer r.Body.Close()
		if err = json.Unmarshal(body, &dataset); err != nil {
			w.WriteHeader(http.StatusInternalServerError)
			w.Write([]byte("500 - Internal Server Error"))
			return
		}
		if dataset.Images == nil {
			dataset.Images = ImageList{}
		}
		dataset.UserID = account.ID
		if err := configs.ORMDB().Create(&dataset).Error; err != nil {
			w.WriteHeader(http.StatusInternalServerError)
			w.Write([]byte("500 - Internal Server Error"))
			return
		}
		w.Header().Set("Content-Type", "application/json; charset=utf-8")
		w.Write(utils.ToJSON(dataset))
	})
}

// 獲取數據集
func DatasetsItemGet(w http.ResponseWriter, r *http.Request) {
	dataset := Dataset{ID: utils.ParamInt(mux.Vars(r)["dataset_id"], 0)}
	if err := configs.ORMDB().Find(&dataset).Error; err != nil {
		w.WriteHeader(http.StatusNotFound)
		w.Write([]byte("404 - Not Found"))
		return
	}
	w.Header().Set("Content-Type", "application/json; charset=utf-8")
	w.Write(utils.ToJSON(dataset))
}

// 修改數據集
func DatasetsItemPatch(w http.ResponseWriter, r *http.Request) {
	models.AccountRead(w, r, func(account *models.Account) {
		var dataset Dataset = Dataset{ID: utils.ParamInt(mux.Vars(r)["dataset_id"], 0)}
		if err := configs.ORMDB().Find(&dataset).Error; err != nil {
			w.WriteHeader(http.StatusNotFound)
			w.Write([]byte("404 - Not Found"))
			return
		}
		if dataset.UserID != account.ID && !account.Admin {
			w.WriteHeader(http.StatusForbidden)
			w.Write([]byte("403 - Forbidden"))
			return
		}
		var form map[string]interface{} = utils.BodyRead(r)
		if name, ok := form["name"].(string); ok {
			dataset.Name = name
		}
		if info, ok := form["info"].(string); ok {
			dataset.Info = info
		}
		if err := configs.ORMDB().Save(&dataset).Error; err != nil {
			w.WriteHeader(http.StatusInternalServerError)
			w.Write([]byte("500 - Internal Server Error"))
			return
		}
		w.Header().Set("Content-Type", "application/json; charset=utf-8")
		w.Write(utils.ToJSON(dataset))
	})
}

// 刪除數據集
func DatasetsItemDelete(w http.ResponseWriter, r *http.Request) {
	models.AccountRead(w, r, func(account *models.Account) {
		// 獲取數據集
		dataset := Dataset{ID: utils.ParamInt(mux.Vars(r)["id"], 0)}
		if err := configs.ORMDB().Find(&dataset).Error; err != nil {
			w.WriteHeader(http.StatusNotFound)
			w.Write([]byte("404 - Not Found"))
			return
		}
		// 只能刪除自己的數據集, 除非是管理員
		if dataset.UserID != account.ID && !account.Admin {
			w.WriteHeader(http.StatusForbidden)
			w.Write([]byte("403 - Forbidden"))
			return
		}
		// 刪除數據集
		if err := configs.ORMDB().Delete(&dataset).Error; err != nil {
			w.WriteHeader(http.StatusNotFound)
			w.Write([]byte("404 - Not Found"))
			return
		}
		w.WriteHeader(http.StatusNoContent)
	})
}