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

@@ -388,37 +388,43 @@ func (pq *PetQuery) sqlAll(ctx context.Context, hooks ...queryHook) ([]*Pet, err
if len(nodes) == 0 {
return nodes, nil
}
if query := pq.withOwner; query != nil {
ids := make([]int, 0, len(nodes))
nodeids := make(map[int][]*Pet)
for i := range nodes {
if nodes[i].user_pets == nil {
continue
}
fk := *nodes[i].user_pets
if _, ok := nodeids[fk]; !ok {
ids = append(ids, fk)
}
nodeids[fk] = append(nodeids[fk], nodes[i])
}
query.Where(user.IDIn(ids...))
neighbors, err := query.All(ctx)
if err != nil {
if err := pq.loadOwner(ctx, query, nodes, nil,
func(n *Pet, e *User) { n.Edges.Owner = e }); err != nil {
return nil, err
}
for _, n := range neighbors {
nodes, ok := nodeids[n.ID]
if !ok {
return nil, fmt.Errorf(`unexpected foreign-key "user_pets" returned %v`, n.ID)
}
for i := range nodes {
nodes[i].Edges.Owner = n
}
}
return nodes, nil
}
func (pq *PetQuery) loadOwner(ctx context.Context, query *UserQuery, nodes []*Pet, init func(*Pet), assign func(*Pet, *User)) error {
ids := make([]int, 0, len(nodes))
nodeids := make(map[int][]*Pet)
for i := range nodes {
if nodes[i].user_pets == nil {
continue
}
fk := *nodes[i].user_pets
if _, ok := nodeids[fk]; !ok {
ids = append(ids, fk)
}
nodeids[fk] = append(nodeids[fk], nodes[i])
}
query.Where(user.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 "user_pets" returned %v`, n.ID)
}
for i := range nodes {
assign(nodes[i], n)
}
}
return nodes, nil
return nil
}
func (pq *PetQuery) sqlCount(ctx context.Context) (int, error) {

View File

@@ -381,37 +381,46 @@ func (uq *UserQuery) sqlAll(ctx context.Context, hooks ...queryHook) ([]*User, e
if len(nodes) == 0 {
return nodes, nil
}
if query := uq.withPets; query != nil {
fks := make([]driver.Value, 0, len(nodes))
nodeids := make(map[int]*User)
for i := range nodes {
fks = append(fks, nodes[i].ID)
nodeids[nodes[i].ID] = nodes[i]
nodes[i].Edges.Pets = []*Pet{}
}
query.withFKs = true
query.Where(predicate.Pet(func(s *sql.Selector) {
s.Where(sql.InValues(user.PetsColumn, fks...))
}))
neighbors, err := query.All(ctx)
if err != nil {
if err := uq.loadPets(ctx, query, nodes,
func(n *User) { n.Edges.Pets = []*Pet{} },
func(n *User, e *Pet) { n.Edges.Pets = append(n.Edges.Pets, e) }); err != nil {
return nil, err
}
for _, n := range neighbors {
fk := n.user_pets
if fk == nil {
return nil, fmt.Errorf(`foreign-key "user_pets" is nil for node %v`, n.ID)
}
node, ok := nodeids[*fk]
if !ok {
return nil, fmt.Errorf(`unexpected foreign-key "user_pets" returned %v for node %v`, *fk, n.ID)
}
node.Edges.Pets = append(node.Edges.Pets, n)
}
return nodes, nil
}
func (uq *UserQuery) loadPets(ctx context.Context, query *PetQuery, nodes []*User, init func(*User), assign func(*User, *Pet)) error {
fks := make([]driver.Value, 0, len(nodes))
nodeids := make(map[int]*User)
for i := range nodes {
fks = append(fks, nodes[i].ID)
nodeids[nodes[i].ID] = nodes[i]
if init != nil {
init(nodes[i])
}
}
return nodes, nil
query.withFKs = true
query.Where(predicate.Pet(func(s *sql.Selector) {
s.Where(sql.InValues(user.PetsColumn, fks...))
}))
neighbors, err := query.All(ctx)
if err != nil {
return err
}
for _, n := range neighbors {
fk := n.user_pets
if fk == nil {
return fmt.Errorf(`foreign-key "user_pets" is nil for node %v`, n.ID)
}
node, ok := nodeids[*fk]
if !ok {
return fmt.Errorf(`unexpected foreign-key "user_pets" returned %v for node %v`, *fk, n.ID)
}
assign(node, n)
}
return nil
}
func (uq *UserQuery) sqlCount(ctx context.Context) (int, error) {