pmt/db/db.go

package db

import (
	"gorm.io/driver/sqlite"
	"gorm.io/gorm"
	"time"
)

var DB *gorm.DB

func init() {
	var err error
	DB, err = gorm.Open(sqlite.Open("gorm.db"), &gorm.Config{})
	if err != nil {
		panic(err)
	}
}

type App struct {
	ID              uint              `json:"id"               gorm:"primarykey"`      // 应用ID
	PID             int               `json:"pid"`                                     // 应用PID
	Name            string            `json:"name"             gorm:"unique"`          // 应用名称
	Description     string            `json:"description"`                             // 应用描述
	Script          string            `json:"script"`                                  // 启动脚本
	Command         []string          `json:"command"          gorm:"serializer:json"` // 启动命令
	Arguments       []string          `json:"arguments"        gorm:"serializer:json"` // 启动参数
	Stdin           string            `json:"stdin"`                                   // 标准输入数据
	Cwd             string            `json:"cwd"`                                     // 工作目录
	Environments    map[string]string `json:"environments"     gorm:"serializer:json"` // 自定义环境变量
	Interpreter     string            `json:"interpreter"`                             // 脚本解释程序
	InterpreterArgs []string          `json:"interpreter_args" gorm:"serializer:json"` // 解释程序参数
	Options         map[string]any    `json:"options"          gorm:"serializer:json"` // 其它参数
	RerunOnError    *RerunOnError     `json:"rerun_on_error"   gorm:"serializer:json"` // 错误重启策略
	StartCount      int               `json:"start_count"`                             // 应用启动次数
	ErrorCount      int               `json:"error_count"`                             // 错误次数，配合 RerunOnError 使用
	Status          uint8             `json:"status"`                                  // 应用状态
	CreatedAt       time.Time         `json:"created_at"`                              // 创建时间
}

// RerunOnError 错误重启策略
type RerunOnError struct {
	Enable bool  `json:"enable"`         // 是否支持错误重试
	Count  int64 `json:"count_on_error"` // 允许重试次数
}

func SaveApp(app *App) error {
	if app.ID > 0 {
		return DB.Save(app).Error
	} else {
		return DB.Create(app).Error
	}
}

func WithAppID(id uint) func(*gorm.DB) *gorm.DB {
	return func(db *gorm.DB) *gorm.DB {
		return db.Where("id=?", id)
	}
}

func WithAppName(name string) func(*gorm.DB) *gorm.DB {
	return func(db *gorm.DB) *gorm.DB {
		return db.Where("name=?", name)
	}
}

func WithAppStatus(statuses ...uint8) func(*gorm.DB) *gorm.DB {
	return func(db *gorm.DB) *gorm.DB {
		switch len(statuses) {
		case 0:
			return db
		case 1:
			return db.Where("status=?", statuses[0])
		default:
			return db.Where("status IN ?", statuses)
		}
	}
}

// FindApp 查找一个 App
func FindApp(opts ...func(*gorm.DB) *gorm.DB) (*App, error) {
	var app App
	err := DB.Model(&App{}).Scopes(opts...).First(&app).Error
	if err != nil {
		return nil, err
	}
	return &app, nil
}

// FindApps 查找多个 App
func FindApps(opts ...func(*gorm.DB) *gorm.DB) ([]*App, error) {
	var apps []*App
	err := DB.Model(&App{}).Scopes(opts...).Find(&apps).Error
	if err != nil {
		return nil, err
	}
	return apps, nil
}

func ConfigPid(pid int) func(*gorm.DB) *gorm.DB {
	return func(db *gorm.DB) *gorm.DB {
		return db.Set("pid", pid)
	}
}

func ConfigStatus(status uint8) func(*gorm.DB) *gorm.DB {
	return func(db *gorm.DB) *gorm.DB {
		return db.Set("status", status)
	}
}

func IncreaseStartCounter(db *gorm.DB) *gorm.DB {
	return db.Set("starts", gorm.Expr("starts + ?", 1))
}

func UpdateApp(id uint, opts ...func(*gorm.DB) *gorm.DB) error {
	return DB.Model(&App{}).Scopes(opts...).Where("id=?", id).Error
}