|
| 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 outputstream |
| 16 | + |
| 17 | +import ( |
| 18 | + "log" |
| 19 | + |
| 20 | + "cloud.google.com/go/spanner" |
| 21 | + "cloud.google.com/go/spanner/apiv1/spannerpb" |
| 22 | + "cloud.google.com/go/spanner/test/cloudexecutor/executor/internal/utility" |
| 23 | + executorpb "cloud.google.com/go/spanner/test/cloudexecutor/proto" |
| 24 | + spb "google.golang.org/genproto/googleapis/rpc/status" |
| 25 | + "google.golang.org/grpc/codes" |
| 26 | + "google.golang.org/grpc/status" |
| 27 | + "google.golang.org/protobuf/types/known/timestamppb" |
| 28 | +) |
| 29 | + |
| 30 | +// if OutcomeSender.rowCount exceed maxRowsPerBatch value, we should send rows back to the client in batch. |
| 31 | +const maxRowsPerBatch = 100 |
| 32 | + |
| 33 | +// OutcomeSender is a utility class used for sending action outcomes back to the client. For read |
| 34 | +// actions, it buffers rows and sends partial read results in batches. |
| 35 | +type OutcomeSender struct { |
| 36 | + actionID int32 |
| 37 | + stream executorpb.SpannerExecutorProxy_ExecuteActionAsyncServer |
| 38 | + |
| 39 | + // partialOutcome accumulates rows and other relevant information |
| 40 | + partialOutcome *executorpb.SpannerActionOutcome |
| 41 | + readResult *executorpb.ReadResult |
| 42 | + queryResult *executorpb.QueryResult |
| 43 | + |
| 44 | + // All the relevant variables below should be set before first outcome is sent back, |
| 45 | + // and unused variables should leave null. |
| 46 | + timestamp *timestamppb.Timestamp |
| 47 | + hasReadResult bool |
| 48 | + hasQueryResult bool |
| 49 | + hasChangeStreamRecords bool |
| 50 | + table string // name of the table being read |
| 51 | + index *string // name of the secondary index used for read |
| 52 | + requestIndex *int32 // request index (for multireads) |
| 53 | + rowType *spannerpb.StructType |
| 54 | + |
| 55 | + // Current row count in read/query result |
| 56 | + rowCount int64 |
| 57 | + // modified row count in dml result |
| 58 | + rowsModified []int64 |
| 59 | +} |
| 60 | + |
| 61 | +// NewOutcomeSender returns an OutcomeSender with default fields set. |
| 62 | +func NewOutcomeSender(actionID int32, stream executorpb.SpannerExecutorProxy_ExecuteActionAsyncServer) *OutcomeSender { |
| 63 | + return &OutcomeSender{ |
| 64 | + actionID: actionID, |
| 65 | + stream: stream, |
| 66 | + hasReadResult: false, |
| 67 | + hasQueryResult: false, |
| 68 | + } |
| 69 | +} |
| 70 | + |
| 71 | +// SetTimestamp sets the timestamp for commit. |
| 72 | +func (s *OutcomeSender) SetTimestamp(timestamp *timestamppb.Timestamp) { |
| 73 | + s.timestamp = timestamp |
| 74 | +} |
| 75 | + |
| 76 | +// SetRowType sets the rowType for appending row. |
| 77 | +func (s *OutcomeSender) SetRowType(rowType *spannerpb.StructType) { |
| 78 | + s.rowType = rowType |
| 79 | +} |
| 80 | + |
| 81 | +// InitForRead init the sender for read action, then set the table and index if there exists. |
| 82 | +func (s *OutcomeSender) InitForRead(table string, index *string) { |
| 83 | + s.hasReadResult = true |
| 84 | + s.table = table |
| 85 | + if index != nil { |
| 86 | + s.index = index |
| 87 | + } |
| 88 | +} |
| 89 | + |
| 90 | +// InitForQuery init the sender for query action |
| 91 | +func (s *OutcomeSender) InitForQuery() { |
| 92 | + s.hasQueryResult = true |
| 93 | +} |
| 94 | + |
| 95 | +// InitForBatchRead init the sender for batch read action, then set the table and index if there exists. |
| 96 | +func (s *OutcomeSender) InitForBatchRead(table string, index *string) { |
| 97 | + s.InitForRead(table, index) |
| 98 | + // Cloud API supports only simple batch reads (not multi reads), so request index is always 0. |
| 99 | + requestIndex := int32(0) |
| 100 | + s.requestIndex = &requestIndex |
| 101 | +} |
| 102 | + |
| 103 | +// AppendDmlRowsModified add rows modified in dml to result |
| 104 | +func (s *OutcomeSender) AppendDmlRowsModified(rowsModified int64) { |
| 105 | + s.rowsModified = append(s.rowsModified, rowsModified) |
| 106 | +} |
| 107 | + |
| 108 | +// FinishSuccessfully sends the last outcome with OK status. |
| 109 | +func (s *OutcomeSender) FinishSuccessfully() error { |
| 110 | + s.buildOutcome() |
| 111 | + s.partialOutcome.Status = &spb.Status{Code: int32(codes.OK)} |
| 112 | + return s.flush() |
| 113 | +} |
| 114 | + |
| 115 | +// FinishWithTransactionRestarted sends the last outcome with aborted error, |
| 116 | +// this will set the TransactionRestarted to true |
| 117 | +func (s *OutcomeSender) FinishWithTransactionRestarted() error { |
| 118 | + s.buildOutcome() |
| 119 | + transactionRestarted := true |
| 120 | + s.partialOutcome.TransactionRestarted = &transactionRestarted |
| 121 | + s.partialOutcome.Status = &spb.Status{Code: int32(codes.OK)} |
| 122 | + return s.flush() |
| 123 | +} |
| 124 | + |
| 125 | +// FinishWithError sends the last outcome with given error status. |
| 126 | +func (s *OutcomeSender) FinishWithError(err error) error { |
| 127 | + s.buildOutcome() |
| 128 | + //TODO(harsha:oct10) uncomment below line and comment s.partialOutcome.Status = errToStatus(err) |
| 129 | + //s.partialOutcome.Status = &status.Status{Code: int32(gstatus.Code(err)), Message: err.Error()} |
| 130 | + s.partialOutcome.Status = utility.ErrToStatus(err) |
| 131 | + return s.flush() |
| 132 | +} |
| 133 | + |
| 134 | +// AppendRow adds another row to buffer. If buffer hits its size limit, the buffered rows will be sent back. |
| 135 | +func (s *OutcomeSender) AppendRow(row *executorpb.ValueList) error { |
| 136 | + if !s.hasReadResult && !s.hasQueryResult { |
| 137 | + return spanner.ToSpannerError(status.Error(codes.InvalidArgument, "either hasReadResult or hasQueryResult should be true")) |
| 138 | + } |
| 139 | + if s.rowType == nil { |
| 140 | + return spanner.ToSpannerError(status.Error(codes.InvalidArgument, "rowType should be set first")) |
| 141 | + } |
| 142 | + s.buildOutcome() |
| 143 | + if s.hasReadResult { |
| 144 | + s.readResult.Row = append(s.readResult.Row, row) |
| 145 | + s.rowCount++ |
| 146 | + } else if s.hasQueryResult { |
| 147 | + s.queryResult.Row = append(s.queryResult.Row, row) |
| 148 | + s.rowCount++ |
| 149 | + } |
| 150 | + if s.rowCount >= maxRowsPerBatch { |
| 151 | + return s.flush() |
| 152 | + } |
| 153 | + return nil |
| 154 | +} |
| 155 | + |
| 156 | +// buildOutcome will build the partialOutcome if not exists using relevant variables. |
| 157 | +func (s *OutcomeSender) buildOutcome() { |
| 158 | + if s.partialOutcome != nil { |
| 159 | + return |
| 160 | + } |
| 161 | + s.partialOutcome = &executorpb.SpannerActionOutcome{ |
| 162 | + CommitTime: s.timestamp, |
| 163 | + } |
| 164 | + if s.hasReadResult { |
| 165 | + s.readResult = &executorpb.ReadResult{ |
| 166 | + Table: s.table, |
| 167 | + Index: s.index, |
| 168 | + RowType: s.rowType, |
| 169 | + RequestIndex: s.requestIndex, |
| 170 | + } |
| 171 | + } else if s.hasQueryResult { |
| 172 | + s.queryResult = &executorpb.QueryResult{ |
| 173 | + RowType: s.rowType, |
| 174 | + } |
| 175 | + } |
| 176 | +} |
| 177 | + |
| 178 | +// flush sends partialOutcome to stream and clear the internal state |
| 179 | +func (s *OutcomeSender) flush() error { |
| 180 | + if s == nil || s.partialOutcome == nil { |
| 181 | + log.Println("outcomeSender.flush() is called when there is no partial outcome to send. This is an internal error that should never happen") |
| 182 | + return spanner.ToSpannerError(status.Error(codes.InvalidArgument, "either outcome sender or partial outcome is nil")) |
| 183 | + } |
| 184 | + s.partialOutcome.DmlRowsModified = s.rowsModified |
| 185 | + if s.hasReadResult { |
| 186 | + s.partialOutcome.ReadResult = s.readResult |
| 187 | + } else if s.hasQueryResult { |
| 188 | + s.partialOutcome.QueryResult = s.queryResult |
| 189 | + } |
| 190 | + err := s.SendOutcome(s.partialOutcome) |
| 191 | + s.partialOutcome = nil |
| 192 | + s.readResult = nil |
| 193 | + s.queryResult = nil |
| 194 | + s.rowCount = 0 |
| 195 | + s.rowsModified = []int64{} |
| 196 | + return err |
| 197 | +} |
| 198 | + |
| 199 | +// SendOutcome sends the given SpannerActionOutcome. |
| 200 | +func (s *OutcomeSender) SendOutcome(outcome *executorpb.SpannerActionOutcome) error { |
| 201 | + log.Printf("sending result %v actionId %d", outcome, s.actionID) |
| 202 | + resp := &executorpb.SpannerAsyncActionResponse{ |
| 203 | + ActionId: s.actionID, |
| 204 | + Outcome: outcome, |
| 205 | + } |
| 206 | + err := s.stream.Send(resp) |
| 207 | + if err != nil { |
| 208 | + log.Printf("Failed to send outcome with error: %s", err.Error()) |
| 209 | + } else { |
| 210 | + log.Printf("Sent result %v actionId %d", outcome, s.actionID) |
| 211 | + } |
| 212 | + return err |
| 213 | +} |
0 commit comments