Skip to content

Commit fc49c78

Browse files
harshachintaharshachinta
authored
feat(spanner): add databoost property for batch transactions (#8152)
* feat(spanner): add databoost property for batch transactions * feat(spanner): add databoost property * feat(spanner): disable databoost property in tests * feat(spanner): comment refactor
1 parent 005d2df commit fc49c78

File tree

3 files changed

+86
-60
lines changed

3 files changed

+86
-60
lines changed

‎spanner/batch.go‎

Lines changed: 36 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -150,13 +150,14 @@ func (t *BatchReadOnlyTransaction) PartitionReadUsingIndexWithOptions(ctx contex
150150
}
151151
// Prepare ReadRequest.
152152
req := &sppb.ReadRequest{
153-
Session: sid,
154-
Transaction: ts,
155-
Table: table,
156-
Index: index,
157-
Columns: columns,
158-
KeySet: kset,
159-
RequestOptions: createRequestOptions(readOptions.Priority, readOptions.RequestTag, ""),
153+
Session: sid,
154+
Transaction: ts,
155+
Table: table,
156+
Index: index,
157+
Columns: columns,
158+
KeySet: kset,
159+
RequestOptions: createRequestOptions(readOptions.Priority, readOptions.RequestTag, ""),
160+
DataBoostEnabled: readOptions.DataBoostEnabled,
160161
}
161162
// Generate partitions.
162163
for _, p := range resp.GetPartitions() {
@@ -212,13 +213,14 @@ func (t *BatchReadOnlyTransaction) partitionQuery(ctx context.Context, statement
212213

213214
// prepare ExecuteSqlRequest
214215
r := &sppb.ExecuteSqlRequest{
215-
Session: sid,
216-
Transaction: ts,
217-
Sql: statement.SQL,
218-
Params: params,
219-
ParamTypes: paramTypes,
220-
QueryOptions: qOpts.Options,
221-
RequestOptions: createRequestOptions(qOpts.Priority, qOpts.RequestTag, ""),
216+
Session: sid,
217+
Transaction: ts,
218+
Sql: statement.SQL,
219+
Params: params,
220+
ParamTypes: paramTypes,
221+
QueryOptions: qOpts.Options,
222+
RequestOptions: createRequestOptions(qOpts.Priority, qOpts.RequestTag, ""),
223+
DataBoostEnabled: qOpts.DataBoostEnabled,
222224
}
223225

224226
// generate Partitions
@@ -308,15 +310,16 @@ func (t *BatchReadOnlyTransaction) Execute(ctx context.Context, p *Partition) *R
308310
if p.rreq != nil {
309311
rpc = func(ctx context.Context, resumeToken []byte) (streamingReceiver, error) {
310312
client, err := client.StreamingRead(ctx, &sppb.ReadRequest{
311-
Session: p.rreq.Session,
312-
Transaction: p.rreq.Transaction,
313-
Table: p.rreq.Table,
314-
Index: p.rreq.Index,
315-
Columns: p.rreq.Columns,
316-
KeySet: p.rreq.KeySet,
317-
PartitionToken: p.pt,
318-
RequestOptions: p.rreq.RequestOptions,
319-
ResumeToken: resumeToken,
313+
Session: p.rreq.Session,
314+
Transaction: p.rreq.Transaction,
315+
Table: p.rreq.Table,
316+
Index: p.rreq.Index,
317+
Columns: p.rreq.Columns,
318+
KeySet: p.rreq.KeySet,
319+
PartitionToken: p.pt,
320+
RequestOptions: p.rreq.RequestOptions,
321+
ResumeToken: resumeToken,
322+
DataBoostEnabled: p.rreq.DataBoostEnabled,
320323
})
321324
if err != nil {
322325
return client, err
@@ -332,15 +335,16 @@ func (t *BatchReadOnlyTransaction) Execute(ctx context.Context, p *Partition) *R
332335
} else {
333336
rpc = func(ctx context.Context, resumeToken []byte) (streamingReceiver, error) {
334337
client, err := client.ExecuteStreamingSql(ctx, &sppb.ExecuteSqlRequest{
335-
Session: p.qreq.Session,
336-
Transaction: p.qreq.Transaction,
337-
Sql: p.qreq.Sql,
338-
Params: p.qreq.Params,
339-
ParamTypes: p.qreq.ParamTypes,
340-
QueryOptions: p.qreq.QueryOptions,
341-
PartitionToken: p.pt,
342-
RequestOptions: p.qreq.RequestOptions,
343-
ResumeToken: resumeToken,
338+
Session: p.qreq.Session,
339+
Transaction: p.qreq.Transaction,
340+
Sql: p.qreq.Sql,
341+
Params: p.qreq.Params,
342+
ParamTypes: p.qreq.ParamTypes,
343+
QueryOptions: p.qreq.QueryOptions,
344+
PartitionToken: p.pt,
345+
RequestOptions: p.qreq.RequestOptions,
346+
ResumeToken: resumeToken,
347+
DataBoostEnabled: p.qreq.DataBoostEnabled,
344348
})
345349
if err != nil {
346350
return client, err

‎spanner/integration_test.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3141,7 +3141,7 @@ func TestIntegration_BatchQuery(t *testing.T) {
31413141
t.Fatal(err)
31423142
}
31433143
defer txn.Cleanup(ctx)
3144-
if partitions, err = txn.PartitionQuery(ctx, stmt, PartitionOptions{0, 3}); err != nil {
3144+
if partitions, err = txn.PartitionQueryWithOptions(ctx, stmt, PartitionOptions{0, 3}, QueryOptions{DataBoostEnabled: false}); err != nil {
31453145
t.Fatal(err)
31463146
}
31473147

@@ -3224,7 +3224,7 @@ func TestIntegration_BatchRead(t *testing.T) {
32243224
t.Fatal(err)
32253225
}
32263226
defer txn.Cleanup(ctx)
3227-
if partitions, err = txn.PartitionRead(ctx, "test", AllKeys(), simpleDBTableColumns, PartitionOptions{0, 3}); err != nil {
3227+
if partitions, err = txn.PartitionReadWithOptions(ctx, "test", AllKeys(), simpleDBTableColumns, PartitionOptions{0, 3}, ReadOptions{DataBoostEnabled: false}); err != nil {
32283228
t.Fatal(err)
32293229
}
32303230

‎spanner/transaction.go‎

Lines changed: 48 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -165,16 +165,21 @@ type ReadOptions struct {
165165

166166
// The request tag to use for this request.
167167
RequestTag string
168+
169+
// If this is for a partitioned read and DataBoostEnabled field is set to true, the request will be executed
170+
// via Spanner independent compute resources. Setting this option for regular read operations has no effect.
171+
DataBoostEnabled bool
168172
}
169173

170174
// merge combines two ReadOptions that the input parameter will have higher
171175
// order of precedence.
172176
func (ro ReadOptions) merge(opts ReadOptions) ReadOptions {
173177
merged := ReadOptions{
174-
Index: ro.Index,
175-
Limit: ro.Limit,
176-
Priority: ro.Priority,
177-
RequestTag: ro.RequestTag,
178+
Index: ro.Index,
179+
Limit: ro.Limit,
180+
Priority: ro.Priority,
181+
RequestTag: ro.RequestTag,
182+
DataBoostEnabled: ro.DataBoostEnabled,
178183
}
179184
if opts.Index != "" {
180185
merged.Index = opts.Index
@@ -188,6 +193,9 @@ func (ro ReadOptions) merge(opts ReadOptions) ReadOptions {
188193
if opts.RequestTag != "" {
189194
merged.RequestTag = opts.RequestTag
190195
}
196+
if opts.DataBoostEnabled {
197+
merged.DataBoostEnabled = opts.DataBoostEnabled
198+
}
191199
return merged
192200
}
193201

@@ -218,13 +226,17 @@ func (t *txReadOnly) ReadWithOptions(ctx context.Context, table string, keys Key
218226
limit := t.ro.Limit
219227
prio := t.ro.Priority
220228
requestTag := t.ro.RequestTag
229+
dataBoostEnabled := t.ro.DataBoostEnabled
221230
if opts != nil {
222231
index = opts.Index
223232
if opts.Limit > 0 {
224233
limit = opts.Limit
225234
}
226235
prio = opts.Priority
227236
requestTag = opts.RequestTag
237+
if opts.DataBoostEnabled {
238+
dataBoostEnabled = opts.DataBoostEnabled
239+
}
228240
}
229241
var setTransactionID func(transactionID)
230242
if _, ok := ts.Selector.(*sppb.TransactionSelector_Begin); ok {
@@ -238,15 +250,16 @@ func (t *txReadOnly) ReadWithOptions(ctx context.Context, table string, keys Key
238250
func(ctx context.Context, resumeToken []byte) (streamingReceiver, error) {
239251
client, err := client.StreamingRead(ctx,
240252
&sppb.ReadRequest{
241-
Session: t.sh.getID(),
242-
Transaction: t.getTransactionSelector(),
243-
Table: table,
244-
Index: index,
245-
Columns: columns,
246-
KeySet: kset,
247-
ResumeToken: resumeToken,
248-
Limit: int64(limit),
249-
RequestOptions: createRequestOptions(prio, requestTag, t.txOpts.TransactionTag),
253+
Session: t.sh.getID(),
254+
Transaction: t.getTransactionSelector(),
255+
Table: table,
256+
Index: index,
257+
Columns: columns,
258+
KeySet: kset,
259+
ResumeToken: resumeToken,
260+
Limit: int64(limit),
261+
RequestOptions: createRequestOptions(prio, requestTag, t.txOpts.TransactionTag),
262+
DataBoostEnabled: dataBoostEnabled,
250263
})
251264
if err != nil {
252265
if _, ok := t.getTransactionSelector().GetSelector().(*sppb.TransactionSelector_Begin); ok {
@@ -357,16 +370,21 @@ type QueryOptions struct {
357370

358371
// The request tag to use for this request.
359372
RequestTag string
373+
374+
// If this is for a partitioned query and DataBoostEnabled field is set to true, the request will be executed
375+
// via Spanner independent compute resources. Setting this option for regular query operations has no effect.
376+
DataBoostEnabled bool
360377
}
361378

362379
// merge combines two QueryOptions that the input parameter will have higher
363380
// order of precedence.
364381
func (qo QueryOptions) merge(opts QueryOptions) QueryOptions {
365382
merged := QueryOptions{
366-
Mode: qo.Mode,
367-
Options: &sppb.ExecuteSqlRequest_QueryOptions{},
368-
RequestTag: qo.RequestTag,
369-
Priority: qo.Priority,
383+
Mode: qo.Mode,
384+
Options: &sppb.ExecuteSqlRequest_QueryOptions{},
385+
RequestTag: qo.RequestTag,
386+
Priority: qo.Priority,
387+
DataBoostEnabled: qo.DataBoostEnabled,
370388
}
371389
if opts.Mode != nil {
372390
merged.Mode = opts.Mode
@@ -377,6 +395,9 @@ func (qo QueryOptions) merge(opts QueryOptions) QueryOptions {
377395
if opts.Priority != sppb.RequestOptions_PRIORITY_UNSPECIFIED {
378396
merged.Priority = opts.Priority
379397
}
398+
if opts.DataBoostEnabled {
399+
merged.DataBoostEnabled = opts.DataBoostEnabled
400+
}
380401
proto.Merge(merged.Options, qo.Options)
381402
proto.Merge(merged.Options, opts.Options)
382403
return merged
@@ -517,15 +538,16 @@ func (t *txReadOnly) prepareExecuteSQL(ctx context.Context, stmt Statement, opti
517538
mode = *options.Mode
518539
}
519540
req := &sppb.ExecuteSqlRequest{
520-
Session: sid,
521-
Transaction: ts,
522-
Sql: stmt.SQL,
523-
QueryMode: mode,
524-
Seqno: atomic.AddInt64(&t.sequenceNumber, 1),
525-
Params: params,
526-
ParamTypes: paramTypes,
527-
QueryOptions: options.Options,
528-
RequestOptions: createRequestOptions(options.Priority, options.RequestTag, t.txOpts.TransactionTag),
541+
Session: sid,
542+
Transaction: ts,
543+
Sql: stmt.SQL,
544+
QueryMode: mode,
545+
Seqno: atomic.AddInt64(&t.sequenceNumber, 1),
546+
Params: params,
547+
ParamTypes: paramTypes,
548+
QueryOptions: options.Options,
549+
RequestOptions: createRequestOptions(options.Priority, options.RequestTag, t.txOpts.TransactionTag),
550+
DataBoostEnabled: options.DataBoostEnabled,
529551
}
530552
return req, sh, nil
531553
}

0 commit comments

Comments
 (0)