diff --git a/spanner/pdml.go b/spanner/pdml.go index eac3b99c5660..e50331235ab3 100644 --- a/spanner/pdml.go +++ b/spanner/pdml.go @@ -48,20 +48,15 @@ func (c *Client) partitionedUpdate(ctx context.Context, statement Statement, opt if err := checkNestedTxn(ctx); err != nil { return 0, err } - var ( - s *session - sh *sessionHandle - ) - // Create session. - s, err = c.sc.createSession(ctx) + + sh, err := c.idleSessions.take(ctx) if err != nil { return 0, toSpannerError(err) } - // Delete the session at the end of the request. If the PDML statement - // timed out or was cancelled, the DeleteSession request might not succeed, - // but the session will eventually be garbage collected by the server. - defer s.delete(ctx) - sh = &sessionHandle{session: s} + if sh != nil { + defer sh.recycle() + } + // Create the parameters and the SQL request, but without a transaction. // The transaction reference will be added by the executePdml method. params, paramTypes, err := statement.convertParams() diff --git a/spanner/pdml_test.go b/spanner/pdml_test.go index 430ad674e89e..bbbc2df31878 100644 --- a/spanner/pdml_test.go +++ b/spanner/pdml_test.go @@ -89,14 +89,13 @@ func TestPartitionedUpdate_Aborted(t *testing.T) { } gotReqs, err := shouldHaveReceived(server.TestSpanner, []interface{}{ - &sppb.CreateSessionRequest{}, + &sppb.BatchCreateSessionsRequest{}, &sppb.BeginTransactionRequest{}, &sppb.ExecuteSqlRequest{}, &sppb.BeginTransactionRequest{}, &sppb.ExecuteSqlRequest{}, &sppb.BeginTransactionRequest{}, &sppb.ExecuteSqlRequest{}, - &sppb.DeleteSessionRequest{}, }) if err != nil { t.Fatal(err)