Merge pull request #1300 from ydb-platform/operation-status

* Fixed bug with fail cast of grpc response to `operation.{Response,Status}`

Author: asmyasnikov

CHANGELOG.md

@@ -1,3 +1,5 @@
+* Fixed bug with fail cast of grpc response to `operation.{Response,Status}`
+
 ## v3.74.3
 * Removed check the node is available for query and table service sessions 
 * Refactored the `balancers.PreferLocations()` function - it is a clean/pure function

internal/conn/conn.go

@@ -15,7 +15,7 @@ import (
 
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/endpoint"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/meta"
-	"github.com/ydb-platform/ydb-go-sdk/v3/internal/response"
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/operation"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/stack"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xcontext"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
@@ -299,109 +299,151 @@ func (c *conn) Close(ctx context.Context) (err error) {
 	return c.wrapError(err)
 }
 
+var (
+	onTransportErrorStub = func(ctx context.Context, err error) {}
+	wrapErrorStub        = func(err error) error { return err }
+)
+
 //nolint:funlen
-func (c *conn) Invoke(
+func invoke(
 	ctx context.Context,
 	method string,
-	req interface{},
-	res interface{},
+	req, reply any,
+	cc grpc.ClientConnInterface,
+	onTransportError func(context.Context, error),
+	address string,
+	wrapError func(err error) error,
 	opts ...grpc.CallOption,
-) (err error) {
-	var (
-		opID        string
-		issues      []trace.Issue
-		useWrapping = UseWrapping(ctx)
-		onDone      = trace.DriverOnConnInvoke(
-			c.config.Trace(), &ctx,
-			stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/3/internal/conn.(*conn).Invoke"),
-			c.endpoint, trace.Method(method),
-		)
-		cc *grpc.ClientConn
-		md = metadata.MD{}
-	)
-	defer func() {
-		meta.CallTrailerCallback(ctx, md)
-		onDone(err, issues, opID, c.GetState(), md)
-	}()
+) (
+	opID string,
+	issues []trace.Issue,
+	_ error,
+) {
+	useWrapping := UseWrapping(ctx)
 
-	cc, err = c.realConn(ctx)
+	ctx, traceID, err := meta.TraceID(ctx)
 	if err != nil {
-		return c.wrapError(err)
+		return opID, issues, xerrors.WithStackTrace(err)
 	}
 
-	stop := c.lastUsage.Start()
-	defer stop()
+	ctx, sentMark := markContext(meta.WithTraceID(ctx, traceID))
 
-	ctx, traceID, err := meta.TraceID(ctx)
-	if err != nil {
-		return xerrors.WithStackTrace(err)
+	if onTransportError == nil {
+		onTransportError = onTransportErrorStub
 	}
 
-	ctx, sentMark := markContext(meta.WithTraceID(ctx, traceID))
+	if wrapError == nil {
+		wrapError = wrapErrorStub
+	}
 
-	err = cc.Invoke(ctx, method, req, res, append(opts, grpc.Trailer(&md))...)
+	err = cc.Invoke(ctx, method, req, reply, opts...)
 	if err != nil {
 		if xerrors.IsContextError(err) {
-			return xerrors.WithStackTrace(err)
+			return opID, issues, xerrors.WithStackTrace(err)
 		}
 
-		defer func() {
-			c.onTransportError(ctx, err)
-		}()
+		defer onTransportError(ctx, err)
 
 		if useWrapping {
 			err = xerrors.Transport(err,
-				xerrors.WithAddress(c.Address()),
+				xerrors.WithAddress(address),
 				xerrors.WithTraceID(traceID),
 			)
 			if sentMark.canRetry() {
-				return c.wrapError(xerrors.Retryable(err, xerrors.WithName("Invoke")))
+				return opID, issues, wrapError(xerrors.Retryable(err, xerrors.WithName("Invoke")))
 			}
 
-			return c.wrapError(err)
+			return opID, issues, wrapError(err)
 		}
 
-		return err
-	}
-
-	err = c.handleResponse(res, &opID, &issues, traceID, useWrapping)
-	if err != nil {
-		return err
+		return opID, issues, err
 	}
 
-	return err
-}
-
-func (c *conn) handleResponse(
-	res interface{},
-	opID *string,
-	issues *[]trace.Issue,
-	traceID string,
-	useWrapping bool,
-) error {
-	if o, ok := res.(response.Response); ok {
-		*opID = o.GetOperation().GetId()
-		for _, issue := range o.GetOperation().GetIssues() {
-			*issues = append(*issues, issue)
+	switch t := reply.(type) {
+	case operation.Response:
+		opID = t.GetOperation().GetId()
+		for _, issue := range t.GetOperation().GetIssues() {
+			issues = append(issues, issue)
 		}
 		if useWrapping {
 			switch {
-			case !o.GetOperation().GetReady():
-				return c.wrapError(errOperationNotReady)
+			case !t.GetOperation().GetReady():
+				return opID, issues, wrapError(errOperationNotReady)
 
-			case o.GetOperation().GetStatus() != Ydb.StatusIds_SUCCESS:
-				return c.wrapError(
+			case t.GetOperation().GetStatus() != Ydb.StatusIds_SUCCESS:
+				return opID, issues, wrapError(
 					xerrors.Operation(
-						xerrors.FromOperation(o.GetOperation()),
-						xerrors.WithAddress(c.Address()),
+						xerrors.FromOperation(t.GetOperation()),
+						xerrors.WithAddress(address),
+						xerrors.WithTraceID(traceID),
+					),
+				)
+			}
+		}
+	case operation.Status:
+		for _, issue := range t.GetIssues() {
+			issues = append(issues, issue)
+		}
+		if useWrapping {
+			if t.GetStatus() != Ydb.StatusIds_SUCCESS {
+				return opID, issues, wrapError(
+					xerrors.Operation(
+						xerrors.FromOperation(t),
+						xerrors.WithAddress(address),
 						xerrors.WithTraceID(traceID),
 					),
 				)
 			}
 		}
 	}
 
-	return nil
+	return opID, issues, nil
+}
+
+func (c *conn) Invoke(
+	ctx context.Context,
+	method string,
+	req interface{},
+	res interface{},
+	opts ...grpc.CallOption,
+) (err error) {
+	var (
+		opID   string
+		issues []trace.Issue
+		onDone = trace.DriverOnConnInvoke(
+			c.config.Trace(), &ctx,
+			stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/3/internal/conn.(*conn).Invoke"),
+			c.endpoint, trace.Method(method),
+		)
+		cc *grpc.ClientConn
+		md = metadata.MD{}
+	)
+	defer func() {
+		meta.CallTrailerCallback(ctx, md)
+		onDone(err, issues, opID, c.GetState(), md)
+	}()
+
+	cc, err = c.realConn(ctx)
+	if err != nil {
+		return c.wrapError(err)
+	}
+
+	stop := c.lastUsage.Start()
+	defer stop()
+
+	opID, issues, err = invoke(
+		ctx,
+		method,
+		req,
+		res,
+		cc,
+		c.onTransportError,
+		c.Address(),
+		c.wrapError,
+		append(opts, grpc.Trailer(&md))...,
+	)
+
+	return err
 }
 
 //nolint:funlen

internal/conn/conn_test.go

@@ -0,0 +1,152 @@
+package conn
+
+import (
+	"context"
+	"fmt"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	"github.com/ydb-platform/ydb-go-genproto/Ydb_Discovery_V1"
+	"github.com/ydb-platform/ydb-go-genproto/Ydb_Query_V1"
+	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb"
+	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb_Discovery"
+	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb_Operations"
+	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb_Query"
+	"go.uber.org/mock/gomock"
+	"google.golang.org/grpc"
+	grpcCodes "google.golang.org/grpc/codes"
+	grpcStatus "google.golang.org/grpc/status"
+
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xtest"
+)
+
+//go:generate mockgen -destination grpc_client_conn_interface_mock_test.go -package conn -write_package_comment=false google.golang.org/grpc ClientConnInterface
+
+var _ grpc.ClientConnInterface = (*connMock)(nil)
+
+type connMock struct {
+	cc grpc.ClientConnInterface
+}
+
+func (c connMock) Invoke(ctx context.Context, method string, args any, reply any, opts ...grpc.CallOption) error {
+	_, _, err := invoke(ctx, method, args, reply, c.cc, nil, "", nil, opts...)
+
+	return err
+}
+
+func (c connMock) NewStream(
+	ctx context.Context, desc *grpc.StreamDesc, method string, opts ...grpc.CallOption,
+) (grpc.ClientStream, error) {
+	return c.cc.NewStream(ctx, desc, method, opts...)
+}
+
+func TestConn(t *testing.T) {
+	t.Run("Invoke", func(t *testing.T) {
+		t.Run("HappyWay", func(t *testing.T) {
+			ctx := xtest.Context(t)
+			ctrl := gomock.NewController(t)
+			cc := NewMockClientConnInterface(ctrl)
+			cc.EXPECT().Invoke(
+				gomock.Any(),
+				Ydb_Discovery_V1.DiscoveryService_WhoAmI_FullMethodName,
+				&Ydb_Discovery.WhoAmIRequest{},
+				&Ydb_Discovery.WhoAmIResponse{},
+			).DoAndReturn(func(ctx context.Context, method string, args any, reply any, opts ...grpc.CallOption) error {
+				res, ok := reply.(*Ydb_Discovery.WhoAmIResponse)
+				if !ok {
+					return fmt.Errorf("reply is not *Ydb_Discovery.WhoAmIResponse: %T", reply)
+				}
+
+				res.Operation = &Ydb_Operations.Operation{
+					Ready:  true,
+					Status: Ydb.StatusIds_SUCCESS,
+				}
+
+				return nil
+			})
+			client := Ydb_Discovery_V1.NewDiscoveryServiceClient(&connMock{
+				cc,
+			})
+			response, err := client.WhoAmI(ctx, &Ydb_Discovery.WhoAmIRequest{})
+			require.NoError(t, err)
+			require.NotNil(t, response)
+		})
+		t.Run("TransportError", func(t *testing.T) {
+			ctx := xtest.Context(t)
+			ctrl := gomock.NewController(t)
+			cc := NewMockClientConnInterface(ctrl)
+			expectedErr := grpcStatus.Error(grpcCodes.Unavailable, "")
+			cc.EXPECT().Invoke(
+				gomock.Any(),
+				Ydb_Discovery_V1.DiscoveryService_WhoAmI_FullMethodName,
+				&Ydb_Discovery.WhoAmIRequest{},
+				&Ydb_Discovery.WhoAmIResponse{},
+			).Return(expectedErr)
+			client := Ydb_Discovery_V1.NewDiscoveryServiceClient(&connMock{
+				cc,
+			})
+			response, err := client.WhoAmI(ctx, &Ydb_Discovery.WhoAmIRequest{})
+			require.Error(t, err)
+			require.True(t, xerrors.IsTransportError(err, grpcCodes.Unavailable))
+			require.Nil(t, response)
+		})
+		t.Run("OperationError", func(t *testing.T) {
+			ctx := xtest.Context(t)
+			ctrl := gomock.NewController(t)
+			t.Run("discovery.WhoAmI", func(t *testing.T) {
+				cc := NewMockClientConnInterface(ctrl)
+				cc.EXPECT().Invoke(
+					gomock.Any(),
+					Ydb_Discovery_V1.DiscoveryService_WhoAmI_FullMethodName,
+					&Ydb_Discovery.WhoAmIRequest{},
+					&Ydb_Discovery.WhoAmIResponse{},
+				).DoAndReturn(func(ctx context.Context, method string, args any, reply any, opts ...grpc.CallOption) error {
+					res, ok := reply.(*Ydb_Discovery.WhoAmIResponse)
+					if !ok {
+						return fmt.Errorf("reply is not *Ydb_Discovery.WhoAmIResponse: %T", reply)
+					}
+
+					res.Operation = &Ydb_Operations.Operation{
+						Ready:  true,
+						Status: Ydb.StatusIds_UNAVAILABLE,
+					}
+
+					return nil
+				})
+				client := Ydb_Discovery_V1.NewDiscoveryServiceClient(&connMock{
+					cc,
+				})
+				response, err := client.WhoAmI(ctx, &Ydb_Discovery.WhoAmIRequest{})
+				require.Error(t, err)
+				require.True(t, xerrors.IsOperationError(err, Ydb.StatusIds_UNAVAILABLE))
+				require.Nil(t, response)
+			})
+			t.Run("query.BeginTransaction", func(t *testing.T) {
+				cc := NewMockClientConnInterface(ctrl)
+				cc.EXPECT().Invoke(
+					gomock.Any(),
+					Ydb_Query_V1.QueryService_BeginTransaction_FullMethodName,
+					&Ydb_Query.BeginTransactionRequest{},
+					&Ydb_Query.BeginTransactionResponse{},
+				).DoAndReturn(func(ctx context.Context, method string, args any, reply any, opts ...grpc.CallOption) error {
+					res, ok := reply.(*Ydb_Query.BeginTransactionResponse)
+					if !ok {
+						return fmt.Errorf("reply is not *Ydb_Query.BeginTransactionResponse: %T", reply)
+					}
+
+					res.Status = Ydb.StatusIds_UNAVAILABLE
+
+					return nil
+				})
+				client := Ydb_Query_V1.NewQueryServiceClient(&connMock{
+					cc,
+				})
+				response, err := client.BeginTransaction(ctx, &Ydb_Query.BeginTransactionRequest{})
+				require.Error(t, err)
+				require.True(t, xerrors.IsOperationError(err, Ydb.StatusIds_UNAVAILABLE))
+				require.Nil(t, response)
+			})
+		})
+	})
+}