entc/gen: move eager-loading to method (#2790)

This is a preparation work for 'WithNamed<E>' API
This commit is contained in:
Ariel Mashraki
2022-07-23 23:46:02 +03:00
committed by GitHub
parent 43ceed9b6f
commit a2b18f24f0
85 changed files with 7066 additions and 5678 deletions

View File

@@ -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) {