From 33a468c0bceff7841abe168a6bc825d6ccf96a6f Mon Sep 17 00:00:00 2001 From: zhangqian <zhangqian@123.com> Date: 星期四, 12 十二月 2024 14:25:28 +0800 Subject: [PATCH] 查询模型任务时只查询启用的模型 --- service/task.go | 15 +++ db/model.go | 186 ++++++++++++++++++++++++++++++++++++++++++++++ db/task.go | 10 ++ 3 files changed, 210 insertions(+), 1 deletions(-) diff --git a/db/model.go b/db/model.go new file mode 100644 index 0000000..97dcef1 --- /dev/null +++ b/db/model.go @@ -0,0 +1,186 @@ +package db + +import ( + "fmt" + "gorm.io/gorm" + "model-engine/pkg/mysqlx" +) + +type Model struct { + BaseModel + Name string `json:"name" gorm:"type:varchar(255)"` //妯″瀷鍚嶇О + Description string `json:"description,omitempty" gorm:"type:varchar(1000)"` //妯″瀷鎻忚堪 + Version string `json:"version" gorm:"type:varchar(255)"` //鐗堟湰鍙� + Enabled bool `json:"enabled"` //鏄惁寮�鍚� +} + +func (m *Model) TableName() string { + return "model" +} + +type ModelSearch struct { + Model + Orm *gorm.DB + PageNum int + PageSize int + Order string + Keyword string +} + +func NewModelSearch() *ModelSearch { + return &ModelSearch{ + Orm: mysqlx.GetDB(), + PageNum: 1, + PageSize: 10, + } +} + +func (slf *ModelSearch) SetOrm(tx *gorm.DB) *ModelSearch { + slf.Orm = tx + return slf +} + +func (slf *ModelSearch) SetPage(page, size int) *ModelSearch { + slf.PageNum, slf.PageSize = page, size + return slf +} + +func (slf *ModelSearch) SetOrder(order string) *ModelSearch { + slf.Order = order + return slf +} + +func (slf *ModelSearch) SetID(id string) *ModelSearch { + slf.ID = id + return slf +} + +func (slf *ModelSearch) SetKeyword(kw string) *ModelSearch { + slf.Keyword = kw + return slf +} + +func (slf *ModelSearch) SetEnabled(enabled bool) *ModelSearch { + slf.Enabled = enabled + return slf +} + +func (slf *ModelSearch) build() *gorm.DB { + var db = slf.Orm.Table(slf.TableName()) + if slf.Order != "" { + db = db.Order(slf.Order) + } + + if slf.ID != "" { + db = db.Where("id = ?", slf.ID) + } + + if slf.Keyword != "" { + kw := "%" + slf.Keyword + "%" + db = db.Where("name like ?", kw) + } + + if slf.Enabled { + db = db.Where("enabled = ?", slf.Enabled) + } + + return db +} + +func (slf *ModelSearch) First() (*Model, error) { + var ( + record = new(Model) + db = slf.build() + ) + + if err := db.First(record).Error; err != nil { + return record, err + } + + return record, nil +} + +func (slf *ModelSearch) Find() ([]*Model, int64, error) { + var ( + records = make([]*Model, 0) + total int64 + db = slf.build() + ) + + if err := db.Count(&total).Error; err != nil { + return records, total, fmt.Errorf("find count err: %v", err) + } + if slf.PageNum*slf.PageSize > 0 { + db = db.Offset((slf.PageNum - 1) * slf.PageSize).Limit(slf.PageSize) + } + if err := db.Find(&records).Error; err != nil { + return records, total, fmt.Errorf("find records err: %v", err) + } + + return records, total, nil +} + +func (slf *ModelSearch) FindAll() ([]*Model, error) { + var ( + records = make([]*Model, 0) + db = slf.build() + ) + if err := db.Find(&records).Error; err != nil { + return records, fmt.Errorf("find records err: %v", err) + } + + return records, nil +} + +func (slf *ModelSearch) Count() int64 { + var ( + count int64 + db = slf.build() + ) + + if err := db.Count(&count).Error; err != nil { + return count + } + + return count +} + +func (slf *ModelSearch) Create(record *Model) error { + var db = slf.build() + + if err := db.Create(record).Error; err != nil { + return fmt.Errorf("create err: %v, record: %+v", err, record) + } + + return nil +} + +func (slf *ModelSearch) Save(record *Model) error { + var db = slf.build() + + if err := db.Omit("CreatedAt").Save(record).Error; err != nil { + return fmt.Errorf("save err: %v, record: %+v", err, record) + } + + return nil +} + +func (slf *ModelSearch) Update(record *Model) error { + var db = slf.build() + + if err := db.Updates(record).Error; err != nil { + return fmt.Errorf("update err: %v, record: %+v", err, record) + } + + return nil +} + +func (slf *ModelSearch) Delete() error { + var db = slf.build() + return db.Delete(&Model{}).Error +} + +const ( + ModelIdDrug = "drug" //娑夋瘨 + ModelIdGather = "gather" //鑱氶泦 +) diff --git a/db/task.go b/db/task.go index 573558c..7f21a2a 100644 --- a/db/task.go +++ b/db/task.go @@ -32,6 +32,7 @@ PageNum int PageSize int Keyword string + ModelIDs []string } ) @@ -102,6 +103,11 @@ return slf } +func (slf *ModelTaskSearch) SetModelIDs(ids []string) *ModelTaskSearch { + slf.ModelIDs = ids + return slf +} + func (slf *ModelTaskSearch) SetKeyword(kw string) *ModelTaskSearch { slf.Keyword = kw return slf @@ -123,6 +129,10 @@ db = db.Where("model_id = ?", slf.ModelID) } + if len(slf.ModelIDs) != 0 { + db = db.Where("model_id in ?", slf.ModelIDs) + } + if slf.Keyword != "" { kw := "%" + slf.Keyword + "%" db = db.Where("name like ?", kw) diff --git a/service/task.go b/service/task.go index b0afcfb..f47d7d2 100644 --- a/service/task.go +++ b/service/task.go @@ -3,5 +3,18 @@ import "model-engine/db" func GetTasks() (tasks []*db.ModelTask, err error) { - return db.NewModelTaskSearch().FindAll() + models, err := db.NewModelSearch().SetEnabled(true).FindAll() + if err != nil { + return nil, err + } + if len(models) == 0 { + return nil, nil + } + + modelIds := make([]string, 0, len(models)) + for _, model := range models { + modelIds = append(modelIds, model.ID) + } + + return db.NewModelTaskSearch().SetModelIDs(modelIds).FindAll() } -- Gitblit v1.8.0