/
hydrator.go
130 lines (123 loc) · 3.34 KB
/
hydrator.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
package query
import (
"context"
"errors"
)
type HasOneHydrator[Root any, Sub any, SubID comparable] interface {
RootToSubID(root Root) (SubID, bool)
GetSubs(ctx context.Context, dbCtx DbContext, ids []SubID) ([]Sub, error)
SubID(sub Sub) SubID
Hydrate(root *Root, sub Sub)
MustSucceed() bool
}
func HydrateHasOne[Root any, Sub any, SubID comparable](h HasOneHydrator[Root, Sub, SubID]) Option {
return func(b OptionBuilder) (OptionBuilder, error) {
return b.Callback(func(ctx context.Context, cbCtx CallbackContext, result any) error {
items, ok := result.([]Root)
if !ok {
return errors.New("invalid result type")
}
return hydrateHasOne(ctx, cbCtx, h, items)
}), nil
}
}
func hydrateHasOne[Root any, Sub any, SubID comparable](
ctx context.Context,
cbCtx CallbackContext,
h HasOneHydrator[Root, Sub, SubID],
roots []Root,
) error {
indexMap := make(map[int]SubID, len(roots))
idMap := make(map[SubID]struct{}, len(roots))
ids := make([]SubID, 0, len(roots))
cbCtx.Lock()
for i, root := range roots {
if subID, ok1 := h.RootToSubID(root); ok1 {
indexMap[i] = subID
if _, ok2 := idMap[subID]; !ok2 {
ids = append(ids, subID)
idMap[subID] = struct{}{}
}
}
}
cbCtx.Unlock()
if len(ids) == 0 {
return nil
}
subs, err := h.GetSubs(ctx, cbCtx, ids)
if err != nil {
return err
}
subsMap := make(map[SubID]Sub, len(subs))
for _, sub := range subs {
subsMap[h.SubID(sub)] = sub
}
cbCtx.Lock()
defer cbCtx.Unlock()
for i, subID := range indexMap {
if sub, ok := subsMap[subID]; ok {
h.Hydrate(&roots[i], sub)
} else if h.MustSucceed() {
return errors.New("failed to hydrateHasOne")
}
}
return nil
}
type HasManyHydrator[Root any, RootID comparable, JoinSub any, Sub any] interface {
RootID(root Root) (RootID, bool)
GetJoinSubs(ctx context.Context, dbCtx DbContext, ids []RootID) ([]JoinSub, error)
JoinSubToRootIDAndSub(j JoinSub) (RootID, Sub)
Hydrate(root *Root, subs []Sub)
}
func HydrateHasMany[Root any, RootID comparable, JoinSub any, Sub any](h HasManyHydrator[Root, RootID, JoinSub, Sub]) Option {
return func(b OptionBuilder) (OptionBuilder, error) {
return b.Callback(func(ctx context.Context, cbCtx CallbackContext, result any) error {
items, ok := result.([]Root)
if !ok {
return errors.New("invalid result type")
}
return hydrateHasMany(ctx, cbCtx, h, items)
}), nil
}
}
func hydrateHasMany[Root any, RootID comparable, JoinSub any, Sub any](
ctx context.Context,
cbCtx CallbackContext,
h HasManyHydrator[Root, RootID, JoinSub, Sub],
roots []Root,
) error {
indexMap := make(map[int]RootID, len(roots))
idMap := make(map[RootID]struct{}, len(roots))
ids := make([]RootID, 0, len(roots))
cbCtx.Lock()
for i, root := range roots {
if rootID, ok1 := h.RootID(root); ok1 {
indexMap[i] = rootID
if _, ok2 := idMap[rootID]; !ok2 {
ids = append(ids, rootID)
idMap[rootID] = struct{}{}
}
}
}
cbCtx.Unlock()
if len(ids) == 0 {
return nil
}
joinSubs, err := h.GetJoinSubs(ctx, cbCtx, ids)
if err != nil {
return err
}
subsMap := make(map[RootID][]Sub, len(roots))
for _, j := range joinSubs {
rootID, sub := h.JoinSubToRootIDAndSub(j)
subsMap[rootID] = append(subsMap[rootID], sub)
}
cbCtx.Lock()
defer cbCtx.Unlock()
for i, rootID := range indexMap {
if thisSubs, ok := subsMap[rootID]; ok {
h.Hydrate(&roots[i], thisSubs)
}
}
return nil
}