mirror of
https://github.com/dutchcoders/transfer.sh.git
synced 2024-11-27 14:40:18 +01:00
1734 lines
49 KiB
Go
1734 lines
49 KiB
Go
|
/*
|
||
|
|
Copyright 2017 Google LLC
|
||
|
|
|
||
|
|
Licensed under the Apache License, Version 2.0 (the "License");
|
||
|
|
you may not use this file except in compliance with the License.
|
||
|
|
You may obtain a copy of the License at
|
||
|
|
|
||
|
|
http://www.apache.org/licenses/LICENSE-2.0
|
||
|
|
|
||
|
|
Unless required by applicable law or agreed to in writing, software
|
||
|
|
distributed under the License is distributed on an "AS IS" BASIS,
|
||
|
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||
|
|
See the License for the specific language governing permissions and
|
||
|
|
limitations under the License.
|
||
|
|
*/
|
||
|
|
|
||
|
|
package spanner
|
||
|
|
|
||
|
|
import (
|
||
|
|
"errors"
|
||
|
|
"fmt"
|
||
|
|
"io"
|
||
|
|
"sync/atomic"
|
||
|
|
"testing"
|
||
|
|
"time"
|
||
|
|
|
||
|
|
"golang.org/x/net/context"
|
||
|
|
"google.golang.org/grpc/status"
|
||
|
|
|
||
|
|
"github.com/golang/protobuf/proto"
|
||
|
|
proto3 "github.com/golang/protobuf/ptypes/struct"
|
||
|
|
|
||
|
|
"cloud.google.com/go/spanner/internal/testutil"
|
||
|
|
"google.golang.org/api/iterator"
|
||
|
|
sppb "google.golang.org/genproto/googleapis/spanner/v1"
|
||
|
|
"google.golang.org/grpc"
|
||
|
|
"google.golang.org/grpc/codes"
|
||
|
|
)
|
||
|
|
|
||
|
|
var (
|
||
|
|
// Mocked transaction timestamp.
|
||
|
|
trxTs = time.Unix(1, 2)
|
||
|
|
// Metadata for mocked KV table, its rows are returned by SingleUse transactions.
|
||
|
|
kvMeta = func() *sppb.ResultSetMetadata {
|
||
|
|
meta := testutil.KvMeta
|
||
|
|
meta.Transaction = &sppb.Transaction{
|
||
|
|
ReadTimestamp: timestampProto(trxTs),
|
||
|
|
}
|
||
|
|
return &meta
|
||
|
|
}()
|
||
|
|
// Metadata for mocked ListKV table, which uses List for its key and value.
|
||
|
|
// Its rows are returned by snapshot readonly transactions, as indicated in the transaction metadata.
|
||
|
|
kvListMeta = &sppb.ResultSetMetadata{
|
||
|
|
RowType: &sppb.StructType{
|
||
|
|
Fields: []*sppb.StructType_Field{
|
||
|
|
{
|
||
|
|
Name: "Key",
|
||
|
|
Type: &sppb.Type{
|
||
|
|
Code: sppb.TypeCode_ARRAY,
|
||
|
|
ArrayElementType: &sppb.Type{
|
||
|
|
Code: sppb.TypeCode_STRING,
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
{
|
||
|
|
Name: "Value",
|
||
|
|
Type: &sppb.Type{
|
||
|
|
Code: sppb.TypeCode_ARRAY,
|
||
|
|
ArrayElementType: &sppb.Type{
|
||
|
|
Code: sppb.TypeCode_STRING,
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
Transaction: &sppb.Transaction{
|
||
|
|
Id: transactionID{5, 6, 7, 8, 9},
|
||
|
|
ReadTimestamp: timestampProto(trxTs),
|
||
|
|
},
|
||
|
|
}
|
||
|
|
// Metadata for mocked schema of a query result set, which has two struct
|
||
|
|
// columns named "Col1" and "Col2", the struct's schema is like the
|
||
|
|
// following:
|
||
|
|
//
|
||
|
|
// STRUCT {
|
||
|
|
// INT
|
||
|
|
// LIST<STRING>
|
||
|
|
// }
|
||
|
|
//
|
||
|
|
// Its rows are returned in readwrite transaction, as indicated in the transaction metadata.
|
||
|
|
kvObjectMeta = &sppb.ResultSetMetadata{
|
||
|
|
RowType: &sppb.StructType{
|
||
|
|
Fields: []*sppb.StructType_Field{
|
||
|
|
{
|
||
|
|
Name: "Col1",
|
||
|
|
Type: &sppb.Type{
|
||
|
|
Code: sppb.TypeCode_STRUCT,
|
||
|
|
StructType: &sppb.StructType{
|
||
|
|
Fields: []*sppb.StructType_Field{
|
||
|
|
{
|
||
|
|
Name: "foo-f1",
|
||
|
|
Type: &sppb.Type{
|
||
|
|
Code: sppb.TypeCode_INT64,
|
||
|
|
},
|
||
|
|
},
|
||
|
|
{
|
||
|
|
Name: "foo-f2",
|
||
|
|
Type: &sppb.Type{
|
||
|
|
Code: sppb.TypeCode_ARRAY,
|
||
|
|
ArrayElementType: &sppb.Type{
|
||
|
|
Code: sppb.TypeCode_STRING,
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
{
|
||
|
|
Name: "Col2",
|
||
|
|
Type: &sppb.Type{
|
||
|
|
Code: sppb.TypeCode_STRUCT,
|
||
|
|
StructType: &sppb.StructType{
|
||
|
|
Fields: []*sppb.StructType_Field{
|
||
|
|
{
|
||
|
|
Name: "bar-f1",
|
||
|
|
Type: &sppb.Type{
|
||
|
|
Code: sppb.TypeCode_INT64,
|
||
|
|
},
|
||
|
|
},
|
||
|
|
{
|
||
|
|
Name: "bar-f2",
|
||
|
|
Type: &sppb.Type{
|
||
|
|
Code: sppb.TypeCode_ARRAY,
|
||
|
|
ArrayElementType: &sppb.Type{
|
||
|
|
Code: sppb.TypeCode_STRING,
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
Transaction: &sppb.Transaction{
|
||
|
|
Id: transactionID{1, 2, 3, 4, 5},
|
||
|
|
},
|
||
|
|
}
|
||
|
|
)
|
||
|
|
|
||
|
|
// String implements fmt.stringer.
|
||
|
|
func (r *Row) String() string {
|
||
|
|
return fmt.Sprintf("{fields: %s, val: %s}", r.fields, r.vals)
|
||
|
|
}
|
||
|
|
|
||
|
|
func describeRows(l []*Row) string {
|
||
|
|
// generate a nice test failure description
|
||
|
|
var s = "["
|
||
|
|
for i, r := range l {
|
||
|
|
if i != 0 {
|
||
|
|
s += ",\n "
|
||
|
|
}
|
||
|
|
s += fmt.Sprint(r)
|
||
|
|
}
|
||
|
|
s += "]"
|
||
|
|
return s
|
||
|
|
}
|
||
|
|
|
||
|
|
// Helper for generating proto3 Value_ListValue instances, making
|
||
|
|
// test code shorter and readable.
|
||
|
|
func genProtoListValue(v ...string) *proto3.Value_ListValue {
|
||
|
|
r := &proto3.Value_ListValue{
|
||
|
|
ListValue: &proto3.ListValue{
|
||
|
|
Values: []*proto3.Value{},
|
||
|
|
},
|
||
|
|
}
|
||
|
|
for _, e := range v {
|
||
|
|
r.ListValue.Values = append(
|
||
|
|
r.ListValue.Values,
|
||
|
|
&proto3.Value{
|
||
|
|
Kind: &proto3.Value_StringValue{StringValue: e},
|
||
|
|
},
|
||
|
|
)
|
||
|
|
}
|
||
|
|
return r
|
||
|
|
}
|
||
|
|
|
||
|
|
// Test Row generation logics of partialResultSetDecoder.
|
||
|
|
func TestPartialResultSetDecoder(t *testing.T) {
|
||
|
|
restore := setMaxBytesBetweenResumeTokens()
|
||
|
|
defer restore()
|
||
|
|
var tests = []struct {
|
||
|
|
input []*sppb.PartialResultSet
|
||
|
|
wantF []*Row
|
||
|
|
wantTxID transactionID
|
||
|
|
wantTs time.Time
|
||
|
|
wantD bool
|
||
|
|
}{
|
||
|
|
{
|
||
|
|
// Empty input.
|
||
|
|
wantD: true,
|
||
|
|
},
|
||
|
|
// String merging examples.
|
||
|
|
{
|
||
|
|
// Single KV result.
|
||
|
|
input: []*sppb.PartialResultSet{
|
||
|
|
{
|
||
|
|
Metadata: kvMeta,
|
||
|
|
Values: []*proto3.Value{
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: "foo"}},
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: "bar"}},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
wantF: []*Row{
|
||
|
|
{
|
||
|
|
fields: kvMeta.RowType.Fields,
|
||
|
|
vals: []*proto3.Value{
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: "foo"}},
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: "bar"}},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
wantTs: trxTs,
|
||
|
|
wantD: true,
|
||
|
|
},
|
||
|
|
{
|
||
|
|
// Incomplete partial result.
|
||
|
|
input: []*sppb.PartialResultSet{
|
||
|
|
{
|
||
|
|
Metadata: kvMeta,
|
||
|
|
Values: []*proto3.Value{
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: "foo"}},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
wantTs: trxTs,
|
||
|
|
wantD: false,
|
||
|
|
},
|
||
|
|
{
|
||
|
|
// Complete splitted result.
|
||
|
|
input: []*sppb.PartialResultSet{
|
||
|
|
{
|
||
|
|
Metadata: kvMeta,
|
||
|
|
Values: []*proto3.Value{
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: "foo"}},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
{
|
||
|
|
Values: []*proto3.Value{
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: "bar"}},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
wantF: []*Row{
|
||
|
|
{
|
||
|
|
fields: kvMeta.RowType.Fields,
|
||
|
|
vals: []*proto3.Value{
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: "foo"}},
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: "bar"}},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
wantTs: trxTs,
|
||
|
|
wantD: true,
|
||
|
|
},
|
||
|
|
{
|
||
|
|
// Multi-row example with splitted row in the middle.
|
||
|
|
input: []*sppb.PartialResultSet{
|
||
|
|
{
|
||
|
|
Metadata: kvMeta,
|
||
|
|
Values: []*proto3.Value{
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: "foo"}},
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: "bar"}},
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: "A"}},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
{
|
||
|
|
Values: []*proto3.Value{
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: "1"}},
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: "B"}},
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: "2"}},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
wantF: []*Row{
|
||
|
|
{
|
||
|
|
fields: kvMeta.RowType.Fields,
|
||
|
|
vals: []*proto3.Value{
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: "foo"}},
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: "bar"}},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
{
|
||
|
|
fields: kvMeta.RowType.Fields,
|
||
|
|
vals: []*proto3.Value{
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: "A"}},
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: "1"}},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
{
|
||
|
|
fields: kvMeta.RowType.Fields,
|
||
|
|
vals: []*proto3.Value{
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: "B"}},
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: "2"}},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
wantTs: trxTs,
|
||
|
|
wantD: true,
|
||
|
|
},
|
||
|
|
{
|
||
|
|
// Merging example in result_set.proto.
|
||
|
|
input: []*sppb.PartialResultSet{
|
||
|
|
{
|
||
|
|
Metadata: kvMeta,
|
||
|
|
Values: []*proto3.Value{
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: "Hello"}},
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: "W"}},
|
||
|
|
},
|
||
|
|
ChunkedValue: true,
|
||
|
|
},
|
||
|
|
{
|
||
|
|
Values: []*proto3.Value{
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: "orl"}},
|
||
|
|
},
|
||
|
|
ChunkedValue: true,
|
||
|
|
},
|
||
|
|
{
|
||
|
|
Values: []*proto3.Value{
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: "d"}},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
wantF: []*Row{
|
||
|
|
{
|
||
|
|
fields: kvMeta.RowType.Fields,
|
||
|
|
vals: []*proto3.Value{
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: "Hello"}},
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: "World"}},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
wantTs: trxTs,
|
||
|
|
wantD: true,
|
||
|
|
},
|
||
|
|
{
|
||
|
|
// More complex example showing completing a merge and
|
||
|
|
// starting a new merge in the same partialResultSet.
|
||
|
|
input: []*sppb.PartialResultSet{
|
||
|
|
{
|
||
|
|
Metadata: kvMeta,
|
||
|
|
Values: []*proto3.Value{
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: "Hello"}},
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: "W"}}, // start split in value
|
||
|
|
},
|
||
|
|
ChunkedValue: true,
|
||
|
|
},
|
||
|
|
{
|
||
|
|
Values: []*proto3.Value{
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: "orld"}}, // complete value
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: "i"}}, // start split in key
|
||
|
|
},
|
||
|
|
ChunkedValue: true,
|
||
|
|
},
|
||
|
|
{
|
||
|
|
Values: []*proto3.Value{
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: "s"}}, // complete key
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: "not"}},
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: "a"}},
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: "qu"}}, // split in value
|
||
|
|
},
|
||
|
|
ChunkedValue: true,
|
||
|
|
},
|
||
|
|
{
|
||
|
|
Values: []*proto3.Value{
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: "estion"}}, // complete value
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
wantF: []*Row{
|
||
|
|
{
|
||
|
|
fields: kvMeta.RowType.Fields,
|
||
|
|
vals: []*proto3.Value{
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: "Hello"}},
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: "World"}},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
{
|
||
|
|
fields: kvMeta.RowType.Fields,
|
||
|
|
vals: []*proto3.Value{
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: "is"}},
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: "not"}},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
{
|
||
|
|
fields: kvMeta.RowType.Fields,
|
||
|
|
vals: []*proto3.Value{
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: "a"}},
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: "question"}},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
wantTs: trxTs,
|
||
|
|
wantD: true,
|
||
|
|
},
|
||
|
|
// List merging examples.
|
||
|
|
{
|
||
|
|
// Non-splitting Lists.
|
||
|
|
input: []*sppb.PartialResultSet{
|
||
|
|
{
|
||
|
|
Metadata: kvListMeta,
|
||
|
|
Values: []*proto3.Value{
|
||
|
|
{
|
||
|
|
Kind: genProtoListValue("foo-1", "foo-2"),
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
{
|
||
|
|
Values: []*proto3.Value{
|
||
|
|
{
|
||
|
|
Kind: genProtoListValue("bar-1", "bar-2"),
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
wantF: []*Row{
|
||
|
|
{
|
||
|
|
fields: kvListMeta.RowType.Fields,
|
||
|
|
vals: []*proto3.Value{
|
||
|
|
{
|
||
|
|
Kind: genProtoListValue("foo-1", "foo-2"),
|
||
|
|
},
|
||
|
|
{
|
||
|
|
Kind: genProtoListValue("bar-1", "bar-2"),
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
wantTxID: transactionID{5, 6, 7, 8, 9},
|
||
|
|
wantTs: trxTs,
|
||
|
|
wantD: true,
|
||
|
|
},
|
||
|
|
{
|
||
|
|
// Simple List merge case: splitted string element.
|
||
|
|
input: []*sppb.PartialResultSet{
|
||
|
|
{
|
||
|
|
Metadata: kvListMeta,
|
||
|
|
Values: []*proto3.Value{
|
||
|
|
{
|
||
|
|
Kind: genProtoListValue("foo-1", "foo-"),
|
||
|
|
},
|
||
|
|
},
|
||
|
|
ChunkedValue: true,
|
||
|
|
},
|
||
|
|
{
|
||
|
|
Values: []*proto3.Value{
|
||
|
|
{
|
||
|
|
Kind: genProtoListValue("2"),
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
{
|
||
|
|
Values: []*proto3.Value{
|
||
|
|
{
|
||
|
|
Kind: genProtoListValue("bar-1", "bar-2"),
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
wantF: []*Row{
|
||
|
|
{
|
||
|
|
fields: kvListMeta.RowType.Fields,
|
||
|
|
vals: []*proto3.Value{
|
||
|
|
{
|
||
|
|
Kind: genProtoListValue("foo-1", "foo-2"),
|
||
|
|
},
|
||
|
|
{
|
||
|
|
Kind: genProtoListValue("bar-1", "bar-2"),
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
wantTxID: transactionID{5, 6, 7, 8, 9},
|
||
|
|
wantTs: trxTs,
|
||
|
|
wantD: true,
|
||
|
|
},
|
||
|
|
{
|
||
|
|
// Struct merging is also implemented by List merging. Note that
|
||
|
|
// Cloud Spanner uses proto.ListValue to encode Structs as well.
|
||
|
|
input: []*sppb.PartialResultSet{
|
||
|
|
{
|
||
|
|
Metadata: kvObjectMeta,
|
||
|
|
Values: []*proto3.Value{
|
||
|
|
{
|
||
|
|
Kind: &proto3.Value_ListValue{
|
||
|
|
ListValue: &proto3.ListValue{
|
||
|
|
Values: []*proto3.Value{
|
||
|
|
{Kind: &proto3.Value_NumberValue{NumberValue: 23}},
|
||
|
|
{Kind: genProtoListValue("foo-1", "fo")},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
ChunkedValue: true,
|
||
|
|
},
|
||
|
|
{
|
||
|
|
Values: []*proto3.Value{
|
||
|
|
{
|
||
|
|
Kind: &proto3.Value_ListValue{
|
||
|
|
ListValue: &proto3.ListValue{
|
||
|
|
Values: []*proto3.Value{
|
||
|
|
{Kind: genProtoListValue("o-2", "f")},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
ChunkedValue: true,
|
||
|
|
},
|
||
|
|
{
|
||
|
|
Values: []*proto3.Value{
|
||
|
|
{
|
||
|
|
Kind: &proto3.Value_ListValue{
|
||
|
|
ListValue: &proto3.ListValue{
|
||
|
|
Values: []*proto3.Value{
|
||
|
|
{Kind: genProtoListValue("oo-3")},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
{
|
||
|
|
Kind: &proto3.Value_ListValue{
|
||
|
|
ListValue: &proto3.ListValue{
|
||
|
|
Values: []*proto3.Value{
|
||
|
|
{Kind: &proto3.Value_NumberValue{NumberValue: 45}},
|
||
|
|
{Kind: genProtoListValue("bar-1")},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
wantF: []*Row{
|
||
|
|
{
|
||
|
|
fields: kvObjectMeta.RowType.Fields,
|
||
|
|
vals: []*proto3.Value{
|
||
|
|
{
|
||
|
|
Kind: &proto3.Value_ListValue{
|
||
|
|
ListValue: &proto3.ListValue{
|
||
|
|
Values: []*proto3.Value{
|
||
|
|
{Kind: &proto3.Value_NumberValue{NumberValue: 23}},
|
||
|
|
{Kind: genProtoListValue("foo-1", "foo-2", "foo-3")},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
{
|
||
|
|
Kind: &proto3.Value_ListValue{
|
||
|
|
ListValue: &proto3.ListValue{
|
||
|
|
Values: []*proto3.Value{
|
||
|
|
{Kind: &proto3.Value_NumberValue{NumberValue: 45}},
|
||
|
|
{Kind: genProtoListValue("bar-1")},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
wantTxID: transactionID{1, 2, 3, 4, 5},
|
||
|
|
wantD: true,
|
||
|
|
},
|
||
|
|
}
|
||
|
|
|
||
|
|
nextTest:
|
||
|
|
for i, test := range tests {
|
||
|
|
var rows []*Row
|
||
|
|
p := &partialResultSetDecoder{}
|
||
|
|
for j, v := range test.input {
|
||
|
|
rs, err := p.add(v)
|
||
|
|
if err != nil {
|
||
|
|
t.Errorf("test %d.%d: partialResultSetDecoder.add(%v) = %v; want nil", i, j, v, err)
|
||
|
|
continue nextTest
|
||
|
|
}
|
||
|
|
rows = append(rows, rs...)
|
||
|
|
}
|
||
|
|
if !testEqual(p.ts, test.wantTs) {
|
||
|
|
t.Errorf("got transaction(%v), want %v", p.ts, test.wantTs)
|
||
|
|
}
|
||
|
|
if !testEqual(rows, test.wantF) {
|
||
|
|
t.Errorf("test %d: rows=\n%v\n; want\n%v\n; p.row:\n%v\n", i, describeRows(rows), describeRows(test.wantF), p.row)
|
||
|
|
}
|
||
|
|
if got := p.done(); got != test.wantD {
|
||
|
|
t.Errorf("test %d: partialResultSetDecoder.done() = %v", i, got)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
const (
|
||
|
|
maxBuffers = 16 // max number of PartialResultSets that will be buffered in tests.
|
||
|
|
)
|
||
|
|
|
||
|
|
// setMaxBytesBetweenResumeTokens sets the global maxBytesBetweenResumeTokens to a smaller
|
||
|
|
// value more suitable for tests. It returns a function which should be called to restore
|
||
|
|
// the maxBytesBetweenResumeTokens to its old value
|
||
|
|
func setMaxBytesBetweenResumeTokens() func() {
|
||
|
|
o := atomic.LoadInt32(&maxBytesBetweenResumeTokens)
|
||
|
|
atomic.StoreInt32(&maxBytesBetweenResumeTokens, int32(maxBuffers*proto.Size(&sppb.PartialResultSet{
|
||
|
|
Metadata: kvMeta,
|
||
|
|
Values: []*proto3.Value{
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: keyStr(0)}},
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: valStr(0)}},
|
||
|
|
},
|
||
|
|
})))
|
||
|
|
return func() {
|
||
|
|
atomic.StoreInt32(&maxBytesBetweenResumeTokens, o)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// keyStr generates key string for kvMeta schema.
|
||
|
|
func keyStr(i int) string {
|
||
|
|
return fmt.Sprintf("foo-%02d", i)
|
||
|
|
}
|
||
|
|
|
||
|
|
// valStr generates value string for kvMeta schema.
|
||
|
|
func valStr(i int) string {
|
||
|
|
return fmt.Sprintf("bar-%02d", i)
|
||
|
|
}
|
||
|
|
|
||
|
|
// Test state transitions of resumableStreamDecoder where state machine
|
||
|
|
// ends up to a non-blocking state(resumableStreamDecoder.Next returns
|
||
|
|
// on non-blocking state).
|
||
|
|
func TestRsdNonblockingStates(t *testing.T) {
|
||
|
|
restore := setMaxBytesBetweenResumeTokens()
|
||
|
|
defer restore()
|
||
|
|
tests := []struct {
|
||
|
|
name string
|
||
|
|
msgs []testutil.MockCtlMsg
|
||
|
|
rpc func(ct context.Context, resumeToken []byte) (streamingReceiver, error)
|
||
|
|
sql string
|
||
|
|
// Expected values
|
||
|
|
want []*sppb.PartialResultSet // PartialResultSets that should be returned to caller
|
||
|
|
queue []*sppb.PartialResultSet // PartialResultSets that should be buffered
|
||
|
|
resumeToken []byte // Resume token that is maintained by resumableStreamDecoder
|
||
|
|
stateHistory []resumableStreamDecoderState // State transition history of resumableStreamDecoder
|
||
|
|
wantErr error
|
||
|
|
}{
|
||
|
|
{
|
||
|
|
// unConnected->queueingRetryable->finished
|
||
|
|
name: "unConnected->queueingRetryable->finished",
|
||
|
|
msgs: []testutil.MockCtlMsg{
|
||
|
|
{},
|
||
|
|
{},
|
||
|
|
{Err: io.EOF, ResumeToken: false},
|
||
|
|
},
|
||
|
|
sql: "SELECT t.key key, t.value value FROM t_mock t",
|
||
|
|
want: []*sppb.PartialResultSet{
|
||
|
|
{
|
||
|
|
Metadata: kvMeta,
|
||
|
|
Values: []*proto3.Value{
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: keyStr(0)}},
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: valStr(0)}},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
queue: []*sppb.PartialResultSet{
|
||
|
|
{
|
||
|
|
Metadata: kvMeta,
|
||
|
|
Values: []*proto3.Value{
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: keyStr(1)}},
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: valStr(1)}},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
stateHistory: []resumableStreamDecoderState{
|
||
|
|
queueingRetryable, // do RPC
|
||
|
|
queueingRetryable, // got foo-00
|
||
|
|
queueingRetryable, // got foo-01
|
||
|
|
finished, // got EOF
|
||
|
|
},
|
||
|
|
},
|
||
|
|
{
|
||
|
|
// unConnected->queueingRetryable->aborted
|
||
|
|
name: "unConnected->queueingRetryable->aborted",
|
||
|
|
msgs: []testutil.MockCtlMsg{
|
||
|
|
{},
|
||
|
|
{Err: nil, ResumeToken: true},
|
||
|
|
{},
|
||
|
|
{Err: errors.New("I quit"), ResumeToken: false},
|
||
|
|
},
|
||
|
|
sql: "SELECT t.key key, t.value value FROM t_mock t",
|
||
|
|
want: []*sppb.PartialResultSet{
|
||
|
|
{
|
||
|
|
Metadata: kvMeta,
|
||
|
|
Values: []*proto3.Value{
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: keyStr(0)}},
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: valStr(0)}},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
{
|
||
|
|
Metadata: kvMeta,
|
||
|
|
Values: []*proto3.Value{
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: keyStr(1)}},
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: valStr(1)}},
|
||
|
|
},
|
||
|
|
ResumeToken: testutil.EncodeResumeToken(1),
|
||
|
|
},
|
||
|
|
},
|
||
|
|
stateHistory: []resumableStreamDecoderState{
|
||
|
|
queueingRetryable, // do RPC
|
||
|
|
queueingRetryable, // got foo-00
|
||
|
|
queueingRetryable, // got foo-01
|
||
|
|
queueingRetryable, // foo-01, resume token
|
||
|
|
queueingRetryable, // got foo-02
|
||
|
|
aborted, // got error
|
||
|
|
},
|
||
|
|
wantErr: status.Errorf(codes.Unknown, "I quit"),
|
||
|
|
},
|
||
|
|
{
|
||
|
|
// unConnected->queueingRetryable->queueingUnretryable->queueingUnretryable
|
||
|
|
name: "unConnected->queueingRetryable->queueingUnretryable->queueingUnretryable",
|
||
|
|
msgs: func() (m []testutil.MockCtlMsg) {
|
||
|
|
for i := 0; i < maxBuffers+1; i++ {
|
||
|
|
m = append(m, testutil.MockCtlMsg{})
|
||
|
|
}
|
||
|
|
return m
|
||
|
|
}(),
|
||
|
|
sql: "SELECT t.key key, t.value value FROM t_mock t",
|
||
|
|
want: func() (s []*sppb.PartialResultSet) {
|
||
|
|
for i := 0; i < maxBuffers+1; i++ {
|
||
|
|
s = append(s, &sppb.PartialResultSet{
|
||
|
|
Metadata: kvMeta,
|
||
|
|
Values: []*proto3.Value{
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: keyStr(i)}},
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: valStr(i)}},
|
||
|
|
},
|
||
|
|
})
|
||
|
|
}
|
||
|
|
return s
|
||
|
|
}(),
|
||
|
|
stateHistory: func() (s []resumableStreamDecoderState) {
|
||
|
|
s = append(s, queueingRetryable) // RPC
|
||
|
|
for i := 0; i < maxBuffers; i++ {
|
||
|
|
s = append(s, queueingRetryable) // the internal queue of resumableStreamDecoder fills up
|
||
|
|
}
|
||
|
|
// the first item fills up the queue and triggers state transition;
|
||
|
|
// the second item is received under queueingUnretryable state.
|
||
|
|
s = append(s, queueingUnretryable)
|
||
|
|
s = append(s, queueingUnretryable)
|
||
|
|
return s
|
||
|
|
}(),
|
||
|
|
},
|
||
|
|
{
|
||
|
|
// unConnected->queueingRetryable->queueingUnretryable->aborted
|
||
|
|
name: "unConnected->queueingRetryable->queueingUnretryable->aborted",
|
||
|
|
msgs: func() (m []testutil.MockCtlMsg) {
|
||
|
|
for i := 0; i < maxBuffers; i++ {
|
||
|
|
m = append(m, testutil.MockCtlMsg{})
|
||
|
|
}
|
||
|
|
m = append(m, testutil.MockCtlMsg{Err: errors.New("Just Abort It"), ResumeToken: false})
|
||
|
|
return m
|
||
|
|
}(),
|
||
|
|
sql: "SELECT t.key key, t.value value FROM t_mock t",
|
||
|
|
want: func() (s []*sppb.PartialResultSet) {
|
||
|
|
for i := 0; i < maxBuffers; i++ {
|
||
|
|
s = append(s, &sppb.PartialResultSet{
|
||
|
|
Metadata: kvMeta,
|
||
|
|
Values: []*proto3.Value{
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: keyStr(i)}},
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: valStr(i)}},
|
||
|
|
},
|
||
|
|
})
|
||
|
|
}
|
||
|
|
return s
|
||
|
|
}(),
|
||
|
|
stateHistory: func() (s []resumableStreamDecoderState) {
|
||
|
|
s = append(s, queueingRetryable) // RPC
|
||
|
|
for i := 0; i < maxBuffers; i++ {
|
||
|
|
s = append(s, queueingRetryable) // internal queue of resumableStreamDecoder fills up
|
||
|
|
}
|
||
|
|
s = append(s, queueingUnretryable) // the last row triggers state change
|
||
|
|
s = append(s, aborted) // Error happens
|
||
|
|
return s
|
||
|
|
}(),
|
||
|
|
wantErr: status.Errorf(codes.Unknown, "Just Abort It"),
|
||
|
|
},
|
||
|
|
}
|
||
|
|
nextTest:
|
||
|
|
for _, test := range tests {
|
||
|
|
ms := testutil.NewMockCloudSpanner(t, trxTs)
|
||
|
|
ms.Serve()
|
||
|
|
mc := sppb.NewSpannerClient(dialMock(t, ms))
|
||
|
|
if test.rpc == nil {
|
||
|
|
test.rpc = func(ct context.Context, resumeToken []byte) (streamingReceiver, error) {
|
||
|
|
return mc.ExecuteStreamingSql(ct, &sppb.ExecuteSqlRequest{
|
||
|
|
Sql: test.sql,
|
||
|
|
ResumeToken: resumeToken,
|
||
|
|
})
|
||
|
|
}
|
||
|
|
}
|
||
|
|
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||
|
|
defer cancel()
|
||
|
|
r := newResumableStreamDecoder(
|
||
|
|
ctx,
|
||
|
|
test.rpc,
|
||
|
|
)
|
||
|
|
st := []resumableStreamDecoderState{}
|
||
|
|
var lastErr error
|
||
|
|
// Once the expected number of state transitions are observed,
|
||
|
|
// send a signal by setting stateDone = true.
|
||
|
|
stateDone := false
|
||
|
|
// Set stateWitness to listen to state changes.
|
||
|
|
hl := len(test.stateHistory) // To avoid data race on test.
|
||
|
|
r.stateWitness = func(rs resumableStreamDecoderState) {
|
||
|
|
if !stateDone {
|
||
|
|
// Record state transitions.
|
||
|
|
st = append(st, rs)
|
||
|
|
if len(st) == hl {
|
||
|
|
lastErr = r.lastErr()
|
||
|
|
stateDone = true
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
// Let mock server stream given messages to resumableStreamDecoder.
|
||
|
|
for _, m := range test.msgs {
|
||
|
|
ms.AddMsg(m.Err, m.ResumeToken)
|
||
|
|
}
|
||
|
|
var rs []*sppb.PartialResultSet
|
||
|
|
for {
|
||
|
|
select {
|
||
|
|
case <-ctx.Done():
|
||
|
|
t.Errorf("context cancelled or timeout during test")
|
||
|
|
continue nextTest
|
||
|
|
default:
|
||
|
|
}
|
||
|
|
if stateDone {
|
||
|
|
// Check if resumableStreamDecoder carried out expected
|
||
|
|
// state transitions.
|
||
|
|
if !testEqual(st, test.stateHistory) {
|
||
|
|
t.Errorf("%v: observed state transitions: \n%v\n, want \n%v\n",
|
||
|
|
test.name, st, test.stateHistory)
|
||
|
|
}
|
||
|
|
// Check if resumableStreamDecoder returns expected array of
|
||
|
|
// PartialResultSets.
|
||
|
|
if !testEqual(rs, test.want) {
|
||
|
|
t.Errorf("%v: received PartialResultSets: \n%v\n, want \n%v\n", test.name, rs, test.want)
|
||
|
|
}
|
||
|
|
// Verify that resumableStreamDecoder's internal buffering is also correct.
|
||
|
|
var q []*sppb.PartialResultSet
|
||
|
|
for {
|
||
|
|
item := r.q.pop()
|
||
|
|
if item == nil {
|
||
|
|
break
|
||
|
|
}
|
||
|
|
q = append(q, item)
|
||
|
|
}
|
||
|
|
if !testEqual(q, test.queue) {
|
||
|
|
t.Errorf("%v: PartialResultSets still queued: \n%v\n, want \n%v\n", test.name, q, test.queue)
|
||
|
|
}
|
||
|
|
// Verify resume token.
|
||
|
|
if test.resumeToken != nil && !testEqual(r.resumeToken, test.resumeToken) {
|
||
|
|
t.Errorf("%v: Resume token is %v, want %v\n", test.name, r.resumeToken, test.resumeToken)
|
||
|
|
}
|
||
|
|
// Verify error message.
|
||
|
|
if !testEqual(lastErr, test.wantErr) {
|
||
|
|
t.Errorf("%v: got error %v, want %v", test.name, lastErr, test.wantErr)
|
||
|
|
}
|
||
|
|
// Proceed to next test
|
||
|
|
continue nextTest
|
||
|
|
}
|
||
|
|
// Receive next decoded item.
|
||
|
|
if r.next() {
|
||
|
|
rs = append(rs, r.get())
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// Test state transitions of resumableStreamDecoder where state machine
|
||
|
|
// ends up to a blocking state(resumableStreamDecoder.Next blocks
|
||
|
|
// on blocking state).
|
||
|
|
func TestRsdBlockingStates(t *testing.T) {
|
||
|
|
restore := setMaxBytesBetweenResumeTokens()
|
||
|
|
defer restore()
|
||
|
|
tests := []struct {
|
||
|
|
name string
|
||
|
|
msgs []testutil.MockCtlMsg
|
||
|
|
rpc func(ct context.Context, resumeToken []byte) (streamingReceiver, error)
|
||
|
|
sql string
|
||
|
|
// Expected values
|
||
|
|
want []*sppb.PartialResultSet // PartialResultSets that should be returned to caller
|
||
|
|
queue []*sppb.PartialResultSet // PartialResultSets that should be buffered
|
||
|
|
resumeToken []byte // Resume token that is maintained by resumableStreamDecoder
|
||
|
|
stateHistory []resumableStreamDecoderState // State transition history of resumableStreamDecoder
|
||
|
|
wantErr error
|
||
|
|
}{
|
||
|
|
{
|
||
|
|
// unConnected -> unConnected
|
||
|
|
name: "unConnected -> unConnected",
|
||
|
|
rpc: func(ct context.Context, resumeToken []byte) (streamingReceiver, error) {
|
||
|
|
return nil, status.Errorf(codes.Unavailable, "trust me: server is unavailable")
|
||
|
|
},
|
||
|
|
sql: "SELECT * from t_whatever",
|
||
|
|
stateHistory: []resumableStreamDecoderState{unConnected, unConnected, unConnected},
|
||
|
|
wantErr: status.Errorf(codes.Unavailable, "trust me: server is unavailable"),
|
||
|
|
},
|
||
|
|
{
|
||
|
|
// unConnected -> queueingRetryable
|
||
|
|
name: "unConnected -> queueingRetryable",
|
||
|
|
sql: "SELECT t.key key, t.value value FROM t_mock t",
|
||
|
|
stateHistory: []resumableStreamDecoderState{queueingRetryable},
|
||
|
|
},
|
||
|
|
{
|
||
|
|
// unConnected->queueingRetryable->queueingRetryable
|
||
|
|
name: "unConnected->queueingRetryable->queueingRetryable",
|
||
|
|
msgs: []testutil.MockCtlMsg{
|
||
|
|
{},
|
||
|
|
{Err: nil, ResumeToken: true},
|
||
|
|
{Err: nil, ResumeToken: true},
|
||
|
|
{},
|
||
|
|
},
|
||
|
|
sql: "SELECT t.key key, t.value value FROM t_mock t",
|
||
|
|
want: []*sppb.PartialResultSet{
|
||
|
|
{
|
||
|
|
Metadata: kvMeta,
|
||
|
|
Values: []*proto3.Value{
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: keyStr(0)}},
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: valStr(0)}},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
{
|
||
|
|
Metadata: kvMeta,
|
||
|
|
Values: []*proto3.Value{
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: keyStr(1)}},
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: valStr(1)}},
|
||
|
|
},
|
||
|
|
ResumeToken: testutil.EncodeResumeToken(1),
|
||
|
|
},
|
||
|
|
{
|
||
|
|
Metadata: kvMeta,
|
||
|
|
Values: []*proto3.Value{
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: keyStr(2)}},
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: valStr(2)}},
|
||
|
|
},
|
||
|
|
ResumeToken: testutil.EncodeResumeToken(2),
|
||
|
|
},
|
||
|
|
},
|
||
|
|
queue: []*sppb.PartialResultSet{
|
||
|
|
{
|
||
|
|
Metadata: kvMeta,
|
||
|
|
Values: []*proto3.Value{
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: keyStr(3)}},
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: valStr(3)}},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
resumeToken: testutil.EncodeResumeToken(2),
|
||
|
|
stateHistory: []resumableStreamDecoderState{
|
||
|
|
queueingRetryable, // do RPC
|
||
|
|
queueingRetryable, // got foo-00
|
||
|
|
queueingRetryable, // got foo-01
|
||
|
|
queueingRetryable, // foo-01, resume token
|
||
|
|
queueingRetryable, // got foo-02
|
||
|
|
queueingRetryable, // foo-02, resume token
|
||
|
|
queueingRetryable, // got foo-03
|
||
|
|
},
|
||
|
|
},
|
||
|
|
{
|
||
|
|
// unConnected->queueingRetryable->queueingUnretryable->queueingRetryable->queueingRetryable
|
||
|
|
name: "unConnected->queueingRetryable->queueingUnretryable->queueingRetryable->queueingRetryable",
|
||
|
|
msgs: func() (m []testutil.MockCtlMsg) {
|
||
|
|
for i := 0; i < maxBuffers+1; i++ {
|
||
|
|
m = append(m, testutil.MockCtlMsg{})
|
||
|
|
}
|
||
|
|
m = append(m, testutil.MockCtlMsg{Err: nil, ResumeToken: true})
|
||
|
|
m = append(m, testutil.MockCtlMsg{})
|
||
|
|
return m
|
||
|
|
}(),
|
||
|
|
sql: "SELECT t.key key, t.value value FROM t_mock t",
|
||
|
|
want: func() (s []*sppb.PartialResultSet) {
|
||
|
|
for i := 0; i < maxBuffers+2; i++ {
|
||
|
|
s = append(s, &sppb.PartialResultSet{
|
||
|
|
Metadata: kvMeta,
|
||
|
|
Values: []*proto3.Value{
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: keyStr(i)}},
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: valStr(i)}},
|
||
|
|
},
|
||
|
|
})
|
||
|
|
}
|
||
|
|
s[maxBuffers+1].ResumeToken = testutil.EncodeResumeToken(maxBuffers + 1)
|
||
|
|
return s
|
||
|
|
}(),
|
||
|
|
resumeToken: testutil.EncodeResumeToken(maxBuffers + 1),
|
||
|
|
queue: []*sppb.PartialResultSet{
|
||
|
|
{
|
||
|
|
Metadata: kvMeta,
|
||
|
|
Values: []*proto3.Value{
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: keyStr(maxBuffers + 2)}},
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: valStr(maxBuffers + 2)}},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
stateHistory: func() (s []resumableStreamDecoderState) {
|
||
|
|
s = append(s, queueingRetryable) // RPC
|
||
|
|
for i := 0; i < maxBuffers; i++ {
|
||
|
|
s = append(s, queueingRetryable) // internal queue of resumableStreamDecoder filles up
|
||
|
|
}
|
||
|
|
for i := maxBuffers - 1; i < maxBuffers+1; i++ {
|
||
|
|
// the first item fills up the queue and triggers state change;
|
||
|
|
// the second item is received under queueingUnretryable state.
|
||
|
|
s = append(s, queueingUnretryable)
|
||
|
|
}
|
||
|
|
s = append(s, queueingUnretryable) // got (maxBuffers+1)th row under Unretryable state
|
||
|
|
s = append(s, queueingRetryable) // (maxBuffers+1)th row has resume token
|
||
|
|
s = append(s, queueingRetryable) // (maxBuffers+2)th row has no resume token
|
||
|
|
return s
|
||
|
|
}(),
|
||
|
|
},
|
||
|
|
{
|
||
|
|
// unConnected->queueingRetryable->queueingUnretryable->finished
|
||
|
|
name: "unConnected->queueingRetryable->queueingUnretryable->finished",
|
||
|
|
msgs: func() (m []testutil.MockCtlMsg) {
|
||
|
|
for i := 0; i < maxBuffers; i++ {
|
||
|
|
m = append(m, testutil.MockCtlMsg{})
|
||
|
|
}
|
||
|
|
m = append(m, testutil.MockCtlMsg{Err: io.EOF, ResumeToken: false})
|
||
|
|
return m
|
||
|
|
}(),
|
||
|
|
sql: "SELECT t.key key, t.value value FROM t_mock t",
|
||
|
|
want: func() (s []*sppb.PartialResultSet) {
|
||
|
|
for i := 0; i < maxBuffers; i++ {
|
||
|
|
s = append(s, &sppb.PartialResultSet{
|
||
|
|
Metadata: kvMeta,
|
||
|
|
Values: []*proto3.Value{
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: keyStr(i)}},
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: valStr(i)}},
|
||
|
|
},
|
||
|
|
})
|
||
|
|
}
|
||
|
|
return s
|
||
|
|
}(),
|
||
|
|
stateHistory: func() (s []resumableStreamDecoderState) {
|
||
|
|
s = append(s, queueingRetryable) // RPC
|
||
|
|
for i := 0; i < maxBuffers; i++ {
|
||
|
|
s = append(s, queueingRetryable) // internal queue of resumableStreamDecoder fills up
|
||
|
|
}
|
||
|
|
s = append(s, queueingUnretryable) // last row triggers state change
|
||
|
|
s = append(s, finished) // query finishes
|
||
|
|
return s
|
||
|
|
}(),
|
||
|
|
},
|
||
|
|
}
|
||
|
|
for _, test := range tests {
|
||
|
|
ms := testutil.NewMockCloudSpanner(t, trxTs)
|
||
|
|
ms.Serve()
|
||
|
|
cc := dialMock(t, ms)
|
||
|
|
mc := sppb.NewSpannerClient(cc)
|
||
|
|
if test.rpc == nil {
|
||
|
|
// Avoid using test.sql directly in closure because for loop changes test.
|
||
|
|
sql := test.sql
|
||
|
|
test.rpc = func(ct context.Context, resumeToken []byte) (streamingReceiver, error) {
|
||
|
|
return mc.ExecuteStreamingSql(ct, &sppb.ExecuteSqlRequest{
|
||
|
|
Sql: sql,
|
||
|
|
ResumeToken: resumeToken,
|
||
|
|
})
|
||
|
|
}
|
||
|
|
}
|
||
|
|
ctx, cancel := context.WithCancel(context.Background())
|
||
|
|
defer cancel()
|
||
|
|
r := newResumableStreamDecoder(
|
||
|
|
ctx,
|
||
|
|
test.rpc,
|
||
|
|
)
|
||
|
|
// Override backoff to make the test run faster.
|
||
|
|
r.backoff = exponentialBackoff{1 * time.Nanosecond, 1 * time.Nanosecond}
|
||
|
|
// st is the set of observed state transitions.
|
||
|
|
st := []resumableStreamDecoderState{}
|
||
|
|
// q is the content of the decoder's partial result queue when expected number of state transitions are done.
|
||
|
|
q := []*sppb.PartialResultSet{}
|
||
|
|
var lastErr error
|
||
|
|
// Once the expected number of state transitions are observed,
|
||
|
|
// send a signal to channel stateDone.
|
||
|
|
stateDone := make(chan int)
|
||
|
|
// Set stateWitness to listen to state changes.
|
||
|
|
hl := len(test.stateHistory) // To avoid data race on test.
|
||
|
|
r.stateWitness = func(rs resumableStreamDecoderState) {
|
||
|
|
select {
|
||
|
|
case <-stateDone:
|
||
|
|
// Noop after expected number of state transitions
|
||
|
|
default:
|
||
|
|
// Record state transitions.
|
||
|
|
st = append(st, rs)
|
||
|
|
if len(st) == hl {
|
||
|
|
lastErr = r.lastErr()
|
||
|
|
q = r.q.dump()
|
||
|
|
close(stateDone)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
// Let mock server stream given messages to resumableStreamDecoder.
|
||
|
|
for _, m := range test.msgs {
|
||
|
|
ms.AddMsg(m.Err, m.ResumeToken)
|
||
|
|
}
|
||
|
|
var rs []*sppb.PartialResultSet
|
||
|
|
go func() {
|
||
|
|
for {
|
||
|
|
if !r.next() {
|
||
|
|
// Note that r.Next also exits on context cancel/timeout.
|
||
|
|
return
|
||
|
|
}
|
||
|
|
rs = append(rs, r.get())
|
||
|
|
}
|
||
|
|
}()
|
||
|
|
// Verify that resumableStreamDecoder reaches expected state.
|
||
|
|
select {
|
||
|
|
case <-stateDone: // Note that at this point, receiver is still blocking on r.next().
|
||
|
|
// Check if resumableStreamDecoder carried out expected
|
||
|
|
// state transitions.
|
||
|
|
if !testEqual(st, test.stateHistory) {
|
||
|
|
t.Errorf("%v: observed state transitions: \n%v\n, want \n%v\n",
|
||
|
|
test.name, st, test.stateHistory)
|
||
|
|
}
|
||
|
|
// Check if resumableStreamDecoder returns expected array of
|
||
|
|
// PartialResultSets.
|
||
|
|
if !testEqual(rs, test.want) {
|
||
|
|
t.Errorf("%v: received PartialResultSets: \n%v\n, want \n%v\n", test.name, rs, test.want)
|
||
|
|
}
|
||
|
|
// Verify that resumableStreamDecoder's internal buffering is also correct.
|
||
|
|
if !testEqual(q, test.queue) {
|
||
|
|
t.Errorf("%v: PartialResultSets still queued: \n%v\n, want \n%v\n", test.name, q, test.queue)
|
||
|
|
}
|
||
|
|
// Verify resume token.
|
||
|
|
if test.resumeToken != nil && !testEqual(r.resumeToken, test.resumeToken) {
|
||
|
|
t.Errorf("%v: Resume token is %v, want %v\n", test.name, r.resumeToken, test.resumeToken)
|
||
|
|
}
|
||
|
|
// Verify error message.
|
||
|
|
if !testEqual(lastErr, test.wantErr) {
|
||
|
|
t.Errorf("%v: got error %v, want %v", test.name, lastErr, test.wantErr)
|
||
|
|
}
|
||
|
|
case <-time.After(1 * time.Second):
|
||
|
|
t.Errorf("%v: Timeout in waiting for state change", test.name)
|
||
|
|
}
|
||
|
|
ms.Stop()
|
||
|
|
cc.Close()
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// sReceiver signals every receiving attempt through a channel,
|
||
|
|
// used by TestResumeToken to determine if the receiving of a certain
|
||
|
|
// PartialResultSet will be attempted next.
|
||
|
|
type sReceiver struct {
|
||
|
|
c chan int
|
||
|
|
rpcReceiver sppb.Spanner_ExecuteStreamingSqlClient
|
||
|
|
}
|
||
|
|
|
||
|
|
// Recv() implements streamingReceiver.Recv for sReceiver.
|
||
|
|
func (sr *sReceiver) Recv() (*sppb.PartialResultSet, error) {
|
||
|
|
sr.c <- 1
|
||
|
|
return sr.rpcReceiver.Recv()
|
||
|
|
}
|
||
|
|
|
||
|
|
// waitn waits for nth receiving attempt from now on, until
|
||
|
|
// the signal for nth Recv() attempts is received or timeout.
|
||
|
|
// Note that because the way stream() works, the signal for the
|
||
|
|
// nth Recv() means that the previous n - 1 PartialResultSets
|
||
|
|
// has already been returned to caller or queued, if no error happened.
|
||
|
|
func (sr *sReceiver) waitn(n int) error {
|
||
|
|
for i := 0; i < n; i++ {
|
||
|
|
select {
|
||
|
|
case <-sr.c:
|
||
|
|
case <-time.After(10 * time.Second):
|
||
|
|
return fmt.Errorf("timeout in waiting for %v-th Recv()", i+1)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
|
||
|
|
// Test the handling of resumableStreamDecoder.bytesBetweenResumeTokens.
|
||
|
|
func TestQueueBytes(t *testing.T) {
|
||
|
|
restore := setMaxBytesBetweenResumeTokens()
|
||
|
|
defer restore()
|
||
|
|
ms := testutil.NewMockCloudSpanner(t, trxTs)
|
||
|
|
ms.Serve()
|
||
|
|
defer ms.Stop()
|
||
|
|
cc := dialMock(t, ms)
|
||
|
|
defer cc.Close()
|
||
|
|
mc := sppb.NewSpannerClient(cc)
|
||
|
|
sr := &sReceiver{
|
||
|
|
c: make(chan int, 1000), // will never block in this test
|
||
|
|
}
|
||
|
|
wantQueueBytes := 0
|
||
|
|
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||
|
|
defer cancel()
|
||
|
|
r := newResumableStreamDecoder(
|
||
|
|
ctx,
|
||
|
|
func(ct context.Context, resumeToken []byte) (streamingReceiver, error) {
|
||
|
|
r, err := mc.ExecuteStreamingSql(ct, &sppb.ExecuteSqlRequest{
|
||
|
|
Sql: "SELECT t.key key, t.value value FROM t_mock t",
|
||
|
|
ResumeToken: resumeToken,
|
||
|
|
})
|
||
|
|
sr.rpcReceiver = r
|
||
|
|
return sr, err
|
||
|
|
},
|
||
|
|
)
|
||
|
|
go func() {
|
||
|
|
for r.next() {
|
||
|
|
}
|
||
|
|
}()
|
||
|
|
// Let server send maxBuffers / 2 rows.
|
||
|
|
for i := 0; i < maxBuffers/2; i++ {
|
||
|
|
wantQueueBytes += proto.Size(&sppb.PartialResultSet{
|
||
|
|
Metadata: kvMeta,
|
||
|
|
Values: []*proto3.Value{
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: keyStr(i)}},
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: valStr(i)}},
|
||
|
|
},
|
||
|
|
})
|
||
|
|
ms.AddMsg(nil, false)
|
||
|
|
}
|
||
|
|
if err := sr.waitn(maxBuffers/2 + 1); err != nil {
|
||
|
|
t.Fatalf("failed to wait for the first %v recv() calls: %v", maxBuffers, err)
|
||
|
|
}
|
||
|
|
if int32(wantQueueBytes) != r.bytesBetweenResumeTokens {
|
||
|
|
t.Errorf("r.bytesBetweenResumeTokens = %v, want %v", r.bytesBetweenResumeTokens, wantQueueBytes)
|
||
|
|
}
|
||
|
|
// Now send a resume token to drain the queue.
|
||
|
|
ms.AddMsg(nil, true)
|
||
|
|
// Wait for all rows to be processes.
|
||
|
|
if err := sr.waitn(1); err != nil {
|
||
|
|
t.Fatalf("failed to wait for rows to be processed: %v", err)
|
||
|
|
}
|
||
|
|
if r.bytesBetweenResumeTokens != 0 {
|
||
|
|
t.Errorf("r.bytesBetweenResumeTokens = %v, want 0", r.bytesBetweenResumeTokens)
|
||
|
|
}
|
||
|
|
// Let server send maxBuffers - 1 rows.
|
||
|
|
wantQueueBytes = 0
|
||
|
|
for i := 0; i < maxBuffers-1; i++ {
|
||
|
|
wantQueueBytes += proto.Size(&sppb.PartialResultSet{
|
||
|
|
Metadata: kvMeta,
|
||
|
|
Values: []*proto3.Value{
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: keyStr(i)}},
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: valStr(i)}},
|
||
|
|
},
|
||
|
|
})
|
||
|
|
ms.AddMsg(nil, false)
|
||
|
|
}
|
||
|
|
if err := sr.waitn(maxBuffers - 1); err != nil {
|
||
|
|
t.Fatalf("failed to wait for %v rows to be processed: %v", maxBuffers-1, err)
|
||
|
|
}
|
||
|
|
if int32(wantQueueBytes) != r.bytesBetweenResumeTokens {
|
||
|
|
t.Errorf("r.bytesBetweenResumeTokens = %v, want 0", r.bytesBetweenResumeTokens)
|
||
|
|
}
|
||
|
|
// Trigger a state transition: queueingRetryable -> queueingUnretryable.
|
||
|
|
ms.AddMsg(nil, false)
|
||
|
|
if err := sr.waitn(1); err != nil {
|
||
|
|
t.Fatalf("failed to wait for state transition: %v", err)
|
||
|
|
}
|
||
|
|
if r.bytesBetweenResumeTokens != 0 {
|
||
|
|
t.Errorf("r.bytesBetweenResumeTokens = %v, want 0", r.bytesBetweenResumeTokens)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// Verify that client can deal with resume token correctly
|
||
|
|
func TestResumeToken(t *testing.T) {
|
||
|
|
restore := setMaxBytesBetweenResumeTokens()
|
||
|
|
defer restore()
|
||
|
|
ms := testutil.NewMockCloudSpanner(t, trxTs)
|
||
|
|
ms.Serve()
|
||
|
|
defer ms.Stop()
|
||
|
|
cc := dialMock(t, ms)
|
||
|
|
defer cc.Close()
|
||
|
|
mc := sppb.NewSpannerClient(cc)
|
||
|
|
sr := &sReceiver{
|
||
|
|
c: make(chan int, 1000), // will never block in this test
|
||
|
|
}
|
||
|
|
rows := []*Row{}
|
||
|
|
done := make(chan error)
|
||
|
|
streaming := func() {
|
||
|
|
// Establish a stream to mock cloud spanner server.
|
||
|
|
iter := stream(context.Background(),
|
||
|
|
func(ct context.Context, resumeToken []byte) (streamingReceiver, error) {
|
||
|
|
r, err := mc.ExecuteStreamingSql(ct, &sppb.ExecuteSqlRequest{
|
||
|
|
Sql: "SELECT t.key key, t.value value FROM t_mock t",
|
||
|
|
ResumeToken: resumeToken,
|
||
|
|
})
|
||
|
|
sr.rpcReceiver = r
|
||
|
|
return sr, err
|
||
|
|
},
|
||
|
|
nil,
|
||
|
|
func(error) {})
|
||
|
|
defer iter.Stop()
|
||
|
|
var err error
|
||
|
|
for {
|
||
|
|
var row *Row
|
||
|
|
row, err = iter.Next()
|
||
|
|
if err == iterator.Done {
|
||
|
|
err = nil
|
||
|
|
break
|
||
|
|
}
|
||
|
|
if err != nil {
|
||
|
|
break
|
||
|
|
}
|
||
|
|
rows = append(rows, row)
|
||
|
|
}
|
||
|
|
done <- err
|
||
|
|
}
|
||
|
|
go streaming()
|
||
|
|
// Server streaming row 0 - 2, only row 1 has resume token.
|
||
|
|
// Client will receive row 0 - 2, so it will try receiving for
|
||
|
|
// 4 times (the last recv will block), and only row 0 - 1 will
|
||
|
|
// be yielded.
|
||
|
|
for i := 0; i < 3; i++ {
|
||
|
|
if i == 1 {
|
||
|
|
ms.AddMsg(nil, true)
|
||
|
|
} else {
|
||
|
|
ms.AddMsg(nil, false)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
// Wait for 4 receive attempts, as explained above.
|
||
|
|
if err := sr.waitn(4); err != nil {
|
||
|
|
t.Fatalf("failed to wait for row 0 - 2: %v", err)
|
||
|
|
}
|
||
|
|
want := []*Row{
|
||
|
|
{
|
||
|
|
fields: kvMeta.RowType.Fields,
|
||
|
|
vals: []*proto3.Value{
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: keyStr(0)}},
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: valStr(0)}},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
{
|
||
|
|
fields: kvMeta.RowType.Fields,
|
||
|
|
vals: []*proto3.Value{
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: keyStr(1)}},
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: valStr(1)}},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
}
|
||
|
|
if !testEqual(rows, want) {
|
||
|
|
t.Errorf("received rows: \n%v\n; but want\n%v\n", rows, want)
|
||
|
|
}
|
||
|
|
// Inject resumable failure.
|
||
|
|
ms.AddMsg(
|
||
|
|
status.Errorf(codes.Unavailable, "mock server unavailable"),
|
||
|
|
false,
|
||
|
|
)
|
||
|
|
// Test if client detects the resumable failure and retries.
|
||
|
|
if err := sr.waitn(1); err != nil {
|
||
|
|
t.Fatalf("failed to wait for client to retry: %v", err)
|
||
|
|
}
|
||
|
|
// Client has resumed the query, now server resend row 2.
|
||
|
|
ms.AddMsg(nil, true)
|
||
|
|
if err := sr.waitn(1); err != nil {
|
||
|
|
t.Fatalf("failed to wait for resending row 2: %v", err)
|
||
|
|
}
|
||
|
|
// Now client should have received row 0 - 2.
|
||
|
|
want = append(want, &Row{
|
||
|
|
fields: kvMeta.RowType.Fields,
|
||
|
|
vals: []*proto3.Value{
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: keyStr(2)}},
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: valStr(2)}},
|
||
|
|
},
|
||
|
|
})
|
||
|
|
if !testEqual(rows, want) {
|
||
|
|
t.Errorf("received rows: \n%v\n, want\n%v\n", rows, want)
|
||
|
|
}
|
||
|
|
// Sending 3rd - (maxBuffers+1)th rows without resume tokens, client should buffer them.
|
||
|
|
for i := 3; i < maxBuffers+2; i++ {
|
||
|
|
ms.AddMsg(nil, false)
|
||
|
|
}
|
||
|
|
if err := sr.waitn(maxBuffers - 1); err != nil {
|
||
|
|
t.Fatalf("failed to wait for row 3-%v: %v", maxBuffers+1, err)
|
||
|
|
}
|
||
|
|
// Received rows should be unchanged.
|
||
|
|
if !testEqual(rows, want) {
|
||
|
|
t.Errorf("receive rows: \n%v\n, want\n%v\n", rows, want)
|
||
|
|
}
|
||
|
|
// Send (maxBuffers+2)th row to trigger state change of resumableStreamDecoder:
|
||
|
|
// queueingRetryable -> queueingUnretryable
|
||
|
|
ms.AddMsg(nil, false)
|
||
|
|
if err := sr.waitn(1); err != nil {
|
||
|
|
t.Fatalf("failed to wait for row %v: %v", maxBuffers+2, err)
|
||
|
|
}
|
||
|
|
// Client should yield row 3rd - (maxBuffers+2)th to application. Therefore, application should
|
||
|
|
// see row 0 - (maxBuffers+2)th so far.
|
||
|
|
for i := 3; i < maxBuffers+3; i++ {
|
||
|
|
want = append(want, &Row{
|
||
|
|
fields: kvMeta.RowType.Fields,
|
||
|
|
vals: []*proto3.Value{
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: keyStr(i)}},
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: valStr(i)}},
|
||
|
|
},
|
||
|
|
})
|
||
|
|
}
|
||
|
|
if !testEqual(rows, want) {
|
||
|
|
t.Errorf("received rows: \n%v\n; want\n%v\n", rows, want)
|
||
|
|
}
|
||
|
|
// Inject resumable error, but since resumableStreamDecoder is already at queueingUnretryable
|
||
|
|
// state, query will just fail.
|
||
|
|
ms.AddMsg(
|
||
|
|
status.Errorf(codes.Unavailable, "mock server wants some sleep"),
|
||
|
|
false,
|
||
|
|
)
|
||
|
|
var gotErr error
|
||
|
|
select {
|
||
|
|
case gotErr = <-done:
|
||
|
|
case <-time.After(10 * time.Second):
|
||
|
|
t.Fatalf("timeout in waiting for failed query to return.")
|
||
|
|
}
|
||
|
|
if wantErr := toSpannerError(status.Errorf(codes.Unavailable, "mock server wants some sleep")); !testEqual(gotErr, wantErr) {
|
||
|
|
t.Fatalf("stream() returns error: %v, but want error: %v", gotErr, wantErr)
|
||
|
|
}
|
||
|
|
|
||
|
|
// Reconnect to mock Cloud Spanner.
|
||
|
|
rows = []*Row{}
|
||
|
|
go streaming()
|
||
|
|
// Let server send two rows without resume token.
|
||
|
|
for i := maxBuffers + 3; i < maxBuffers+5; i++ {
|
||
|
|
ms.AddMsg(nil, false)
|
||
|
|
}
|
||
|
|
if err := sr.waitn(3); err != nil {
|
||
|
|
t.Fatalf("failed to wait for row %v - %v: %v", maxBuffers+3, maxBuffers+5, err)
|
||
|
|
}
|
||
|
|
if len(rows) > 0 {
|
||
|
|
t.Errorf("client received some rows unexpectedly: %v, want nothing", rows)
|
||
|
|
}
|
||
|
|
// Let server end the query.
|
||
|
|
ms.AddMsg(io.EOF, false)
|
||
|
|
select {
|
||
|
|
case gotErr = <-done:
|
||
|
|
case <-time.After(10 * time.Second):
|
||
|
|
t.Fatalf("timeout in waiting for failed query to return")
|
||
|
|
}
|
||
|
|
if gotErr != nil {
|
||
|
|
t.Fatalf("stream() returns unexpected error: %v, but want no error", gotErr)
|
||
|
|
}
|
||
|
|
// Verify if a normal server side EOF flushes all queued rows.
|
||
|
|
want = []*Row{
|
||
|
|
{
|
||
|
|
fields: kvMeta.RowType.Fields,
|
||
|
|
vals: []*proto3.Value{
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: keyStr(maxBuffers + 3)}},
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: valStr(maxBuffers + 3)}},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
{
|
||
|
|
fields: kvMeta.RowType.Fields,
|
||
|
|
vals: []*proto3.Value{
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: keyStr(maxBuffers + 4)}},
|
||
|
|
{Kind: &proto3.Value_StringValue{StringValue: valStr(maxBuffers + 4)}},
|
||
|
|
},
|
||
|
|
},
|
||
|
|
}
|
||
|
|
if !testEqual(rows, want) {
|
||
|
|
t.Errorf("received rows: \n%v\n; but want\n%v\n", rows, want)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// Verify that streaming query get retried upon real gRPC server transport failures.
|
||
|
|
func TestGrpcReconnect(t *testing.T) {
|
||
|
|
restore := setMaxBytesBetweenResumeTokens()
|
||
|
|
defer restore()
|
||
|
|
ms := testutil.NewMockCloudSpanner(t, trxTs)
|
||
|
|
ms.Serve()
|
||
|
|
defer ms.Stop()
|
||
|
|
cc := dialMock(t, ms)
|
||
|
|
defer cc.Close()
|
||
|
|
mc := sppb.NewSpannerClient(cc)
|
||
|
|
retry := make(chan int)
|
||
|
|
row := make(chan int)
|
||
|
|
var err error
|
||
|
|
go func() {
|
||
|
|
r := 0
|
||
|
|
// Establish a stream to mock cloud spanner server.
|
||
|
|
iter := stream(context.Background(),
|
||
|
|
func(ct context.Context, resumeToken []byte) (streamingReceiver, error) {
|
||
|
|
if r > 0 {
|
||
|
|
// This RPC attempt is a retry, signal it.
|
||
|
|
retry <- r
|
||
|
|
}
|
||
|
|
r++
|
||
|
|
return mc.ExecuteStreamingSql(ct, &sppb.ExecuteSqlRequest{
|
||
|
|
Sql: "SELECT t.key key, t.value value FROM t_mock t",
|
||
|
|
ResumeToken: resumeToken,
|
||
|
|
})
|
||
|
|
|
||
|
|
},
|
||
|
|
nil,
|
||
|
|
func(error) {})
|
||
|
|
defer iter.Stop()
|
||
|
|
for {
|
||
|
|
_, err = iter.Next()
|
||
|
|
if err == iterator.Done {
|
||
|
|
err = nil
|
||
|
|
break
|
||
|
|
}
|
||
|
|
if err != nil {
|
||
|
|
break
|
||
|
|
}
|
||
|
|
row <- 0
|
||
|
|
}
|
||
|
|
}()
|
||
|
|
// Add a message and wait for the receipt.
|
||
|
|
ms.AddMsg(nil, true)
|
||
|
|
select {
|
||
|
|
case <-row:
|
||
|
|
case <-time.After(10 * time.Second):
|
||
|
|
t.Fatalf("expect stream to be established within 10 seconds, but it didn't")
|
||
|
|
}
|
||
|
|
// Error injection: force server to close all connections.
|
||
|
|
ms.Stop()
|
||
|
|
// Test to see if client respond to the real RPC failure correctly by
|
||
|
|
// retrying RPC.
|
||
|
|
select {
|
||
|
|
case r, ok := <-retry:
|
||
|
|
if ok && r == 1 {
|
||
|
|
break
|
||
|
|
}
|
||
|
|
t.Errorf("retry count = %v, want 1", r)
|
||
|
|
case <-time.After(10 * time.Second):
|
||
|
|
t.Errorf("client library failed to respond after 10 seconds, aborting")
|
||
|
|
return
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// Test cancel/timeout for client operations.
|
||
|
|
func TestCancelTimeout(t *testing.T) {
|
||
|
|
restore := setMaxBytesBetweenResumeTokens()
|
||
|
|
defer restore()
|
||
|
|
ms := testutil.NewMockCloudSpanner(t, trxTs)
|
||
|
|
ms.Serve()
|
||
|
|
defer ms.Stop()
|
||
|
|
cc := dialMock(t, ms)
|
||
|
|
defer cc.Close()
|
||
|
|
mc := sppb.NewSpannerClient(cc)
|
||
|
|
done := make(chan int)
|
||
|
|
go func() {
|
||
|
|
for {
|
||
|
|
ms.AddMsg(nil, true)
|
||
|
|
}
|
||
|
|
}()
|
||
|
|
// Test cancelling query.
|
||
|
|
ctx, cancel := context.WithCancel(context.Background())
|
||
|
|
var err error
|
||
|
|
go func() {
|
||
|
|
// Establish a stream to mock cloud spanner server.
|
||
|
|
iter := stream(ctx,
|
||
|
|
func(ct context.Context, resumeToken []byte) (streamingReceiver, error) {
|
||
|
|
return mc.ExecuteStreamingSql(ct, &sppb.ExecuteSqlRequest{
|
||
|
|
Sql: "SELECT t.key key, t.value value FROM t_mock t",
|
||
|
|
ResumeToken: resumeToken,
|
||
|
|
})
|
||
|
|
},
|
||
|
|
nil,
|
||
|
|
func(error) {})
|
||
|
|
defer iter.Stop()
|
||
|
|
for {
|
||
|
|
_, err = iter.Next()
|
||
|
|
if err == iterator.Done {
|
||
|
|
break
|
||
|
|
}
|
||
|
|
if err != nil {
|
||
|
|
done <- 0
|
||
|
|
break
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}()
|
||
|
|
cancel()
|
||
|
|
select {
|
||
|
|
case <-done:
|
||
|
|
if ErrCode(err) != codes.Canceled {
|
||
|
|
t.Errorf("streaming query is canceled and returns error %v, want error code %v", err, codes.Canceled)
|
||
|
|
}
|
||
|
|
case <-time.After(1 * time.Second):
|
||
|
|
t.Errorf("query doesn't exit timely after being cancelled")
|
||
|
|
}
|
||
|
|
// Test query timeout.
|
||
|
|
ctx, cancel = context.WithTimeout(context.Background(), 1*time.Second)
|
||
|
|
go func() {
|
||
|
|
// Establish a stream to mock cloud spanner server.
|
||
|
|
iter := stream(ctx,
|
||
|
|
func(ct context.Context, resumeToken []byte) (streamingReceiver, error) {
|
||
|
|
return mc.ExecuteStreamingSql(ct, &sppb.ExecuteSqlRequest{
|
||
|
|
Sql: "SELECT t.key key, t.value value FROM t_mock t",
|
||
|
|
ResumeToken: resumeToken,
|
||
|
|
})
|
||
|
|
},
|
||
|
|
nil,
|
||
|
|
func(error) {})
|
||
|
|
defer iter.Stop()
|
||
|
|
for {
|
||
|
|
_, err = iter.Next()
|
||
|
|
if err == iterator.Done {
|
||
|
|
err = nil
|
||
|
|
break
|
||
|
|
}
|
||
|
|
if err != nil {
|
||
|
|
break
|
||
|
|
}
|
||
|
|
}
|
||
|
|
done <- 0
|
||
|
|
}()
|
||
|
|
select {
|
||
|
|
case <-done:
|
||
|
|
if wantErr := codes.DeadlineExceeded; ErrCode(err) != wantErr {
|
||
|
|
t.Errorf("streaming query timeout returns error %v, want error code %v", err, wantErr)
|
||
|
|
}
|
||
|
|
case <-time.After(2 * time.Second):
|
||
|
|
t.Errorf("query doesn't timeout as expected")
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestRowIteratorDo(t *testing.T) {
|
||
|
|
restore := setMaxBytesBetweenResumeTokens()
|
||
|
|
defer restore()
|
||
|
|
ms := testutil.NewMockCloudSpanner(t, trxTs)
|
||
|
|
ms.Serve()
|
||
|
|
defer ms.Stop()
|
||
|
|
cc := dialMock(t, ms)
|
||
|
|
defer cc.Close()
|
||
|
|
mc := sppb.NewSpannerClient(cc)
|
||
|
|
|
||
|
|
for i := 0; i < 3; i++ {
|
||
|
|
ms.AddMsg(nil, false)
|
||
|
|
}
|
||
|
|
ms.AddMsg(io.EOF, true)
|
||
|
|
nRows := 0
|
||
|
|
iter := stream(context.Background(),
|
||
|
|
func(ct context.Context, resumeToken []byte) (streamingReceiver, error) {
|
||
|
|
return mc.ExecuteStreamingSql(ct, &sppb.ExecuteSqlRequest{
|
||
|
|
Sql: "SELECT t.key key, t.value value FROM t_mock t",
|
||
|
|
ResumeToken: resumeToken,
|
||
|
|
})
|
||
|
|
},
|
||
|
|
nil,
|
||
|
|
func(error) {})
|
||
|
|
err := iter.Do(func(r *Row) error { nRows++; return nil })
|
||
|
|
if err != nil {
|
||
|
|
t.Errorf("Using Do: %v", err)
|
||
|
|
}
|
||
|
|
if nRows != 3 {
|
||
|
|
t.Errorf("got %d rows, want 3", nRows)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestRowIteratorDoWithError(t *testing.T) {
|
||
|
|
restore := setMaxBytesBetweenResumeTokens()
|
||
|
|
defer restore()
|
||
|
|
ms := testutil.NewMockCloudSpanner(t, trxTs)
|
||
|
|
ms.Serve()
|
||
|
|
defer ms.Stop()
|
||
|
|
cc := dialMock(t, ms)
|
||
|
|
defer cc.Close()
|
||
|
|
mc := sppb.NewSpannerClient(cc)
|
||
|
|
|
||
|
|
for i := 0; i < 3; i++ {
|
||
|
|
ms.AddMsg(nil, false)
|
||
|
|
}
|
||
|
|
ms.AddMsg(io.EOF, true)
|
||
|
|
iter := stream(context.Background(),
|
||
|
|
func(ct context.Context, resumeToken []byte) (streamingReceiver, error) {
|
||
|
|
return mc.ExecuteStreamingSql(ct, &sppb.ExecuteSqlRequest{
|
||
|
|
Sql: "SELECT t.key key, t.value value FROM t_mock t",
|
||
|
|
ResumeToken: resumeToken,
|
||
|
|
})
|
||
|
|
},
|
||
|
|
nil,
|
||
|
|
func(error) {})
|
||
|
|
injected := errors.New("Failed iterator")
|
||
|
|
err := iter.Do(func(r *Row) error { return injected })
|
||
|
|
if err != injected {
|
||
|
|
t.Errorf("got <%v>, want <%v>", err, injected)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestIteratorStopEarly(t *testing.T) {
|
||
|
|
ctx := context.Background()
|
||
|
|
restore := setMaxBytesBetweenResumeTokens()
|
||
|
|
defer restore()
|
||
|
|
ms := testutil.NewMockCloudSpanner(t, trxTs)
|
||
|
|
ms.Serve()
|
||
|
|
defer ms.Stop()
|
||
|
|
cc := dialMock(t, ms)
|
||
|
|
defer cc.Close()
|
||
|
|
mc := sppb.NewSpannerClient(cc)
|
||
|
|
|
||
|
|
ms.AddMsg(nil, false)
|
||
|
|
ms.AddMsg(nil, false)
|
||
|
|
ms.AddMsg(io.EOF, true)
|
||
|
|
|
||
|
|
iter := stream(ctx,
|
||
|
|
func(ct context.Context, resumeToken []byte) (streamingReceiver, error) {
|
||
|
|
return mc.ExecuteStreamingSql(ct, &sppb.ExecuteSqlRequest{
|
||
|
|
Sql: "SELECT t.key key, t.value value FROM t_mock t",
|
||
|
|
ResumeToken: resumeToken,
|
||
|
|
})
|
||
|
|
},
|
||
|
|
nil,
|
||
|
|
func(error) {})
|
||
|
|
_, err := iter.Next()
|
||
|
|
if err != nil {
|
||
|
|
t.Fatalf("before Stop: %v", err)
|
||
|
|
}
|
||
|
|
iter.Stop()
|
||
|
|
// Stop sets r.err to the FailedPrecondition error "Next called after Stop".
|
||
|
|
// Override that here so this test can observe the Canceled error from the stream.
|
||
|
|
iter.err = nil
|
||
|
|
iter.Next()
|
||
|
|
if ErrCode(iter.streamd.lastErr()) != codes.Canceled {
|
||
|
|
t.Errorf("after Stop: got %v, wanted Canceled", err)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestIteratorWithError(t *testing.T) {
|
||
|
|
injected := errors.New("Failed iterator")
|
||
|
|
iter := RowIterator{err: injected}
|
||
|
|
defer iter.Stop()
|
||
|
|
if _, err := iter.Next(); err != injected {
|
||
|
|
t.Fatalf("Expected error: %v, got %v", injected, err)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func dialMock(t *testing.T, ms *testutil.MockCloudSpanner) *grpc.ClientConn {
|
||
|
|
cc, err := grpc.Dial(ms.Addr(), grpc.WithInsecure(), grpc.WithBlock())
|
||
|
|
if err != nil {
|
||
|
|
t.Fatalf("Dial(%q) = %v", ms.Addr(), err)
|
||
|
|
}
|
||
|
|
return cc
|
||
|
|
}
|