mirror of
https://github.com/Luzifer/nginx-sso.git
synced 2024-12-21 05:11:17 +00:00
102 lines
3.2 KiB
Go
102 lines
3.2 KiB
Go
// Copyright 2018 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 (
|
|
"context"
|
|
"time"
|
|
|
|
"cloud.google.com/go/internal/trace"
|
|
"google.golang.org/api/iterator"
|
|
sppb "google.golang.org/genproto/googleapis/spanner/v1"
|
|
"google.golang.org/grpc/codes"
|
|
)
|
|
|
|
// PartitionedUpdate executes a DML statement in parallel across the database, using
|
|
// separate, internal transactions that commit independently. The DML statement must
|
|
// be fully partitionable: it must be expressible as the union of many statements
|
|
// each of which accesses only a single row of the table. The statement should also be
|
|
// idempotent, because it may be applied more than once.
|
|
//
|
|
// PartitionedUpdate returns an estimated count of the number of rows affected. The actual
|
|
// number of affected rows may be greater than the estimate.
|
|
func (c *Client) PartitionedUpdate(ctx context.Context, statement Statement) (count int64, err error) {
|
|
ctx = trace.StartSpan(ctx, "cloud.google.com/go/spanner.PartitionedUpdate")
|
|
defer func() { trace.EndSpan(ctx, err) }()
|
|
if err := checkNestedTxn(ctx); err != nil {
|
|
return 0, err
|
|
}
|
|
|
|
var (
|
|
tx transactionID
|
|
s *session
|
|
sh *sessionHandle
|
|
)
|
|
// create session
|
|
sc := c.rrNext()
|
|
s, err = createSession(ctx, sc, c.database, c.sessionLabels, c.md)
|
|
if err != nil {
|
|
return 0, toSpannerError(err)
|
|
}
|
|
defer s.delete(ctx)
|
|
sh = &sessionHandle{session: s}
|
|
// begin transaction
|
|
err = runRetryable(contextWithOutgoingMetadata(ctx, sh.getMetadata()), func(ctx context.Context) error {
|
|
res, e := sc.BeginTransaction(ctx, &sppb.BeginTransactionRequest{
|
|
Session: sh.getID(),
|
|
Options: &sppb.TransactionOptions{
|
|
Mode: &sppb.TransactionOptions_PartitionedDml_{PartitionedDml: &sppb.TransactionOptions_PartitionedDml{}},
|
|
},
|
|
})
|
|
if e != nil {
|
|
return e
|
|
}
|
|
tx = res.Id
|
|
return nil
|
|
})
|
|
if err != nil {
|
|
return 0, toSpannerError(err)
|
|
}
|
|
req := &sppb.ExecuteSqlRequest{
|
|
Session: sh.getID(),
|
|
Transaction: &sppb.TransactionSelector{
|
|
Selector: &sppb.TransactionSelector_Id{Id: tx},
|
|
},
|
|
Sql: statement.SQL,
|
|
}
|
|
rpc := func(ctx context.Context, resumeToken []byte) (streamingReceiver, error) {
|
|
req.ResumeToken = resumeToken
|
|
return sc.ExecuteStreamingSql(ctx, req)
|
|
}
|
|
iter := stream(contextWithOutgoingMetadata(ctx, sh.getMetadata()),
|
|
rpc, func(time.Time) {}, func(error) {})
|
|
// TODO(jba): factor out the following code from here and ReadWriteTransaction.Update.
|
|
defer iter.Stop()
|
|
for {
|
|
_, err := iter.Next()
|
|
if err == iterator.Done {
|
|
break
|
|
}
|
|
if err != nil {
|
|
return 0, toSpannerError(err)
|
|
}
|
|
time.Sleep(time.Second)
|
|
}
|
|
|
|
if !iter.sawStats {
|
|
return 0, spannerErrorf(codes.InvalidArgument, "query passed to Update: %q", statement.SQL)
|
|
}
|
|
return iter.RowCount, nil
|
|
}
|