Skip to content

Commit 8b8ca9e

Browse files
bhshkhbhshkh
authored
feat(bigtable): Add methods to read ResultRow (#11924)
* feat(bigtable): Add methods to read ResultRow * revert Scan and Unmarshal * revert GetAllByNames * integration tests for GetByName(string) any * test(bigtable): Add unit tests
1 parent 35e0774 commit 8b8ca9e

File tree

6 files changed

+2294
-81
lines changed

6 files changed

+2294
-81
lines changed

‎bigtable/bigtable.go‎

Lines changed: 30 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -438,6 +438,17 @@ type preparedQueryData struct {
438438
// A token may become invalid early due to changes in the data being read, but
439439
// it provides a guideline to refresh query plans asynchronously.
440440
validUntil *timestamppb.Timestamp
441+
442+
Metadata *ResultRowMetadata
443+
}
444+
445+
func (pqd *preparedQueryData) initializeMetadataAndMap() error {
446+
rrMetadata, err := newResultRowMetadata(pqd.metadata)
447+
if err != nil {
448+
return err
449+
}
450+
pqd.Metadata = rrMetadata
451+
return nil
441452
}
442453

443454
// PrepareOption can be passed while preparing a query statement.
@@ -481,6 +492,9 @@ func (c *Client) prepareStatement(ctx context.Context, mt *builtinMetricsTracer,
481492
if v == nil {
482493
return nil, errors.New("bigtable: invalid SQLType: nil")
483494
}
495+
if !v.isValidPrepareParamType() {
496+
return nil, fmt.Errorf("bigtable: %T cannot be used as parameter type", v)
497+
}
484498
tpb, err := v.typeProto()
485499
if err != nil {
486500
return nil, err
@@ -524,17 +538,23 @@ func (c *Client) prepareStatement(ctx context.Context, mt *builtinMetricsTracer,
524538
// Allowed parameter value types are []byte, string, int64, float32, float64, bool,
525539
// time.Time, civil.Date, array, slice and nil
526540
func (ps *PreparedStatement) Bind(values map[string]any) (*BoundStatement, error) {
527-
bs := BoundStatement{
528-
ps: ps,
529-
params: map[string]*btpb.Value{},
541+
if ps == nil {
542+
return nil, errors.New("bigtable: nil prepared statement")
543+
}
544+
// check that every parameter is bound
545+
for paramName := range ps.paramTypes {
546+
_, found := values[paramName]
547+
if !found {
548+
return nil, fmt.Errorf("bigtable: parameter %q not bound in call to Bind", paramName)
549+
}
530550
}
531551

532552
boundParams := map[string]*btpb.Value{}
533553
for paramName, paramVal := range values {
534554
// Validate that the parameter was specified during prepare
535555
psType, found := ps.paramTypes[paramName]
536556
if !found {
537-
return &bs, errors.New("bigtable: no parameter with name " + paramName + " in prepared statement")
557+
return nil, errors.New("bigtable: no parameter with name " + paramName + " in prepared statement")
538558
}
539559

540560
// Convert value specified by user to *btpb.Value
@@ -544,13 +564,6 @@ func (ps *PreparedStatement) Bind(values map[string]any) (*BoundStatement, error
544564
}
545565
boundParams[paramName] = pbVal
546566
}
547-
// check that every parameter is bound
548-
for paramName := range ps.paramTypes {
549-
_, found := boundParams[paramName]
550-
if !found {
551-
return &bs, errors.New("bigtable: parameter " + paramName + " not bound in prepared statement")
552-
}
553-
}
554567

555568
return &BoundStatement{
556569
ps: ps,
@@ -626,12 +639,6 @@ type BoundStatement struct {
626639
params map[string]*btpb.Value
627640
}
628641

629-
// ResultRow represents a single row in the result set returned on executing a GoogleSQL query.
630-
type ResultRow struct {
631-
values []*btpb.Value
632-
metadata *btpb.ResultSetMetadata
633-
}
634-
635642
// ExecuteOption is an optional argument to Execute.
636643
type ExecuteOption interface{}
637644

@@ -724,7 +731,6 @@ func (bs *BoundStatement) execute(ctx context.Context, f func(ResultRow) bool, m
724731
Params: bs.params,
725732
}
726733
stream, err := bs.ps.c.client.ExecuteQuery(ctx, req)
727-
728734
if err != nil {
729735
prevError = err
730736
return err
@@ -803,6 +809,7 @@ func (bs *BoundStatement) execute(ctx context.Context, f func(ResultRow) bool, m
803809
if !receivedResumeToken {
804810
// first ResumeToken received
805811
finalizedStmt = candFinalizedStmt
812+
finalizedStmt.initializeMetadataAndMap()
806813
receivedResumeToken = true
807814
}
808815

@@ -828,11 +835,12 @@ func (bs *BoundStatement) execute(ctx context.Context, f func(ResultRow) bool, m
828835

829836
completeRowValues, valuesBuffer = valuesBuffer[0:numCols], valuesBuffer[numCols:]
830837

831-
rr := ResultRow{
832-
values: completeRowValues,
833-
metadata: finalizedStmt.metadata,
838+
// Construct ResultRow
839+
rr, err := newResultRow(completeRowValues, finalizedStmt.metadata, finalizedStmt.Metadata)
840+
if err != nil {
841+
return err
834842
}
835-
continueReading := f(rr)
843+
continueReading := f(*rr)
836844
if !continueReading {
837845
// Cancel and drain stream.
838846
cancel()

‎bigtable/bigtable_test.go‎

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1573,7 +1573,9 @@ func TestAnySQLTypeToPbVal(t *testing.T) {
15731573
},
15741574
},
15751575
Kind: &btpb.Value_ArrayValue{
1576-
ArrayValue: &btpb.ArrayValue{},
1576+
ArrayValue: &btpb.ArrayValue{
1577+
Values: []*btpb.Value{},
1578+
},
15771579
},
15781580
},
15791581
},
@@ -1607,7 +1609,7 @@ func TestAnySQLTypeToPbVal(t *testing.T) {
16071609
return
16081610
}
16091611
if tt.wantErr {
1610-
if err != nil && err.Error() != tt.wantErrMsg {
1612+
if err != nil && !strings.Contains(err.Error(), tt.wantErrMsg) {
16111613
t.Errorf("error got: %v, want: %v", err, tt.wantErrMsg)
16121614
}
16131615
return
@@ -1662,14 +1664,14 @@ func TestPreparedStatementBind(t *testing.T) {
16621664
paramTypes: map[string]SQLType{"param1": StringSQLType{}, "param2": StringSQLType{}},
16631665
values: map[string]any{"param1": "value1"},
16641666
wantErr: true,
1665-
wantErrMsg: "bigtable: parameter param2 not bound in prepared statement",
1667+
wantErrMsg: "bigtable: parameter \"param2\" not bound in call to Bind",
16661668
},
16671669
{
16681670
testName: "not bound error - all missing",
16691671
paramTypes: map[string]SQLType{"param1": StringSQLType{}, "param2": StringSQLType{}},
16701672
values: nil,
16711673
wantErr: true,
1672-
wantErrMsg: "not bound in prepared statement",
1674+
wantErrMsg: "not bound in call to Bind",
16731675
},
16741676
}
16751677
for _, tt := range tests {
@@ -2010,7 +2012,7 @@ func TestExecuteQuery(t *testing.T) {
20102012

20112013
// Execute query
20122014
err = bs.Execute(ctx, func(rr ResultRow) bool {
2013-
vals := rr.values
2015+
vals := rr.pbValues
20142016

20152017
// Compare row values
20162018
if gotRowCount > len(tc.wantResultRowValues) ||

0 commit comments

Comments
 (0)