Files
lone-services/pkg/modelbase/base.go
T
2026-08-05 17:37:32 +08:00

331 lines
7.2 KiB
Go

package modelbase
import (
"encoding/json"
"errors"
"fmt"
"gorm.io/gorm"
)
const (
DefaultPage = 1
DefaultSize = 20
)
var (
db *gorm.DB
prefix string
debug bool
)
// Config for Init.
type Config struct {
Prefix string
Debug bool
}
// Init binds the process DB used by Base helpers. Call once after mysql.New.
func Init(conn *gorm.DB, c Config) {
db = conn
prefix = c.Prefix
debug = c.Debug
}
// DB returns the package-level connection (nil before Init).
func DB() *gorm.DB {
return db
}
// Prefix returns the configured table prefix (e.g. lone_).
func Prefix() string {
return prefix
}
type GroupCountResult struct {
Item string `gorm:"column:item"`
Count int `gorm:"column:count"`
}
type GroupSumCountIntResult struct {
Item string `gorm:"column:item"`
Count int `gorm:"column:count"`
Sum int `gorm:"column:sum"`
}
type GroupSumCountFloat32Result struct {
Item string `gorm:"column:item"`
Count int `gorm:"column:count"`
Sum float32 `gorm:"column:sum"`
}
type Params struct {
Eq map[string]string
Not map[string]string
Or []map[string]string
In map[string][]string
Other map[string]string
Like map[string]string
Between map[string][2]string
Join []string
Order string
Page int
Size int
}
type PageResult struct {
Count int64
Items interface{}
}
// Base is embedded by per-table models in each service.
type Base struct {
Table string
tx *gorm.DB
inTx bool
}
// WithTX runs subsequent ops on an outer transaction (e.g. gorm.Transaction).
func (m Base) WithTX(tx *gorm.DB) Base {
m.tx = tx
m.inTx = true
return m
}
func (m *Base) Begin() *gorm.DB {
m.tx = db.Begin()
m.inTx = true
return m.tx
}
func (m *Base) Rollback() *gorm.DB {
if m.tx != nil {
m.tx.Rollback()
}
m.inTx = false
return m.tx
}
func (m *Base) Commit() *gorm.DB {
if m.tx != nil {
m.tx.Commit()
}
m.inTx = false
return m.tx
}
func (m Base) session() *gorm.DB {
var o *gorm.DB
if m.inTx && m.tx != nil {
o = m.tx.Table(m.Table)
} else {
o = db.Table(m.Table)
}
if debug {
o = o.Debug()
}
return o
}
// Create inserts a row.
func (m Base) Create(data interface{}) error {
return m.session().Create(data).Error
}
// Edit updates rows matching Params.
func (m Base) Edit(w Params, data interface{}) (int64, error) {
o := m.setWhere(w)
update := make(map[string]interface{})
jsonBytes, _ := json.Marshal(data)
_ = json.Unmarshal(jsonBytes, &update)
ret := o.Updates(update)
return ret.RowsAffected, ignoreNotFound(ret.Error)
}
// Del deletes rows matching Params.
func (m Base) Del(w Params, data interface{}) (int64, error) {
o := m.setWhere(w)
ret := o.Delete(data)
return ret.RowsAffected, ignoreNotFound(ret.Error)
}
// Sql runs a raw query.
func (m Base) Sql(sql string, retData interface{}) error {
err := m.session().Raw(sql).Scan(retData).Error
return ignoreNotFound(err)
}
// GetById loads by primary key fields already set on retData.
func (m Base) GetById(retData interface{}) error {
err := m.session().First(retData).Error
return ignoreNotFound(err)
}
// GetOne loads one row matching Params.
func (m Base) GetOne(w Params, retData interface{}) error {
o := m.setWhere(w)
err := o.Order(w.Order).First(retData).Error
return ignoreNotFound(err)
}
// GetIDsByKeyword plucks a column into []string.
func (m Base) GetIDsByKeyword(w Params, column string) ([]string, error) {
var results []string
o := m.setWhere(w)
err := o.Pluck(column, &results).Error
if err != nil {
return nil, err
}
return results, nil
}
// Count finds into retData then returns count (legacy behavior).
func (m Base) Count(w Params, retData interface{}) (int64, error) {
o := m.setWhere(w)
var count int64
err := o.Find(retData).Count(&count).Error
return count, ignoreNotFound(err)
}
func (m Base) CountNumber(w Params) (int64, error) {
o := m.setWhere(w)
var count int64
err := o.Count(&count).Error
return count, ignoreNotFound(err)
}
func (m Base) Page(w Params, retData interface{}) (PageResult, error) {
var items PageResult
o := m.setWhere(w)
var count int64
err := o.Count(&count).Error
if err != nil {
return items, ignoreNotFound(err)
}
if w.Page < DefaultPage {
w.Page = DefaultPage
}
if w.Size < DefaultPage {
w.Size = DefaultSize
}
err = o.Offset((w.Page - 1) * w.Size).Limit(w.Size).Order(w.Order).Find(retData).Error
items.Count = count
items.Items = retData
return items, ignoreNotFound(err)
}
func (m Base) CountSumNumber(w Params) (int64, error) {
return m.CountNumber(w)
}
func (m Base) CountSumFloat32(w Params, key string) (float32, error) {
o := m.setWhere(w)
var count float32
fs := fmt.Sprintf("IFNULL(SUM(%s), 0) AS %s", key, key)
err := o.Select(fs).Find(&count).Error
return count, ignoreNotFound(err)
}
func (m Base) CountGroupNumber(w Params, key string) ([]GroupCountResult, error) {
o := m.setWhere(w)
var count []GroupCountResult
fs := fmt.Sprintf("%s AS item, COUNT(*) AS count", key)
err := o.Select(fs).Group(key).Find(&count).Error
return count, ignoreNotFound(err)
}
func (m Base) GroupSumNumber(w Params, group, key string) (GroupSumCountIntResult, error) {
o := m.setWhere(w)
var count GroupSumCountIntResult
fs := fmt.Sprintf("%s AS item, SUM(%s) AS sum", group, key)
err := o.Select(fs).Group(key).Scan(&count).Error
return count, ignoreNotFound(err)
}
func (m Base) CountIntNumber(w Params, key string) (GroupSumCountIntResult, error) {
o := m.setWhere(w)
var count GroupSumCountIntResult
fs := fmt.Sprintf("COUNT(*) AS count, SUM(%s) AS sum", key)
err := o.Select(fs).Find(&count).Error
return count, ignoreNotFound(err)
}
func (m Base) CountFloat32Number(w Params, key string) (GroupSumCountFloat32Result, error) {
o := m.setWhere(w)
var count GroupSumCountFloat32Result
fs := fmt.Sprintf("COUNT(*) AS count, IFNULL(SUM(%s), 0) AS sum", key)
err := o.Select(fs).Find(&count).Error
return count, ignoreNotFound(err)
}
func (m Base) SumFloat32(w Params, key string) (float32, error) {
o := m.setWhere(w)
var sum float32
key = fmt.Sprintf("IFNULL(SUM(%s), 0)", key)
err := o.Select(key).Scan(&sum).Error
return sum, ignoreNotFound(err)
}
func (m Base) SumInt(w Params, key string) (int, error) {
o := m.setWhere(w)
var sum int
key = fmt.Sprintf("IFNULL(SUM(%s), 0)", key)
err := o.Select(key).Scan(&sum).Error
return sum, ignoreNotFound(err)
}
// Items loads all matching rows into retData.
func (m Base) Items(w Params, retData interface{}) error {
o := m.setWhere(w)
err := o.Order(w.Order).Find(retData).Error
return ignoreNotFound(err)
}
func (m Base) setWhere(w Params) *gorm.DB {
o := m.session().Where(w.Eq)
if w.Or != nil {
for i := 0; i < len(w.Or); i++ {
o = o.Or(w.Or[i])
}
}
if w.In != nil {
for k, v := range w.In {
o = o.Where(k, v)
}
}
if w.Like != nil {
for k, v := range w.Like {
o = o.Where(k, v)
}
}
if w.Join != nil {
for i := 0; i < len(w.Join); i++ {
o = o.Joins(w.Join[i])
}
}
if w.Other != nil {
for k, v := range w.Other {
o = o.Where(k, v)
}
}
if w.Between != nil {
for k, v := range w.Between {
if len(v) == 2 {
o = o.Where(k+" BETWEEN ? AND ?", v[0], v[1])
}
}
}
if w.Not != nil {
o = o.Not(w.Not)
}
return o
}
func ignoreNotFound(err error) error {
if errors.Is(err, gorm.ErrRecordNotFound) {
return nil
}
return err
}