Skip to content

Commit ca76671

Browse files
harshachintaharshachinta
authored
feat(spanner): Add DML, DQL, Mutation, Txn Actions and Utility methods for executor framework (#8976)
* feat(spanner): add code for executor framework * feat(spanner): lint fix * feat(spanner): check test flakiness * feat(spanner): revert - check test flakiness * feat(spanner): update context * feat(spanner): add start and finish transaction actions * feat(spanner): remove the fixed TODO * feat(spanner): add input stream handling * feat(spanner): code fixes
1 parent 1a16cbf commit ca76671

File tree

12 files changed

+1690
-46
lines changed

12 files changed

+1690
-46
lines changed

‎spanner/test/cloudexecutor/executor/actions/dml.go

+64
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
1+
// Copyright 2023 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+
15+
package actions
16+
17+
import (
18+
"context"
19+
"log"
20+
21+
"cloud.google.com/go/spanner"
22+
"cloud.google.com/go/spanner/test/cloudexecutor/executor/internal/outputstream"
23+
"cloud.google.com/go/spanner/test/cloudexecutor/executor/internal/utility"
24+
executorpb "cloud.google.com/go/spanner/test/cloudexecutor/proto"
25+
"google.golang.org/grpc/codes"
26+
"google.golang.org/grpc/status"
27+
)
28+
29+
// DmlActionHandler holds the necessary components required for DML action.
30+
type DmlActionHandler struct {
31+
Action *executorpb.DmlAction
32+
FlowContext *ExecutionFlowContext
33+
OutcomeSender *outputstream.OutcomeSender
34+
}
35+
36+
// ExecuteAction executes a DML update action request, store the results in the outputstream.OutcomeSender.
37+
func (h *DmlActionHandler) ExecuteAction(ctx context.Context) error {
38+
h.FlowContext.mu.Lock()
39+
defer h.FlowContext.mu.Unlock()
40+
log.Printf("Executing dml update %s\n %v\n", h.FlowContext.transactionSeed, h.Action)
41+
stmt, err := utility.BuildQuery(h.Action.GetUpdate())
42+
if err != nil {
43+
return h.OutcomeSender.FinishWithError(err)
44+
}
45+
46+
var iter *spanner.RowIterator
47+
txn, err := h.FlowContext.getTransactionForWrite()
48+
if err != nil {
49+
return h.OutcomeSender.FinishWithError(err)
50+
}
51+
h.OutcomeSender.InitForQuery()
52+
iter = txn.Query(ctx, stmt)
53+
defer iter.Stop()
54+
log.Printf("Parsing DML result %s\n", h.FlowContext.transactionSeed)
55+
err = processResults(iter, 0, h.OutcomeSender, h.FlowContext)
56+
if err != nil {
57+
if status.Code(err) == codes.Aborted {
58+
return h.OutcomeSender.FinishWithTransactionRestarted()
59+
}
60+
return h.OutcomeSender.FinishWithError(err)
61+
}
62+
h.OutcomeSender.AppendDmlRowsModified(iter.RowCount)
63+
return h.OutcomeSender.FinishSuccessfully()
64+
}

‎spanner/test/cloudexecutor/executor/actions/dql.go

+166
Original file line numberDiff line numberDiff line change
@@ -15,13 +15,179 @@
1515
package actions
1616

1717
import (
18+
"context"
19+
"fmt"
20+
"log"
21+
1822
"cloud.google.com/go/spanner"
1923
"cloud.google.com/go/spanner/apiv1/spannerpb"
2024
"cloud.google.com/go/spanner/test/cloudexecutor/executor/internal/outputstream"
2125
"cloud.google.com/go/spanner/test/cloudexecutor/executor/internal/utility"
26+
executorpb "cloud.google.com/go/spanner/test/cloudexecutor/proto"
27+
"google.golang.org/api/iterator"
28+
"google.golang.org/grpc/codes"
29+
"google.golang.org/grpc/status"
2230
)
2331

32+
// ReadActionHandler holds the necessary components required for executorpb.ReadAction.
33+
type ReadActionHandler struct {
34+
Action *executorpb.ReadAction
35+
FlowContext *ExecutionFlowContext
36+
OutcomeSender *outputstream.OutcomeSender
37+
}
38+
39+
// ExecuteAction executes a read action request, store the results in the OutcomeSender.
40+
func (h *ReadActionHandler) ExecuteAction(ctx context.Context) error {
41+
h.FlowContext.mu.Lock()
42+
defer h.FlowContext.mu.Unlock()
43+
log.Printf("Executing read %s:\n %v", h.FlowContext.transactionSeed, h.Action)
44+
action := h.Action
45+
var err error
46+
47+
var typeList []*spannerpb.Type
48+
if action.Index != nil {
49+
typeList, err = extractTypes(action.GetTable(), action.GetColumn(), h.FlowContext.tableMetadata)
50+
if err != nil {
51+
return h.OutcomeSender.FinishWithError(status.Error(codes.InvalidArgument, fmt.Sprintf("Can't extract types from metadata: %s", err)))
52+
}
53+
} else {
54+
typeList, err = h.FlowContext.tableMetadata.GetKeyColumnTypes(action.GetTable())
55+
if err != nil {
56+
return h.OutcomeSender.FinishWithError(status.Error(codes.InvalidArgument, fmt.Sprintf("Can't extract types from metadata: %s", err)))
57+
}
58+
}
59+
60+
keySet, err := utility.KeySetProtoToCloudKeySet(action.GetKeys(), typeList)
61+
if err != nil {
62+
return h.OutcomeSender.FinishWithError(status.Error(codes.InvalidArgument, fmt.Sprintf("Can't convert rowSet: %s", err)))
63+
}
64+
65+
var iter *spanner.RowIterator
66+
if h.FlowContext.currentActiveTransaction == None {
67+
return h.OutcomeSender.FinishWithError(status.Error(codes.InvalidArgument, "no active transaction"))
68+
} else if h.FlowContext.currentActiveTransaction == Batch {
69+
return h.OutcomeSender.FinishWithError(status.Error(codes.InvalidArgument, "can't execute regular read in a batch transaction"))
70+
} else if h.FlowContext.currentActiveTransaction == Read {
71+
txn, err := h.FlowContext.getTransactionForRead()
72+
if err != nil {
73+
return h.OutcomeSender.FinishWithError(err)
74+
}
75+
h.OutcomeSender.InitForRead(action.GetTable(), action.Index)
76+
h.FlowContext.numPendingReads++
77+
if action.Index != nil {
78+
iter = txn.ReadUsingIndex(ctx, action.GetTable(), action.GetIndex(), keySet, action.GetColumn())
79+
} else {
80+
iter = txn.Read(ctx, action.GetTable(), keySet, action.GetColumn())
81+
}
82+
} else if h.FlowContext.currentActiveTransaction == ReadWrite {
83+
txn, err := h.FlowContext.getTransactionForWrite()
84+
if err != nil {
85+
return h.OutcomeSender.FinishWithError(err)
86+
}
87+
h.OutcomeSender.InitForRead(action.GetTable(), action.Index)
88+
h.FlowContext.numPendingReads++
89+
if action.Index != nil {
90+
iter = txn.ReadUsingIndex(ctx, action.GetTable(), action.GetIndex(), keySet, action.GetColumn())
91+
} else {
92+
iter = txn.Read(ctx, action.GetTable(), keySet, action.GetColumn())
93+
}
94+
}
95+
defer iter.Stop()
96+
log.Printf("Parsing read result %s\n", h.FlowContext.transactionSeed)
97+
err = processResults(iter, int64(action.GetLimit()), h.OutcomeSender, h.FlowContext)
98+
if err != nil {
99+
h.FlowContext.finishRead(status.Code(err))
100+
if status.Code(err) == codes.Aborted {
101+
return h.OutcomeSender.FinishWithTransactionRestarted()
102+
}
103+
return h.OutcomeSender.FinishWithError(err)
104+
}
105+
h.FlowContext.finishRead(codes.OK)
106+
return h.OutcomeSender.FinishSuccessfully()
107+
}
108+
109+
// QueryActionHandler holds the necessary components required for executorpb.QueryAction.
110+
type QueryActionHandler struct {
111+
Action *executorpb.QueryAction
112+
FlowContext *ExecutionFlowContext
113+
OutcomeSender *outputstream.OutcomeSender
114+
}
115+
116+
// ExecuteAction executes a query action request, store the results in the OutcomeSender.
117+
func (h *QueryActionHandler) ExecuteAction(ctx context.Context) error {
118+
h.FlowContext.mu.Lock()
119+
defer h.FlowContext.mu.Unlock()
120+
log.Printf("Executing query %s\n %v\n", h.FlowContext.transactionSeed, h.Action)
121+
stmt, err := utility.BuildQuery(h.Action)
122+
if err != nil {
123+
return h.OutcomeSender.FinishWithError(err)
124+
}
125+
126+
var iter *spanner.RowIterator
127+
if h.FlowContext.currentActiveTransaction == None {
128+
return h.OutcomeSender.FinishWithError(status.Error(codes.InvalidArgument, "no active transaction"))
129+
} else if h.FlowContext.currentActiveTransaction == Batch {
130+
return h.OutcomeSender.FinishWithError(status.Error(codes.InvalidArgument, "can't execute regular read in a batch transaction"))
131+
} else if h.FlowContext.currentActiveTransaction == Read {
132+
txn, err := h.FlowContext.getTransactionForRead()
133+
if err != nil {
134+
return h.OutcomeSender.FinishWithError(err)
135+
}
136+
h.OutcomeSender.InitForQuery()
137+
h.FlowContext.numPendingReads++
138+
iter = txn.Query(ctx, stmt)
139+
} else if h.FlowContext.currentActiveTransaction == ReadWrite {
140+
txn, err := h.FlowContext.getTransactionForWrite()
141+
if err != nil {
142+
return h.OutcomeSender.FinishWithError(err)
143+
}
144+
h.OutcomeSender.InitForQuery()
145+
h.FlowContext.numPendingReads++
146+
iter = txn.Query(ctx, stmt)
147+
}
148+
defer iter.Stop()
149+
log.Printf("Parsing query result %s\n", h.FlowContext.transactionSeed)
150+
err = processResults(iter, 0, h.OutcomeSender, h.FlowContext)
151+
if err != nil {
152+
h.FlowContext.finishRead(status.Code(err))
153+
if status.Code(err) == codes.Aborted {
154+
return h.OutcomeSender.FinishWithTransactionRestarted()
155+
}
156+
return h.OutcomeSender.FinishWithError(err)
157+
}
158+
h.FlowContext.finishRead(codes.OK)
159+
return h.OutcomeSender.FinishSuccessfully()
160+
}
161+
162+
// processResults processes a ResultSet from a read/query/dml and store the results in the OutcomeSender.
24163
func processResults(iter *spanner.RowIterator, limit int64, outcomeSender *outputstream.OutcomeSender, flowContext *ExecutionFlowContext) error {
164+
var rowCount int64 = 0
165+
log.Printf("Iterating result set: %s\n", flowContext.transactionSeed)
166+
for {
167+
row, err := iter.Next()
168+
if err == iterator.Done {
169+
return nil
170+
}
171+
if err != nil {
172+
return err
173+
}
174+
spannerRow, rowType, err := utility.ConvertSpannerRow(row)
175+
if err != nil {
176+
return err
177+
}
178+
outcomeSender.SetRowType(rowType)
179+
// outcomeSender.rowType = rowType
180+
err = outcomeSender.AppendRow(spannerRow)
181+
if err != nil {
182+
return err
183+
}
184+
rowCount++
185+
if limit > 0 && rowCount >= limit {
186+
log.Printf("Stopping at row limit: %d", limit)
187+
break
188+
}
189+
}
190+
log.Printf("Successfully processed result: %s\n", flowContext.transactionSeed)
25191
return nil
26192
}
27193

‎spanner/test/cloudexecutor/executor/actions/execution_flow_context.go

+71
Original file line numberDiff line numberDiff line change
@@ -17,10 +17,13 @@ package actions
1717
import (
1818
"context"
1919
"log"
20+
"strings"
2021
"sync"
22+
"time"
2123

2224
"cloud.google.com/go/spanner"
2325
"cloud.google.com/go/spanner/test/cloudexecutor/executor/internal/utility"
26+
executorpb "cloud.google.com/go/spanner/test/cloudexecutor/proto"
2427
"google.golang.org/grpc/codes"
2528
"google.golang.org/grpc/status"
2629
)
@@ -122,3 +125,71 @@ func (c *ExecutionFlowContext) initReadState() {
122125
c.readAborted = false
123126
c.numPendingReads = 0
124127
}
128+
129+
// clear clears the transaction related variables.
130+
func (c *ExecutionFlowContext) clear() {
131+
c.roTxn = nil
132+
c.rwTxn = nil
133+
c.tableMetadata = nil
134+
}
135+
136+
// finish attempts to finish the transaction by either committing it or exiting without committing.
137+
// In order to follow the ExecuteActions protocol, we must distinguish Spanner-generated errors
138+
// (e.g. RowNotFound) and internal errors (e.g. a precondition is not met). Errors returned from
139+
// Spanner populate the status of SpannerActionOutcome. Internal errors, in contrast, break the
140+
// stubby call. For this reason, finish() has two return values dedicated to errors. If any of
141+
// these errors is not nil, other return values are undefined.
142+
// Return values in order:
143+
// 1. Whether or not the transaction is restarted. It will be true if commit has been attempted,
144+
// but Spanner returned aborted and restarted instead. When that happens, all reads and writes
145+
// should be replayed, followed by another commit attempt.
146+
// 2. Commit timestamp. It's returned only if commit has succeeded.
147+
// 3. Spanner error -- an error that Spanner client returned.
148+
// 4. Internal error.
149+
func (c *ExecutionFlowContext) finish(ctx context.Context, txnFinishMode executorpb.FinishTransactionAction_Mode) (bool, *time.Time, error, error) {
150+
if txnFinishMode == executorpb.FinishTransactionAction_COMMIT {
151+
var err error
152+
ts, err := c.rwTxn.Commit(ctx)
153+
if err != nil {
154+
log.Printf("Transaction finished with error %v", err)
155+
if status.Code(err) == codes.Aborted {
156+
log.Println("Transaction Aborted. Sending status to outcome sender to restart the transaction.")
157+
return true, nil, nil, nil
158+
}
159+
// Filter expected errors
160+
if status.Code(err) == codes.Unknown && strings.Contains(err.Error(), "Transaction outcome unknown") {
161+
return false, nil, spanner.ToSpannerError(status.Error(codes.DeadlineExceeded, "Transaction outcome unknown")), nil
162+
}
163+
// TODO(harsha): check if this is an internal or spanner error
164+
return false, nil, err, nil
165+
}
166+
return false, &ts, nil, nil
167+
} else if txnFinishMode == executorpb.FinishTransactionAction_ABANDON {
168+
log.Printf("Transaction Abandoned")
169+
c.rwTxn.Rollback(ctx)
170+
return false, nil, nil, nil
171+
} else {
172+
return false, nil, nil, spanner.ToSpannerError(status.Errorf(codes.InvalidArgument, "Unsupported finish mode %s", txnFinishMode.String()))
173+
}
174+
}
175+
176+
// CloseOpenTransactions cleans up all the active transactions if the stubby call is closing.
177+
func (c *ExecutionFlowContext) CloseOpenTransactions() {
178+
c.mu.Lock()
179+
defer c.mu.Unlock()
180+
if c.roTxn != nil {
181+
log.Println("A read only transaction was active when stubby call closed.")
182+
c.roTxn.Close()
183+
}
184+
if c.rwTxn != nil {
185+
log.Println("Abandon a read-write transaction that was active when stubby call closed.")
186+
_, _, _, err := c.finish(context.Background(), executorpb.FinishTransactionAction_ABANDON)
187+
if err != nil {
188+
log.Printf("Failed to abandon a read-write transaction: %v\n", err)
189+
}
190+
}
191+
if c.batchTxn != nil {
192+
log.Println("A batch transaction was active when stubby call closed.")
193+
c.batchTxn.Close()
194+
}
195+
}

0 commit comments

Comments
 (0)