65 lines
1.7 KiB
Go
65 lines
1.7 KiB
Go
package models
|
|
|
|
import (
|
|
"gorm.io/gorm"
|
|
)
|
|
|
|
type Tree struct {
|
|
BaseModel
|
|
// ID int64 `json:"id" gorm:"column:id;primary_key;type:bigint(20) auto_increment;"`
|
|
Name string `json:"name" gorm:"type:varchar(255);not null"`
|
|
ParentID string `json:"parent_id" gorm:"column:parent_id; defualt:'-1' "`
|
|
Level int `json:"level" gorm:"column:level; default:0 "`
|
|
}
|
|
|
|
// 获取子节点
|
|
func (t *Tree) Children(db *gorm.DB) ([]Tree, error) {
|
|
var children []Tree
|
|
err := db.Where("parent_id = ?", t.ID).Find(&children).Error
|
|
return children, err
|
|
}
|
|
|
|
// 获取所有后代节点
|
|
//
|
|
// func (t *Tree) Descendants(db *gorm.DB) ([]Tree, error) {
|
|
// var descendants []Tree
|
|
// err := db.Where("level > ? AND (parent_id = ? OR parent_id IN (SELECT id FROM trees WHERE level > ?))",
|
|
// t.Level, t.ID, t.Level).Find(&descendants).Error
|
|
// return descendants, err
|
|
// }
|
|
func (t *Tree) Descendants(db *gorm.DB) ([]Tree, error) {
|
|
var descendants []Tree
|
|
|
|
// 使用GORM递归获取所有后代节点
|
|
var getAllDescendants func(parentID string) error
|
|
getAllDescendants = func(parentID string) error {
|
|
var children []Tree
|
|
if err := db.Where("parent_id = ?", parentID).Find(&children).Error; err != nil {
|
|
return err
|
|
}
|
|
|
|
for _, child := range children {
|
|
descendants = append(descendants, child)
|
|
if err := getAllDescendants(child.ID); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// 从当前节点开始递归
|
|
if err := getAllDescendants(t.ID); err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return descendants, nil
|
|
}
|
|
|
|
func (t *Tree) ChildrenWithMinLevel(db *gorm.DB, minLevel int) ([]Tree, error) {
|
|
var children []Tree
|
|
if err := db.Where("parent_id = ? AND level >= ?", t.ID, minLevel).Find(&children).Error; err != nil {
|
|
return nil, err
|
|
}
|
|
return children, nil
|
|
}
|