1// Copyright 2018 Google LLC
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15package spanner
16
17import (
18	"bytes"
19	"context"
20	"io/ioutil"
21	"log"
22	"os"
23	"testing"
24	"time"
25
26	. "cloud.google.com/go/spanner/internal/testutil"
27	sppb "google.golang.org/genproto/googleapis/spanner/v1"
28	"google.golang.org/grpc/codes"
29	"google.golang.org/grpc/status"
30)
31
32func TestMockPartitionedUpdate(t *testing.T) {
33	t.Parallel()
34	ctx := context.Background()
35	_, client, teardown := setupMockedTestServer(t)
36	defer teardown()
37
38	stmt := NewStatement(UpdateBarSetFoo)
39	rowCount, err := client.PartitionedUpdate(ctx, stmt)
40	if err != nil {
41		t.Fatal(err)
42	}
43	want := int64(UpdateBarSetFooRowCount)
44	if rowCount != want {
45		t.Errorf("got %d, want %d", rowCount, want)
46	}
47}
48
49func TestMockPartitionedUpdateWithQuery(t *testing.T) {
50	t.Parallel()
51	ctx := context.Background()
52	_, client, teardown := setupMockedTestServer(t)
53	defer teardown()
54
55	stmt := NewStatement(SelectFooFromBar)
56	_, err := client.PartitionedUpdate(ctx, stmt)
57	wantCode := codes.InvalidArgument
58	var serr *Error
59	if !errorAs(err, &serr) {
60		t.Errorf("got error %v, want spanner.Error", err)
61	}
62	if ErrCode(serr) != wantCode {
63		t.Errorf("got error %v, want code %s", serr, wantCode)
64	}
65}
66
67// PDML should be retried if the transaction is aborted.
68func TestPartitionedUpdate_Aborted(t *testing.T) {
69	t.Parallel()
70	ctx := context.Background()
71	server, client, teardown := setupMockedTestServer(t)
72	defer teardown()
73
74	server.TestSpanner.PutExecutionTime(MethodExecuteSql,
75		SimulatedExecutionTime{
76			Errors: []error{
77				status.Error(codes.Aborted, "Transaction aborted"),
78				status.Error(codes.Internal, "Received unexpected EOS on DATA frame from server"),
79			},
80		})
81	stmt := NewStatement(UpdateBarSetFoo)
82	rowCount, err := client.PartitionedUpdate(ctx, stmt)
83	if err != nil {
84		t.Fatal(err)
85	}
86	want := int64(UpdateBarSetFooRowCount)
87	if rowCount != want {
88		t.Errorf("Row count mismatch\ngot: %d\nwant: %d", rowCount, want)
89	}
90
91	gotReqs, err := shouldHaveReceived(server.TestSpanner, []interface{}{
92		&sppb.BatchCreateSessionsRequest{},
93		&sppb.BeginTransactionRequest{},
94		&sppb.ExecuteSqlRequest{},
95		&sppb.BeginTransactionRequest{},
96		&sppb.ExecuteSqlRequest{},
97		&sppb.BeginTransactionRequest{},
98		&sppb.ExecuteSqlRequest{},
99	})
100	if err != nil {
101		t.Fatal(err)
102	}
103	id1 := gotReqs[2].(*sppb.ExecuteSqlRequest).Transaction.GetId()
104	id2 := gotReqs[4].(*sppb.ExecuteSqlRequest).Transaction.GetId()
105	if bytes.Equal(id1, id2) {
106		t.Errorf("same transaction used twice, expected two different transactions\ngot tx1: %q\ngot tx2: %q", id1, id2)
107	}
108}
109
110// Test that a deadline is respected by PDML, and that the session that was
111// created is also deleted, even though the update timed out.
112func TestPartitionedUpdate_WithDeadline(t *testing.T) {
113	t.Parallel()
114	logger := log.New(os.Stderr, "", log.LstdFlags)
115	server, client, teardown := setupMockedTestServerWithConfig(t, ClientConfig{
116		SessionPoolConfig: DefaultSessionPoolConfig,
117		logger:            logger,
118	})
119	defer teardown()
120
121	ctx := context.Background()
122	ctx, cancel := context.WithDeadline(ctx, time.Now().Add(50*time.Millisecond))
123	defer cancel()
124	server.TestSpanner.PutExecutionTime(MethodExecuteSql,
125		SimulatedExecutionTime{
126			MinimumExecutionTime: 100 * time.Millisecond,
127		})
128	stmt := NewStatement(UpdateBarSetFoo)
129	// The following update will cause a 'Failed to delete session' warning to
130	// be logged. This is expected. To avoid spamming the log, we temporarily
131	// set the output to be discarded.
132	logger.SetOutput(ioutil.Discard)
133	_, err := client.PartitionedUpdate(ctx, stmt)
134	logger.SetOutput(os.Stderr)
135	if err == nil {
136		t.Fatalf("missing expected error")
137	}
138	wantCode := codes.DeadlineExceeded
139	if status.Code(err) != wantCode {
140		t.Fatalf("got error %v, want code %s", err, wantCode)
141	}
142}
143
144func TestPartitionedUpdate_QueryOptions(t *testing.T) {
145	for _, tt := range queryOptionsTestCases() {
146		t.Run(tt.name, func(t *testing.T) {
147			if tt.env.Options != nil {
148				os.Setenv("SPANNER_OPTIMIZER_VERSION", tt.env.Options.OptimizerVersion)
149				defer os.Setenv("SPANNER_OPTIMIZER_VERSION", "")
150			}
151
152			ctx := context.Background()
153			server, client, teardown := setupMockedTestServerWithConfig(t, ClientConfig{QueryOptions: tt.client})
154			defer teardown()
155
156			var err error
157			if tt.query.Options == nil {
158				_, err = client.PartitionedUpdate(ctx, NewStatement(UpdateBarSetFoo))
159			} else {
160				_, err = client.PartitionedUpdateWithOptions(ctx, NewStatement(UpdateBarSetFoo), tt.query)
161			}
162			if err != nil {
163				t.Fatalf("expect no errors, but got %v", err)
164			}
165			checkReqsForQueryOptions(t, server.TestSpanner, tt.want)
166		})
167	}
168}
169