xorm/preload.go

package xorm

import (
	"fmt"
	"reflect"
	"strings"
	"xorm.io/builder"
	"xorm.io/xorm/schemas"
)

type Preload struct {
	path    []string
	cols    []string
	cond    builder.Cond
	noPrune bool
}

func NewPreload(path string) *Preload {
	return &Preload{
		path: strings.Split(path, "."),
		cond: builder.NewCond(),
	}
}

func (p *Preload) Cols(cols ...string) *Preload {
	p.cols = append(p.cols, cols...)
	return p
}

func (p *Preload) Where(cond builder.Cond) *Preload {
	p.cond = p.cond.And(cond)
	return p
}

func (p *Preload) NoPrune() *Preload {
	p.noPrune = true
	return p
}

type PreloadNode struct {
	preload     *Preload
	children    map[string]*PreloadNode
	association *schemas.Association
	ExtraCols   []string
}

func NewPreloadNode() *PreloadNode {
	return &PreloadNode{
		children: make(map[string]*PreloadNode),
	}
}

func (pn *PreloadNode) Add(preload *Preload) error {
	return pn.add(preload, 0)
}

func (pn *PreloadNode) add(preload *Preload, index int) error {
	if index == len(preload.path) {
		if pn.preload != nil {
			return fmt.Errorf("preload: duplicated path: %s", strings.Join(preload.path, ","))
		}
		pn.preload = preload
		return nil
	}
	child, ok := pn.children[preload.path[index]]
	if !ok {
		child = NewPreloadNode()
		pn.children[preload.path[index]] = child
	}
	return child.add(preload, index+1)
}

func (pn *PreloadNode) Validate(table *schemas.Table) error {
	if pn.preload != nil {
		for _, col := range pn.preload.cols {
			if table.GetColumn(col) == nil {
				return fmt.Errorf("preload: missing col %s in table %s", col, table.Name)
			}
		}
	}
	for name, node := range pn.children {
		column := table.GetColumn(name)
		if column == nil {
			return fmt.Errorf("preload: missing field %s in struct %s", name, table.Type.Name())
		}
		if column.Association == nil {
			return fmt.Errorf("preload: missing association in field %s", name)
		}
		if column.Association.JoinTable == nil && len(column.Association.SourceCol) > 0 {
			pn.ExtraCols = append(pn.ExtraCols, column.Association.SourceCol)
		}
		if len(column.Association.TargetCol) > 0 {
			pn.ExtraCols = append(pn.ExtraCols, table.PrimaryKeys[0]) // pk might be added many times, but that's ok
		}
		if column.Association.JoinTable == nil && len(column.Association.TargetCol) > 0 {
			node.ExtraCols = append(node.ExtraCols, column.Association.TargetCol)
		}
		node.association = column.Association
		if err := node.Validate(column.Association.RefTable); err != nil {
			return err
		}
	}
	return nil
}

func (pn *PreloadNode) Compute(session *Session, ownMap reflect.Value) error {
	for _, node := range pn.children {
		if err := node.compute(session, ownMap, reflect.Value{}); err != nil {
			return err
		}
	}
	return nil
}

func (pn *PreloadNode) compute(session *Session, ownMap, pruneMap reflect.Value) error {
	// non-root node: pn.association is not nil
	if err := pn.association.ValidateOwnMap(ownMap); err != nil {
		return err
	}

	var joinMap reflect.Value
	cond := pn.association.GetCond(ownMap)
	if pn.association.JoinTable != nil {
		var err error
		cond, joinMap, err = pn.preloadJoin(session, cond)
		if err != nil {
			return err
		}
	}

	refMap := pn.association.MakeRefMap()
	preloadSession := session.Engine().Cols(pn.ExtraCols...).Where(cond)
	if pn.preload != nil {
		preloadSession.Cols(pn.preload.cols...).Where(pn.preload.cond)
	}
	if err := preloadSession.Find(refMap.Interface()); err != nil {
		return err
	}

	var refPruneMap reflect.Value
	if len(pn.children) > 0 && !(pn.preload != nil && (len(pn.preload.cols) > 0 || pn.preload.noPrune)) {
		refPruneMap = reflect.MakeMap(reflect.MapOf(refMap.Type().Key(), reflect.TypeOf(true)))
		refIter := refMap.MapRange()
		for refIter.Next() {
			refPruneMap.SetMapIndex(refIter.Key(), reflect.ValueOf(true))
		}
	}

	for _, node := range pn.children {
		if err := node.compute(session, refMap, refPruneMap); err != nil {
			return err
		}
	}

	if refPruneMap.IsValid() {
		pruneIter := refPruneMap.MapRange()
		for pruneIter.Next() {
			refMap.SetMapIndex(pruneIter.Key(), reflect.Value{})
		}
	}

	pn.association.Link(ownMap, refMap, pruneMap, joinMap)
	return nil
}

func (pn *PreloadNode) preloadJoin(session *Session, cond builder.Cond) (builder.Cond, reflect.Value, error) {
	joinSlicePtr := pn.association.MakeJoinSlice()
	if err := session.Engine().
		Table(pn.association.JoinTable.Name).Where(cond).
		Cols(pn.association.SourceCol, pn.association.TargetCol).
		Find(joinSlicePtr.Interface()); err != nil {
		return nil, reflect.Value{}, err
	}
	joinSlice := joinSlicePtr.Elem()

	joinMap := pn.association.MakeJoinMap()
	for i := 0; i < joinSlice.Len(); i++ {
		entry := joinSlice.Index(i)
		pkSlice := joinMap.MapIndex(entry.Field(1))
		if !pkSlice.IsValid() {
			pkSlice = reflect.MakeSlice(reflect.SliceOf(pn.association.OwnPkType), 0, 0)
		}
		joinMap.SetMapIndex(entry.Field(1), reflect.Append(pkSlice, entry.Field(0)))
	}

	var refPks []interface{}
	iter := joinMap.MapRange()
	joinMap.MapKeys()
	for iter.Next() {
		refPks = append(refPks, iter.Key().Interface())
	}
	cond = builder.In(pn.association.RefTable.PrimaryKeys[0], refPks)
	return cond, joinMap, nil
}