Skip to content

Commit 3d1cef3

Browse files
ggreerggreer
authored
Add actual bulk insertion. Hook it up to PutGrants(). (#269)
* Add actual bulk insertion. Hook it up to PutGrants(). * Remove page size from store requests. These were added when our default page size was 100. * Chunk rows in bulk inserts. * Use bulkPutConnectorObject for all Put operations.
1 parent ef1106d commit 3d1cef3

File tree

6 files changed

+79
-85
lines changed

6 files changed

+79
-85
lines changed

‎pkg/dotc1z/entitlements.go‎

Lines changed: 9 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -84,22 +84,15 @@ func (c *C1File) GetEntitlement(ctx context.Context, request *reader_v2.Entitlem
8484
}
8585

8686
func (c *C1File) PutEntitlements(ctx context.Context, entitlementObjs ...*v2.Entitlement) error {
87-
err := c.db.WithTx(func(tx *goqu.TxDatabase) error {
88-
err := bulkPutConnectorObjectTx(ctx, c, tx, entitlements.Name(),
89-
func(entitlement *v2.Entitlement) (goqu.Record, error) {
90-
return goqu.Record{
91-
"resource_id": entitlement.Resource.Id.Resource,
92-
"resource_type_id": entitlement.Resource.Id.ResourceType,
93-
}, nil
94-
},
95-
entitlementObjs...,
96-
)
97-
if err != nil {
98-
return err
99-
}
100-
101-
return nil
102-
})
87+
err := bulkPutConnectorObject(ctx, c, entitlements.Name(),
88+
func(entitlement *v2.Entitlement) (goqu.Record, error) {
89+
return goqu.Record{
90+
"resource_id": entitlement.Resource.Id.Resource,
91+
"resource_type_id": entitlement.Resource.Id.ResourceType,
92+
}, nil
93+
},
94+
entitlementObjs...,
95+
)
10396
if err != nil {
10497
return err
10598
}

‎pkg/dotc1z/grants.go‎

Lines changed: 12 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -168,24 +168,18 @@ func (c *C1File) ListGrantsForResourceType(
168168
}
169169

170170
func (c *C1File) PutGrants(ctx context.Context, bulkGrants ...*v2.Grant) error {
171-
err := c.db.WithTx(func(tx *goqu.TxDatabase) error {
172-
err := bulkPutConnectorObjectTx(ctx, c, tx, grants.Name(),
173-
func(grant *v2.Grant) (goqu.Record, error) {
174-
return goqu.Record{
175-
"resource_type_id": grant.Entitlement.Resource.Id.ResourceType,
176-
"resource_id": grant.Entitlement.Resource.Id.Resource,
177-
"entitlement_id": grant.Entitlement.Id,
178-
"principal_resource_type_id": grant.Principal.Id.ResourceType,
179-
"principal_resource_id": grant.Principal.Id.Resource,
180-
}, nil
181-
},
182-
bulkGrants...,
183-
)
184-
if err != nil {
185-
return err
186-
}
187-
return nil
188-
})
171+
err := bulkPutConnectorObject(ctx, c, grants.Name(),
172+
func(grant *v2.Grant) (goqu.Record, error) {
173+
return goqu.Record{
174+
"resource_type_id": grant.Entitlement.Resource.Id.ResourceType,
175+
"resource_id": grant.Entitlement.Resource.Id.Resource,
176+
"entitlement_id": grant.Entitlement.Id,
177+
"principal_resource_type_id": grant.Principal.Id.ResourceType,
178+
"principal_resource_id": grant.Principal.Id.Resource,
179+
}, nil
180+
},
181+
bulkGrants...,
182+
)
189183
if err != nil {
190184
return err
191185
}

‎pkg/dotc1z/resouce_types.go‎

Lines changed: 6 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -79,18 +79,12 @@ func (c *C1File) GetResourceType(ctx context.Context, request *reader_v2.Resourc
7979
}
8080

8181
func (c *C1File) PutResourceTypes(ctx context.Context, resourceTypesObjs ...*v2.ResourceType) error {
82-
err := c.db.WithTx(func(tx *goqu.TxDatabase) error {
83-
err := bulkPutConnectorObjectTx(ctx, c, tx, resourceTypes.Name(),
84-
func(resource *v2.ResourceType) (goqu.Record, error) {
85-
return nil, nil
86-
},
87-
resourceTypesObjs...,
88-
)
89-
if err != nil {
90-
return err
91-
}
92-
return nil
93-
})
82+
err := bulkPutConnectorObject(ctx, c, resourceTypes.Name(),
83+
func(resource *v2.ResourceType) (goqu.Record, error) {
84+
return nil, nil
85+
},
86+
resourceTypesObjs...,
87+
)
9488
if err != nil {
9589
return err
9690
}

‎pkg/dotc1z/resources.go‎

Lines changed: 15 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -98,27 +98,21 @@ func (c *C1File) GetResource(ctx context.Context, request *reader_v2.ResourcesRe
9898
}
9999

100100
func (c *C1File) PutResources(ctx context.Context, resourceObjs ...*v2.Resource) error {
101-
err := c.db.WithTx(func(tx *goqu.TxDatabase) error {
102-
err := bulkPutConnectorObjectTx(ctx, c, tx, resources.Name(),
103-
func(resource *v2.Resource) (goqu.Record, error) {
104-
fields := goqu.Record{
105-
"resource_type_id": resource.Id.ResourceType,
106-
"external_id": fmt.Sprintf("%s:%s", resource.Id.ResourceType, resource.Id.Resource),
107-
}
108-
109-
if resource.ParentResourceId != nil {
110-
fields["parent_resource_type_id"] = resource.ParentResourceId.ResourceType
111-
fields["parent_resource_id"] = resource.ParentResourceId.Resource
112-
}
113-
return fields, nil
114-
},
115-
resourceObjs...,
116-
)
117-
if err != nil {
118-
return err
119-
}
120-
return nil
121-
})
101+
err := bulkPutConnectorObject(ctx, c, resources.Name(),
102+
func(resource *v2.Resource) (goqu.Record, error) {
103+
fields := goqu.Record{
104+
"resource_type_id": resource.Id.ResourceType,
105+
"external_id": fmt.Sprintf("%s:%s", resource.Id.ResourceType, resource.Id.Resource),
106+
}
107+
108+
if resource.ParentResourceId != nil {
109+
fields["parent_resource_type_id"] = resource.ParentResourceId.ResourceType
110+
fields["parent_resource_id"] = resource.ParentResourceId.Resource
111+
}
112+
return fields, nil
113+
},
114+
resourceObjs...,
115+
)
122116
if err != nil {
123117
return err
124118
}

‎pkg/dotc1z/sql_helpers.go‎

Lines changed: 29 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -233,20 +233,21 @@ func (c *C1File) listConnectorObjects(ctx context.Context, tableName string, req
233233

234234
var protoMarshaler = proto.MarshalOptions{Deterministic: true}
235235

236-
func bulkPutConnectorObjectTx[T proto.Message](ctx context.Context, c *C1File,
237-
tx *goqu.TxDatabase,
236+
func bulkPutConnectorObject[T proto.Message](ctx context.Context, c *C1File,
238237
tableName string,
239238
extractFields func(m T) (goqu.Record, error),
240239
msgs ...T) error {
240+
if len(msgs) == 0 {
241+
return nil
242+
}
243+
241244
err := c.validateSyncDb(ctx)
242245
if err != nil {
243246
return err
244247
}
245248

246-
baseQ := tx.Insert(tableName).Prepared(true)
247-
baseQ = baseQ.OnConflict(goqu.DoUpdate("external_id, sync_id", goqu.C("data").Set(goqu.I("EXCLUDED.data"))))
248-
249-
for _, m := range msgs {
249+
rows := make([]*goqu.Record, len(msgs))
250+
for i, m := range msgs {
250251
messageBlob, err := protoMarshaler.Marshal(m)
251252
if err != nil {
252253
return err
@@ -270,16 +271,35 @@ func bulkPutConnectorObjectTx[T proto.Message](ctx context.Context, c *C1File,
270271
fields["data"] = messageBlob
271272
fields["sync_id"] = c.currentSyncID
272273
fields["discovered_at"] = time.Now().Format("2006-01-02 15:04:05.999999999")
273-
q := baseQ.Rows(fields)
274-
query, args, err := q.ToSQL()
274+
rows[i] = &fields
275+
}
276+
chunkSize := 100
277+
chunks := len(rows) / chunkSize
278+
if len(rows)%chunkSize != 0 {
279+
chunks++
280+
}
281+
282+
for i := 0; i < chunks; i++ {
283+
start := i * chunkSize
284+
end := (i + 1) * chunkSize
285+
if end > len(rows) {
286+
end = len(rows)
287+
}
288+
chunkedRows := rows[start:end]
289+
query, args, err := c.db.Insert(tableName).
290+
OnConflict(goqu.DoUpdate("external_id, sync_id", goqu.C("data").Set(goqu.I("EXCLUDED.data")))).
291+
Rows(chunkedRows).
292+
Prepared(true).
293+
ToSQL()
275294
if err != nil {
276295
return err
277296
}
278-
_, err = tx.Exec(query, args...)
297+
_, err = c.db.Exec(query, args...)
279298
if err != nil {
280299
return err
281300
}
282301
}
302+
283303
return nil
284304
}
285305

‎pkg/sync/syncer.go‎

Lines changed: 8 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1120,7 +1120,6 @@ func (s *syncer) fetchEtaggedGrantsForResource(
11201120
Resource: resource,
11211121
Annotations: storeAnnos,
11221122
PageToken: npt,
1123-
PageSize: 1000,
11241123
})
11251124
if err != nil {
11261125
return nil, false, err
@@ -1273,14 +1272,14 @@ func (s *syncer) runGrantExpandActions(ctx context.Context) (bool, error) {
12731272
// Fetch a page of source grants
12741273
sourceGrants, err := s.store.ListGrantsForEntitlement(ctx, &reader_v2.GrantsReaderServiceListGrantsForEntitlementRequest{
12751274
Entitlement: sourceEntitlement.GetEntitlement(),
1276-
PageSize: 1000,
12771275
PageToken: action.PageToken,
12781276
})
12791277
if err != nil {
12801278
l.Error("runGrantExpandActions: error fetching source grants", zap.Error(err))
12811279
return false, fmt.Errorf("runGrantExpandActions: error fetching source grants: %w", err)
12821280
}
12831281

1282+
var newGrants []*v2.Grant = make([]*v2.Grant, 0)
12841283
for _, sourceGrant := range sourceGrants.List {
12851284
// Skip this grant if it is not for a resource type we care about
12861285
if len(action.ResourceTypeIDs) > 0 {
@@ -1322,7 +1321,6 @@ func (s *syncer) runGrantExpandActions(ctx context.Context) (bool, error) {
13221321
req := &reader_v2.GrantsReaderServiceListGrantsForEntitlementRequest{
13231322
Entitlement: descendantEntitlement.GetEntitlement(),
13241323
PrincipalId: sourceGrant.GetPrincipal().GetId(),
1325-
PageSize: 1000,
13261324
PageToken: pageToken,
13271325
Annotations: nil,
13281326
}
@@ -1383,13 +1381,14 @@ func (s *syncer) runGrantExpandActions(ctx context.Context) (bool, error) {
13831381
zap.String("grant_id", descendantGrant.GetId()),
13841382
zap.Any("sources", sources),
13851383
)
1386-
1387-
err = s.store.PutGrants(ctx, descendantGrant)
1388-
if err != nil {
1389-
l.Error("runGrantExpandActions: error updating descendant grant", zap.Error(err))
1390-
return false, fmt.Errorf("runGrantExpandActions: error updating descendant grant: %w", err)
1391-
}
13921384
}
1385+
newGrants = append(newGrants, descendantGrants...)
1386+
}
1387+
1388+
err = s.store.PutGrants(ctx, newGrants...)
1389+
if err != nil {
1390+
l.Error("runGrantExpandActions: error updating descendant grants", zap.Error(err))
1391+
return false, fmt.Errorf("runGrantExpandActions: error updating descendant grants: %w", err)
13931392
}
13941393

13951394
// If we have no more pages of work, pop the action off the stack and mark this edge in the graph as done

0 commit comments

Comments
 (0)