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