Skip to content

Commit 4e04b7e

Browse files
olavloiteolavloite
authored
fix(spanner): enforce only one resource header (#12618)
Only one resource header (google-cloud-resource-prefix) should be included in requests that are sent to Spanner. If the context that is used for a request already contains a resource header, then that header value will be removed, and only the resource header that is set by the Spanner client will be included in the request.
1 parent 420edfe commit 4e04b7e

5 files changed

Lines changed: 98 additions & 24 deletions

File tree

‎spanner/client.go‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -390,6 +390,8 @@ type openTelemetryConfig struct {
390390
func contextWithOutgoingMetadata(ctx context.Context, md metadata.MD, disableRouteToLeader bool) context.Context {
391391
existing, ok := metadata.FromOutgoingContext(ctx)
392392
if ok {
393+
// Make sure that we only send one resource header.
394+
existing.Delete(resourcePrefixHeader)
393395
md = metadata.Join(existing, md)
394396
}
395397
if !disableRouteToLeader {

‎spanner/client_test.go‎

Lines changed: 49 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ import (
2626
"net"
2727
"os"
2828
"reflect"
29+
"regexp"
2930
"strconv"
3031
"strings"
3132
"sync"
@@ -93,6 +94,19 @@ func setupMockedTestServerWithConfigAndGCPMultiendpointPool(t *testing.T, config
9394
},
9495
},
9596
}
97+
expectedResourceHeaderFormat := regexp.MustCompile("projects/.+/instances/.+/databases/.+.*")
98+
grpcHeaderChecker.Checkers = append(grpcHeaderChecker.Checkers, &itestutil.HeaderChecker{
99+
Key: resourcePrefixHeader,
100+
ValuesValidator: func(token ...string) error {
101+
if len(token) != 1 {
102+
return status.Errorf(codes.Internal, "unexpected number of resource headers: %v", len(token))
103+
}
104+
if !expectedResourceHeaderFormat.MatchString(token[0]) {
105+
return status.Errorf(codes.Internal, "invalid resource header value: %v", token[0])
106+
}
107+
return nil
108+
},
109+
})
96110
if config.Compression == gzip.Name {
97111
grpcHeaderChecker.Checkers = append(grpcHeaderChecker.Checkers, &itestutil.HeaderChecker{
98112
Key: "x-response-encoding",
@@ -658,6 +672,30 @@ func TestClient_Single(t *testing.T) {
658672
}
659673
}
660674

675+
func TestClient_NonConformingHeader(t *testing.T) {
676+
t.Parallel()
677+
ctx := context.Background()
678+
md := metadata.Pairs(resourcePrefixHeader, "projects/foo/documents/bar")
679+
ctx = metadata.NewOutgoingContext(ctx, md)
680+
err := testSingleQueryWithContext(ctx, t, nil)
681+
if err != nil {
682+
t.Fatal(err)
683+
}
684+
// Verify that even though the request is sent without the non-conforming header,
685+
// it is still present in the original context.
686+
if md, ok := metadata.FromOutgoingContext(ctx); ok {
687+
header := md.Get(resourcePrefixHeader)
688+
if g, w := len(header), 1; g != w {
689+
t.Fatalf("header length mismatch\n Got: %v\nWant: %v", g, w)
690+
}
691+
if g, w := header[0], "projects/foo/documents/bar"; g != w {
692+
t.Fatalf("header mismatch\n Got: %v\nWant: %v", g, w)
693+
}
694+
} else {
695+
t.Fatal("could not get metadata from context")
696+
}
697+
}
698+
661699
func TestClient_Single_Unavailable(t *testing.T) {
662700
t.Parallel()
663701
err := testSingleQuery(t, serverErrorWithMinimalRetryDelay(codes.Unavailable, "Temporary unavailable"))
@@ -1191,8 +1229,11 @@ func checkReqsForTransactionOptions(t *testing.T, server InMemSpannerServer, txo
11911229
}
11921230

11931231
func testSingleQuery(t *testing.T, serverError error) error {
1194-
ctx := context.Background()
1195-
server, client, teardown := setupMockedTestServerWithConfig(t, ClientConfig{SessionPoolConfig: SessionPoolConfig{MinOpened: 1}})
1232+
return testSingleQueryWithContext(context.Background(), t, serverError)
1233+
}
1234+
1235+
func testSingleQueryWithContext(ctx context.Context, t *testing.T, serverError error) error {
1236+
server, client, teardown := setupMockedTestServerWithConfig(t, ClientConfig{DisableNativeMetrics: true, SessionPoolConfig: SessionPoolConfig{MinOpened: 1}})
11961237
defer teardown()
11971238
// Wait until all sessions have been created, so we know that those requests will not interfere with the test.
11981239
sp := client.idleSessions
@@ -1872,7 +1913,7 @@ func validateIsolationLevelForRWTransactions(t *testing.T, server *MockedSpanner
18721913

18731914
func TestClient_ReadWriteTransactionWithNoIsolationLevelForRWTransactionAtClientConfig(t *testing.T) {
18741915
t.Parallel()
1875-
server, teardown, err := testReadWriteTransactionWithConfig(t, ClientConfig{SessionPoolConfig: DefaultSessionPoolConfig}, make(map[string]SimulatedExecutionTime), 1)
1916+
server, teardown, err := testReadWriteTransactionWithConfig(t, ClientConfig{DisableNativeMetrics: true, SessionPoolConfig: DefaultSessionPoolConfig}, make(map[string]SimulatedExecutionTime), 1)
18761917
if err != nil {
18771918
t.Fatal(err)
18781919
}
@@ -1882,7 +1923,7 @@ func TestClient_ReadWriteTransactionWithNoIsolationLevelForRWTransactionAtClient
18821923

18831924
func TestClient_ReadWriteTransactionWithIsolationLevelForRWTransactionAtClientConfig(t *testing.T) {
18841925
t.Parallel()
1885-
server, teardown, err := testReadWriteTransactionWithConfig(t, ClientConfig{SessionPoolConfig: DefaultSessionPoolConfig, TransactionOptions: TransactionOptions{IsolationLevel: sppb.TransactionOptions_REPEATABLE_READ}}, make(map[string]SimulatedExecutionTime), 1)
1926+
server, teardown, err := testReadWriteTransactionWithConfig(t, ClientConfig{DisableNativeMetrics: true, SessionPoolConfig: DefaultSessionPoolConfig, TransactionOptions: TransactionOptions{IsolationLevel: sppb.TransactionOptions_REPEATABLE_READ}}, make(map[string]SimulatedExecutionTime), 1)
18861927
if err != nil {
18871928
t.Fatal(err)
18881929
}
@@ -1892,7 +1933,7 @@ func TestClient_ReadWriteTransactionWithIsolationLevelForRWTransactionAtClientCo
18921933

18931934
func TestClient_ReadWriteTransactionWithIsolationLevelForRWTransactionAtTransactionLevel(t *testing.T) {
18941935
t.Parallel()
1895-
server, teardown, err := testReadWriteTransactionWithConfigWithTransactionOptions(t, ClientConfig{SessionPoolConfig: DefaultSessionPoolConfig, TransactionOptions: TransactionOptions{IsolationLevel: sppb.TransactionOptions_SERIALIZABLE}}, TransactionOptions{IsolationLevel: sppb.TransactionOptions_REPEATABLE_READ}, make(map[string]SimulatedExecutionTime), 1)
1936+
server, teardown, err := testReadWriteTransactionWithConfigWithTransactionOptions(t, ClientConfig{DisableNativeMetrics: true, SessionPoolConfig: DefaultSessionPoolConfig, TransactionOptions: TransactionOptions{IsolationLevel: sppb.TransactionOptions_SERIALIZABLE}}, TransactionOptions{IsolationLevel: sppb.TransactionOptions_REPEATABLE_READ}, make(map[string]SimulatedExecutionTime), 1)
18961937
if err != nil {
18971938
t.Fatal(err)
18981939
}
@@ -2004,7 +2045,7 @@ func TestClient_ReadWriteStmtBasedTransactionWithIsolationLevelAtClientConfigLev
20042045
}
20052046

20062047
func testClientReadWriteStmtBasedTransactionWithIsolationLevelAtClientConfigLevel(t *testing.T, beginTransactionOption BeginTransactionOption) {
2007-
server, client, teardown := setupMockedTestServerWithConfig(t, ClientConfig{TransactionOptions: TransactionOptions{IsolationLevel: sppb.TransactionOptions_SERIALIZABLE}})
2048+
server, client, teardown := setupMockedTestServerWithConfig(t, ClientConfig{DisableNativeMetrics: true, TransactionOptions: TransactionOptions{IsolationLevel: sppb.TransactionOptions_SERIALIZABLE}})
20082049
defer teardown()
20092050
ctx := context.Background()
20102051
tx, err := NewReadWriteStmtBasedTransactionWithOptions(
@@ -2034,7 +2075,7 @@ func TestClient_ReadWriteStmtBasedTransactionWithNoIsolationLevelWithInlineBegin
20342075
}
20352076

20362077
func testClientReadWriteStmtBasedTransactionWithNoIsolationLevel(t *testing.T, beginTransactionOption BeginTransactionOption) {
2037-
server, client, teardown := setupMockedTestServerWithConfig(t, ClientConfig{TransactionOptions: TransactionOptions{}})
2078+
server, client, teardown := setupMockedTestServerWithConfig(t, ClientConfig{DisableNativeMetrics: true, TransactionOptions: TransactionOptions{}})
20382079
defer teardown()
20392080
ctx := context.Background()
20402081
tx, err := NewReadWriteStmtBasedTransactionWithOptions(
@@ -3734,7 +3775,7 @@ func TestClient_ReadWriteTransaction_WithCancelledContext(t *testing.T) {
37343775
}
37353776

37363777
func testReadWriteTransaction(t *testing.T, executionTimes map[string]SimulatedExecutionTime, expectedAttempts int) error {
3737-
_, teardown, err := testReadWriteTransactionWithConfig(t, ClientConfig{SessionPoolConfig: DefaultSessionPoolConfig}, executionTimes, expectedAttempts)
3778+
_, teardown, err := testReadWriteTransactionWithConfig(t, ClientConfig{DisableNativeMetrics: true, SessionPoolConfig: DefaultSessionPoolConfig}, executionTimes, expectedAttempts)
37383779
defer teardown()
37393780
return err
37403781
}

‎spanner/read_test.go‎

Lines changed: 37 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@ import (
3232
"google.golang.org/api/iterator"
3333
"google.golang.org/genproto/googleapis/rpc/errdetails"
3434
"google.golang.org/grpc/codes"
35+
"google.golang.org/grpc/metadata"
3536
"google.golang.org/grpc/status"
3637
"google.golang.org/protobuf/proto"
3738
"google.golang.org/protobuf/types/known/durationpb"
@@ -806,7 +807,9 @@ func TestRsdNonblockingStates(t *testing.T) {
806807
}, opts...)
807808
}
808809
}
809-
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
810+
md := metadata.Pairs(resourcePrefixHeader, "projects/p/instances/i/databases/d")
811+
ctx := metadata.NewOutgoingContext(context.Background(), md)
812+
ctx, cancel := context.WithTimeout(ctx, 10*time.Second)
810813
defer cancel()
811814
mt := c.metricsTracerFactory.createBuiltinMetricsTracer(ctx)
812815
r := newResumableStreamDecoder(
@@ -1108,7 +1111,9 @@ func TestRsdBlockingStates(t *testing.T) {
11081111
}, opts...)
11091112
}
11101113
}
1111-
ctx, cancel := context.WithCancel(context.Background())
1114+
md := metadata.Pairs(resourcePrefixHeader, "projects/p/instances/i/databases/d")
1115+
ctx := metadata.NewOutgoingContext(context.Background(), md)
1116+
ctx, cancel := context.WithCancel(ctx)
11121117
defer cancel()
11131118
mt := c.metricsTracerFactory.createBuiltinMetricsTracer(ctx)
11141119
r := newResumableStreamDecoder(
@@ -1276,7 +1281,9 @@ func TestQueueBytes(t *testing.T) {
12761281
sr := &sReceiver{
12771282
c: make(chan int, 1000), // will never block in this test
12781283
}
1279-
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
1284+
md := metadata.Pairs(resourcePrefixHeader, "projects/p/instances/i/databases/d")
1285+
ctx := metadata.NewOutgoingContext(context.Background(), md)
1286+
ctx, cancel := context.WithTimeout(ctx, 10*time.Second)
12801287
defer cancel()
12811288
mt := c.metricsTracerFactory.createBuiltinMetricsTracer(ctx)
12821289
decoder := newResumableStreamDecoder(
@@ -1379,8 +1386,10 @@ func TestResumeToken(t *testing.T) {
13791386
}
13801387
rows := []*Row{}
13811388

1389+
md := metadata.Pairs(resourcePrefixHeader, "projects/p/instances/i/databases/d")
1390+
ctx := metadata.NewOutgoingContext(context.Background(), md)
13821391
streaming := func() *RowIterator {
1383-
return stream(context.Background(), nil,
1392+
return stream(ctx, nil,
13841393
c.metricsTracerFactory,
13851394
func(ct context.Context, resumeToken []byte, opts ...gax.CallOption) (streamingReceiver, error) {
13861395
r, err := mc.ExecuteStreamingSql(ct, &sppb.ExecuteSqlRequest{
@@ -1523,10 +1532,12 @@ func TestGrpcReconnect(t *testing.T) {
15231532
},
15241533
)
15251534

1535+
md := metadata.Pairs(resourcePrefixHeader, "projects/p/instances/i/databases/d")
1536+
ctx := metadata.NewOutgoingContext(context.Background(), md)
15261537
// The retry is counted from the second call.
15271538
r := -1
15281539
// Establish a stream to mock cloud spanner server.
1529-
iter := stream(context.Background(), nil, c.metricsTracerFactory,
1540+
iter := stream(ctx, nil, c.metricsTracerFactory,
15301541
func(ct context.Context, resumeToken []byte, opts ...gax.CallOption) (streamingReceiver, error) {
15311542
r++
15321543
return mc.ExecuteStreamingSql(ct, &sppb.ExecuteSqlRequest{
@@ -1580,10 +1591,12 @@ func TestRetryResourceExhaustedWithoutRetryInfo(t *testing.T) {
15801591
},
15811592
)
15821593

1594+
md := metadata.Pairs(resourcePrefixHeader, "projects/p/instances/i/databases/d")
1595+
ctx := metadata.NewOutgoingContext(context.Background(), md)
15831596
// The retry is counted from the second call.
15841597
r := -1
15851598
// Establish a stream to mock cloud spanner server.
1586-
iter := stream(context.Background(), nil, c.metricsTracerFactory,
1599+
iter := stream(ctx, nil, c.metricsTracerFactory,
15871600
func(ct context.Context, resumeToken []byte, opts ...gax.CallOption) (streamingReceiver, error) {
15881601
r++
15891602
return mc.ExecuteStreamingSql(ct, &sppb.ExecuteSqlRequest{
@@ -1644,10 +1657,12 @@ func TestRetryResourceExhaustedWithRetryInfo(t *testing.T) {
16441657
},
16451658
)
16461659

1660+
md := metadata.Pairs(resourcePrefixHeader, "projects/p/instances/i/databases/d")
1661+
ctx := metadata.NewOutgoingContext(context.Background(), md)
16471662
// The retry is counted from the second call.
16481663
r := -1
16491664
// Establish a stream to mock cloud spanner server.
1650-
iter := stream(context.Background(), nil, c.metricsTracerFactory,
1665+
iter := stream(ctx, nil, c.metricsTracerFactory,
16511666
func(ct context.Context, resumeToken []byte, opts ...gax.CallOption) (streamingReceiver, error) {
16521667
r++
16531668
return mc.ExecuteStreamingSql(ct, &sppb.ExecuteSqlRequest{
@@ -1696,8 +1711,9 @@ func TestCancelTimeout(t *testing.T) {
16961711
}
16971712
done := make(chan int)
16981713

1714+
md := metadata.Pairs(resourcePrefixHeader, "projects/p/instances/i/databases/d")
16991715
// Test cancelling query.
1700-
ctx, cancel := context.WithCancel(context.Background())
1716+
ctx, cancel := context.WithCancel(metadata.NewOutgoingContext(context.Background(), md))
17011717
go func() {
17021718
// Establish a stream to mock cloud spanner server.
17031719
iter := stream(ctx, nil, c.metricsTracerFactory,
@@ -1733,7 +1749,7 @@ func TestCancelTimeout(t *testing.T) {
17331749
}
17341750

17351751
// Test query timeout.
1736-
ctx, cancel = context.WithTimeout(context.Background(), 100*time.Millisecond)
1752+
ctx, cancel = context.WithTimeout(metadata.NewOutgoingContext(context.Background(), md), 100*time.Millisecond)
17371753
defer cancel()
17381754
go func() {
17391755
// Establish a stream to mock cloud spanner server.
@@ -1886,7 +1902,9 @@ func TestRowIteratorDo(t *testing.T) {
18861902
}
18871903

18881904
nRows := 0
1889-
iter := stream(context.Background(), nil, c.metricsTracerFactory,
1905+
md := metadata.Pairs(resourcePrefixHeader, "projects/p/instances/i/databases/d")
1906+
ctx := metadata.NewOutgoingContext(context.Background(), md)
1907+
iter := stream(ctx, nil, c.metricsTracerFactory,
18901908
func(ct context.Context, resumeToken []byte, opts ...gax.CallOption) (streamingReceiver, error) {
18911909
return mc.ExecuteStreamingSql(ct, &sppb.ExecuteSqlRequest{
18921910
Session: session.Name,
@@ -1921,7 +1939,9 @@ func TestRowIteratorDoWithError(t *testing.T) {
19211939
t.Fatalf("failed to create a session")
19221940
}
19231941

1924-
iter := stream(context.Background(), nil, c.metricsTracerFactory,
1942+
md := metadata.Pairs(resourcePrefixHeader, "projects/p/instances/i/databases/d")
1943+
ctx := metadata.NewOutgoingContext(context.Background(), md)
1944+
iter := stream(ctx, nil, c.metricsTracerFactory,
19251945
func(ct context.Context, resumeToken []byte, opts ...gax.CallOption) (streamingReceiver, error) {
19261946
return mc.ExecuteStreamingSql(ct, &sppb.ExecuteSqlRequest{
19271947
Session: session.Name,
@@ -1955,6 +1975,8 @@ func TestIteratorStopEarly(t *testing.T) {
19551975
t.Fatalf("failed to create a session")
19561976
}
19571977

1978+
md := metadata.Pairs(resourcePrefixHeader, "projects/p/instances/i/databases/d")
1979+
ctx = metadata.NewOutgoingContext(ctx, md)
19581980
iter := stream(ctx, nil, c.metricsTracerFactory,
19591981
func(ct context.Context, resumeToken []byte, opts ...gax.CallOption) (streamingReceiver, error) {
19601982
return mc.ExecuteStreamingSql(ct, &sppb.ExecuteSqlRequest{
@@ -1996,5 +2018,8 @@ func createSession(client spannerClient) (*sppb.Session, error) {
19962018
Database: formattedDatabase,
19972019
Session: &sppb.Session{},
19982020
}
1999-
return client.CreateSession(context.Background(), request)
2021+
ctx := context.Background()
2022+
md := metadata.Pairs(resourcePrefixHeader, "projects/p/instances/i/databases/d")
2023+
ctx = metadata.NewOutgoingContext(ctx, md)
2024+
return client.CreateSession(ctx, request)
20002025
}

‎spanner/session_test.go‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@ import (
3636
"google.golang.org/api/iterator"
3737
"google.golang.org/genproto/googleapis/rpc/errdetails"
3838
"google.golang.org/grpc/codes"
39+
"google.golang.org/grpc/metadata"
3940
"google.golang.org/grpc/status"
4041
)
4142

@@ -129,8 +130,10 @@ loop:
129130
}
130131
}
131132
}
133+
md := metadata.Pairs(resourcePrefixHeader, "projects/p/instances/i/databases/d")
134+
ctx = metadata.NewOutgoingContext(ctx, md)
132135
for _, sh := range shs {
133-
if _, err := sh.getClient().GetSession(context.Background(), &sppb.GetSessionRequest{
136+
if _, err := sh.getClient().GetSession(ctx, &sppb.GetSessionRequest{
134137
Name: sh.getID(),
135138
}); err != nil {
136139
t.Fatalf("error getting expected session from server: %v", err)

‎spanner/transaction_test.go‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1690,7 +1690,8 @@ func TestReadWriteStmtBasedTransaction_UsesMultiplexedSession(t *testing.T) {
16901690
enableMultiplexedSessionForRW: true,
16911691
}
16921692
server, client, teardown := setupMockedTestServerWithConfig(t, ClientConfig{
1693-
SessionPoolConfig: cfg,
1693+
SessionPoolConfig: cfg,
1694+
DisableNativeMetrics: true,
16941695
})
16951696
defer teardown()
16961697
server.TestSpanner.PutExecutionTime(MethodCommitTransaction,
@@ -1741,7 +1742,8 @@ func TestReadWriteStmtBasedTransaction_UsesPreviousTransactionIDForMultiplexedSe
17411742
enableMultiplexedSessionForRW: true,
17421743
}
17431744
server, client, teardown := setupMockedTestServerWithConfig(t, ClientConfig{
1744-
SessionPoolConfig: cfg,
1745+
SessionPoolConfig: cfg,
1746+
DisableNativeMetrics: true,
17451747
})
17461748
defer teardown()
17471749
server.TestSpanner.PutExecutionTime(MethodCommitTransaction,
@@ -1804,7 +1806,8 @@ func TestReadWriteStmtBasedTransaction_SetsPrecommitToken(t *testing.T) {
18041806
enableMultiplexedSessionForRW: true,
18051807
}
18061808
server, client, teardown := setupMockedTestServerWithConfig(t, ClientConfig{
1807-
SessionPoolConfig: cfg,
1809+
SessionPoolConfig: cfg,
1810+
DisableNativeMetrics: true,
18081811
})
18091812
defer teardown()
18101813

0 commit comments

Comments
 (0)