Merge pull request #1413 from ydb-platform/operation-service

* Added operation service client through `db.Operation()` method (sup…

Author: asmyasnikov

CHANGELOG.md

@@ -8,6 +8,7 @@
 * Removed experimental methods `query.Session.ReadResultSet` and  `query.Session.ReadRows`
 * Removed experimental methods `query.TxActor.ReadResultSet` and  `query.TxActor.ReadRows`
 * Removed experimental method `query.Client.Stats`
+* Added experimental support for operation service client through `db.Operation()` method (supports methods `Get`, `List`, `Cancel` and `Forget`)
 
 ## v3.76.6
 * Replaced requirements from go1.22 + experimantal flag to go1.23 for experimental range-over interface

driver.go

@@ -38,6 +38,7 @@ import (
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xsql"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xsync"
 	"github.com/ydb-platform/ydb-go-sdk/v3/log"
+	"github.com/ydb-platform/ydb-go-sdk/v3/operation"
 	"github.com/ydb-platform/ydb-go-sdk/v3/ratelimiter"
 	"github.com/ydb-platform/ydb-go-sdk/v3/scheme"
 	"github.com/ydb-platform/ydb-go-sdk/v3/scripting"
@@ -67,6 +68,8 @@ type Driver struct {
 	discovery        *xsync.Once[*internalDiscovery.Client]
 	discoveryOptions []discoveryConfig.Option
 
+	operation *xsync.Once[*operation.Client]
+
 	table        *xsync.Once[*internalTable.Client]
 	tableOptions []tableConfig.Option
 
@@ -148,8 +151,10 @@ func (d *Driver) Close(ctx context.Context) (finalErr error) {
 		d.scheme.Close,
 		d.scripting.Close,
 		d.table.Close,
+		d.operation.Close,
 		d.query.Close,
 		d.topic.Close,
+		d.discovery.Close,
 		d.balancer.Close,
 		d.pool.Release,
 	)
@@ -215,6 +220,13 @@ func (d *Driver) Discovery() discovery.Client {
 	return d.discovery.Must()
 }
 
+// Operation returns operation client
+//
+// Experimental: https://github.com/ydb-platform/ydb-go-sdk/blob/master/VERSIONING.md#experimental
+func (d *Driver) Operation() *operation.Client {
+	return d.operation.Must()
+}
+
 // Scripting returns scripting client
 func (d *Driver) Scripting() scripting.Client {
 	return d.scripting.Must()
@@ -505,6 +517,12 @@ func (d *Driver) connect(ctx context.Context) (err error) {
 		), nil
 	})
 
+	d.operation = xsync.OnceValue(func() (*operation.Client, error) {
+		return operation.New(xcontext.ValueOnly(ctx),
+			d.pool.Get(endpoint.New(d.config.Endpoint())),
+		), nil
+	})
+
 	d.scripting = xsync.OnceValue(func() (*internalScripting.Client, error) {
 		return internalScripting.New(xcontext.ValueOnly(ctx),
 			d.balancer,

internal/operation/options/options.go

@@ -0,0 +1,10 @@
+package options
+
+import "github.com/ydb-platform/ydb-go-genproto/protos/Ydb_Operations"
+
+type (
+	ListOperationsRequest struct {
+		Ydb_Operations.ListOperationsRequest
+	}
+	List func(r *ListOperationsRequest)
+)

internal/query/client.go

@@ -5,16 +5,13 @@ import (
 	"sync/atomic"
 	"time"
 
-	"github.com/ydb-platform/ydb-go-genproto/Ydb_Operation_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_Operations"
 	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb_Query"
 	"google.golang.org/grpc"
 
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/allocator"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/closer"
-	"github.com/ydb-platform/ydb-go-sdk/v3/internal/conn"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/operation"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/pool"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/query/config"
@@ -49,37 +46,14 @@ type (
 		InUse         atomic.Int32
 	}
 	Client struct {
-		config                 *config.Config
-		queryServiceClient     Ydb_Query_V1.QueryServiceClient
-		operationServiceClient Ydb_Operation_V1.OperationServiceClient
-		pool                   sessionPool
+		config             *config.Config
+		queryServiceClient Ydb_Query_V1.QueryServiceClient
+		pool               sessionPool
 
 		done chan struct{}
 	}
 )
 
-func fetchScriptResultsWithFallback(ctx context.Context,
-	queryServiceClient Ydb_Query_V1.QueryServiceClient,
-	operationServiceClient Ydb_Operation_V1.OperationServiceClient,
-	opID string, opts ...options.FetchScriptOption,
-) (*options.FetchScriptResult, error) {
-	r, err := fetchScriptResults(ctx, queryServiceClient, opID, opts...)
-	if err != nil {
-		if xerrors.IsOperationError(err, Ydb.StatusIds_BAD_REQUEST) {
-			r, err = scriptOperationStatus(ctx, operationServiceClient, opID)
-			if err != nil {
-				return nil, xerrors.WithStackTrace(err)
-			}
-
-			return r, nil
-		}
-
-		return nil, xerrors.WithStackTrace(err)
-	}
-
-	return r, nil
-}
-
 func fetchScriptResults(ctx context.Context,
 	client Ydb_Query_V1.QueryServiceClient,
 	opID string, opts ...options.FetchScriptOption,
@@ -115,13 +89,9 @@ func fetchScriptResults(ctx context.Context,
 		}
 
 		return &options.FetchScriptResult{
-			Ready:  true,
-			Status: response.GetStatus().String(),
-			Data: &options.FetchScriptResultData{
-				ResultSetIndex: response.GetResultSetIndex(),
-				ResultSet:      MaterializedResultSet(int(response.GetResultSetIndex()), columnNames, columnTypes, rows),
-				NextToken:      response.GetNextFetchToken(),
-			},
+			ResultSetIndex: response.GetResultSetIndex(),
+			ResultSet:      MaterializedResultSet(int(response.GetResultSetIndex()), columnNames, columnTypes, rows),
+			NextToken:      response.GetNextFetchToken(),
 		}, nil
 	}, retry.WithIdempotent(true))
 	if err != nil {
@@ -131,40 +101,11 @@ func fetchScriptResults(ctx context.Context,
 	return r, nil
 }
 
-func scriptOperationStatus(
-	ctx context.Context, client Ydb_Operation_V1.OperationServiceClient, opID string,
-) (*options.FetchScriptResult, error) {
-	status, err := retry.RetryWithResult(ctx, func(ctx context.Context) (*options.FetchScriptResult, error) {
-		response, err := client.GetOperation(conn.WithoutWrapping(ctx), &Ydb_Operations.GetOperationRequest{
-			Id: opID,
-		})
-		if err != nil {
-			return nil, xerrors.WithStackTrace(err)
-		}
-
-		var md Ydb_Query.ExecuteScriptMetadata
-		err = response.GetOperation().GetMetadata().UnmarshalTo(&md)
-		if err != nil {
-			return nil, xerrors.WithStackTrace(err)
-		}
-
-		return &options.FetchScriptResult{
-			Ready:  response.GetOperation().GetReady(),
-			Status: response.GetOperation().GetStatus().String(),
-		}, nil
-	})
-	if err != nil {
-		return nil, xerrors.WithStackTrace(err)
-	}
-
-	return status, nil
-}
-
 func (c *Client) FetchScriptResults(ctx context.Context,
 	opID string, opts ...options.FetchScriptOption,
 ) (*options.FetchScriptResult, error) {
 	r, err := retry.RetryWithResult(ctx, func(ctx context.Context) (*options.FetchScriptResult, error) {
-		r, err := fetchScriptResultsWithFallback(ctx, c.queryServiceClient, c.operationServiceClient, opID,
+		r, err := fetchScriptResults(ctx, c.queryServiceClient, opID,
 			append(opts, func(request *options.FetchScriptResultsRequest) {
 				request.Trace = c.config.Trace()
 			})...,
@@ -688,10 +629,9 @@ func New(ctx context.Context, balancer grpc.ClientConnInterface, cfg *config.Con
 	grpcClient := Ydb_Query_V1.NewQueryServiceClient(balancer)
 
 	client := &Client{
-		config:                 cfg,
-		queryServiceClient:     grpcClient,
-		operationServiceClient: Ydb_Operation_V1.NewOperationServiceClient(balancer),
-		done:                   make(chan struct{}),
+		config:             cfg,
+		queryServiceClient: grpcClient,
+		done:               make(chan struct{}),
 		pool: newPool(ctx, cfg, func(ctx context.Context) (_ *Session, err error) {
 			var (
 				createCtx    context.Context

internal/query/client_test.go

@@ -1263,10 +1263,10 @@ func TestQueryScript(t *testing.T) {
 		require.EqualValues(t, "123", op.ID)
 		r, err := fetchScriptResults(ctx, service, op.ID)
 		require.NoError(t, err)
-		require.EqualValues(t, 0, r.Data.ResultSetIndex)
-		require.Equal(t, "456", r.Data.NextToken)
-		require.NotNil(t, r.Data.ResultSet)
-		row, err := r.Data.ResultSet.NextRow(ctx)
+		require.EqualValues(t, 0, r.ResultSetIndex)
+		require.Equal(t, "456", r.NextToken)
+		require.NotNil(t, r.ResultSet)
+		row, err := r.ResultSet.NextRow(ctx)
 		require.NoError(t, err)
 		var (
 			a int

internal/query/options/execute_script.go

@@ -36,16 +36,11 @@ type (
 		ConsumedUnits float64
 		Metadata      ExecuteScriptOperationMetadata
 	}
-	FetchScriptResultData struct {
+	FetchScriptResult struct {
 		ResultSetIndex int64
 		ResultSet      result.Set
 		NextToken      string
 	}
-	FetchScriptResult struct {
-		Ready  bool
-		Status string
-		Data   *FetchScriptResultData
-	}
 )
 
 func WithFetchToken(fetchToken string) FetchScriptOption {