mirror of
https://github.com/ent/ent.git
synced 2026-05-24 09:31:56 +03:00
entc/gen: move eager-loading to method (#2790)
This is a preparation work for 'WithNamed<E>' API
This commit is contained in:
@@ -416,59 +416,74 @@ func (fq *FileQuery) sqlAll(ctx context.Context, hooks ...queryHook) ([]*File, e
|
||||
if len(nodes) == 0 {
|
||||
return nodes, nil
|
||||
}
|
||||
|
||||
if query := fq.withParent; query != nil {
|
||||
ids := make([]int, 0, len(nodes))
|
||||
nodeids := make(map[int][]*File)
|
||||
for i := range nodes {
|
||||
fk := nodes[i].ParentID
|
||||
if _, ok := nodeids[fk]; !ok {
|
||||
ids = append(ids, fk)
|
||||
}
|
||||
nodeids[fk] = append(nodeids[fk], nodes[i])
|
||||
}
|
||||
query.Where(file.IDIn(ids...))
|
||||
neighbors, err := query.All(ctx)
|
||||
if err != nil {
|
||||
if err := fq.loadParent(ctx, query, nodes, nil,
|
||||
func(n *File, e *File) { n.Edges.Parent = e }); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
for _, n := range neighbors {
|
||||
nodes, ok := nodeids[n.ID]
|
||||
if !ok {
|
||||
return nil, fmt.Errorf(`unexpected foreign-key "parent_id" returned %v`, n.ID)
|
||||
}
|
||||
for i := range nodes {
|
||||
nodes[i].Edges.Parent = n
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if query := fq.withChildren; query != nil {
|
||||
fks := make([]driver.Value, 0, len(nodes))
|
||||
nodeids := make(map[int]*File)
|
||||
for i := range nodes {
|
||||
fks = append(fks, nodes[i].ID)
|
||||
nodeids[nodes[i].ID] = nodes[i]
|
||||
nodes[i].Edges.Children = []*File{}
|
||||
}
|
||||
query.Where(predicate.File(func(s *sql.Selector) {
|
||||
s.Where(sql.InValues(file.ChildrenColumn, fks...))
|
||||
}))
|
||||
neighbors, err := query.All(ctx)
|
||||
if err != nil {
|
||||
if err := fq.loadChildren(ctx, query, nodes,
|
||||
func(n *File) { n.Edges.Children = []*File{} },
|
||||
func(n *File, e *File) { n.Edges.Children = append(n.Edges.Children, e) }); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
for _, n := range neighbors {
|
||||
fk := n.ParentID
|
||||
node, ok := nodeids[fk]
|
||||
if !ok {
|
||||
return nil, fmt.Errorf(`unexpected foreign-key "parent_id" returned %v for node %v`, fk, n.ID)
|
||||
}
|
||||
node.Edges.Children = append(node.Edges.Children, n)
|
||||
}
|
||||
return nodes, nil
|
||||
}
|
||||
|
||||
func (fq *FileQuery) loadParent(ctx context.Context, query *FileQuery, nodes []*File, init func(*File), assign func(*File, *File)) error {
|
||||
ids := make([]int, 0, len(nodes))
|
||||
nodeids := make(map[int][]*File)
|
||||
for i := range nodes {
|
||||
fk := nodes[i].ParentID
|
||||
if _, ok := nodeids[fk]; !ok {
|
||||
ids = append(ids, fk)
|
||||
}
|
||||
nodeids[fk] = append(nodeids[fk], nodes[i])
|
||||
}
|
||||
query.Where(file.IDIn(ids...))
|
||||
neighbors, err := query.All(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
for _, n := range neighbors {
|
||||
nodes, ok := nodeids[n.ID]
|
||||
if !ok {
|
||||
return fmt.Errorf(`unexpected foreign-key "parent_id" returned %v`, n.ID)
|
||||
}
|
||||
for i := range nodes {
|
||||
assign(nodes[i], n)
|
||||
}
|
||||
}
|
||||
|
||||
return nodes, nil
|
||||
return nil
|
||||
}
|
||||
func (fq *FileQuery) loadChildren(ctx context.Context, query *FileQuery, nodes []*File, init func(*File), assign func(*File, *File)) error {
|
||||
fks := make([]driver.Value, 0, len(nodes))
|
||||
nodeids := make(map[int]*File)
|
||||
for i := range nodes {
|
||||
fks = append(fks, nodes[i].ID)
|
||||
nodeids[nodes[i].ID] = nodes[i]
|
||||
if init != nil {
|
||||
init(nodes[i])
|
||||
}
|
||||
}
|
||||
query.Where(predicate.File(func(s *sql.Selector) {
|
||||
s.Where(sql.InValues(file.ChildrenColumn, fks...))
|
||||
}))
|
||||
neighbors, err := query.All(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
for _, n := range neighbors {
|
||||
fk := n.ParentID
|
||||
node, ok := nodeids[fk]
|
||||
if !ok {
|
||||
return fmt.Errorf(`unexpected foreign-key "parent_id" returned %v for node %v`, fk, n.ID)
|
||||
}
|
||||
assign(node, n)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (fq *FileQuery) sqlCount(ctx context.Context) (int, error) {
|
||||
|
||||
Reference in New Issue
Block a user