spanner/request-id: complete TODOs and tests

spanner/batch.go

@@ -209,9 +209,13 @@ func (t *BatchReadOnlyTransaction) partitionQuery(ctx context.Context, statement
 		ParamTypes:       paramTypes,
 	}
 	sh.updateLastUseTime()
-	// TODO: (@odeke-em) retrieve the requestID and increment the RPC number
-	// then send it along in every call per retry.
+
+	// PartitionQuery does not retry automatically so we don't need to retrieve
+	// the injected requestID to increment the RPC number on retries.
 	resp, err := client.PartitionQuery(contextWithOutgoingMetadata(ctx, sh.getMetadata(), t.disableRouteToLeader), req, gax.WithGRPCOptions(grpc.Header(&md)))
+	if gcl, ok := client.(*grpcSpannerClient); ok {
+		gcl.setOrResetRPCID()
+	}
 
 	if getGFELatencyMetricsFlag() && md != nil && t.ct != nil {
 		if err := createContextAndCaptureGFELatencyMetrics(ctx, t.ct, md, "partitionQuery"); err != nil {

spanner/errors.go

@@ -123,6 +123,10 @@ func (e *Error) decorate(info string) {
 // APIError error having given error code as its status.
 func spannerErrorf(code codes.Code, format string, args ...interface{}) error {
 	msg := fmt.Sprintf(format, args...)
+	return spannerError(code, msg)
+}
+
+func spannerError(code codes.Code, msg string) error {
 	wrapped, _ := apierror.FromError(status.Error(code, msg))
 	return &Error{
 		Code: code,

spanner/internal/testutil/inmem_spanner_server.go

@@ -90,6 +90,7 @@ const (
 	MethodExecuteBatchDml     string = "EXECUTE_BATCH_DML"
 	MethodStreamingRead       string = "EXECUTE_STREAMING_READ"
 	MethodBatchWrite          string = "BATCH_WRITE"
+	MethodPartitionQuery      string = "PARTITION_QUERY"
 )
 
 // StatementResult represents a mocked result on the test server. The result is
@@ -1107,6 +1108,9 @@ func (s *inMemSpannerServer) Rollback(ctx context.Context, req *spannerpb.Rollba
 }
 
 func (s *inMemSpannerServer) PartitionQuery(ctx context.Context, req *spannerpb.PartitionQueryRequest) (*spannerpb.PartitionResponse, error) {
+	if err := s.simulateExecutionTime(MethodPartitionQuery, req); err != nil {
+		return nil, err
+	}
 	s.mu.Lock()
 	if s.stopped {
 		s.mu.Unlock()

spanner/pdml.go

@@ -110,13 +110,19 @@ func executePdml(ctx context.Context, sh *sessionHandle, req *sppb.ExecuteSqlReq
 	var md metadata.MD
 	sh.updateLastUseTime()
 	// Begin transaction.
-	res, err := sh.getClient().BeginTransaction(ctx, &sppb.BeginTransactionRequest{
+	client := sh.getClient()
+	res, err := client.BeginTransaction(ctx, &sppb.BeginTransactionRequest{
 		Session: sh.getID(),
 		Options: &sppb.TransactionOptions{
 			Mode:                        &sppb.TransactionOptions_PartitionedDml_{PartitionedDml: &sppb.TransactionOptions_PartitionedDml{}},
 			ExcludeTxnFromChangeStreams: options.ExcludeTxnFromChangeStreams,
 		},
 	})
+	// This function is invoked afresh on every retry and it retrieves a fresh client
+	// each time hence does not need an extraction and increment of the injected spanner requestId.
+	if gcl, ok := client.(*grpcSpannerClient); ok {
+		defer gcl.setOrResetRPCID()
+	}
 	if err != nil {
 		return 0, ToSpannerError(err)
 	}
@@ -126,9 +132,8 @@ func executePdml(ctx context.Context, sh *sessionHandle, req *sppb.ExecuteSqlReq
 	}
 
 	sh.updateLastUseTime()
-	// TODO: (@odeke-em) retrieve the requestID and increment the RPC number
-	// then send it along in every call per retry.
-	resultSet, err := sh.getClient().ExecuteSql(ctx, req, gax.WithGRPCOptions(grpc.Header(&md)))
+
+	resultSet, err := client.ExecuteSql(ctx, req, gax.WithGRPCOptions(grpc.Header(&md)))
 	if getGFELatencyMetricsFlag() && md != nil && sh.session.pool != nil {
 		err := captureGFELatencyStats(tag.NewContext(ctx, sh.session.pool.tagMap), md, "executePdml_ExecuteSql")
 		if err != nil {