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