Merge pull request #1689 from ydb-platform/query-client

* Switched internal type of result `ydb.Driver.Query()` from `*internal/query.Client` to `query.Client` interface

Author: asmyasnikov

CHANGELOG.md

@@ -1,3 +1,5 @@
+* Switched internal type of result `ydb.Driver.Query()` from `*internal/query.Client` to `query.Client` interface 
+
 ## v3.101.3
 * Added `query.TransactionActor` type alias to `query.TxActor` for compatibility with `table.Client` API's 
 * Removed comment `experimental` from `ydb.ParamsBuilder` and `ydb.ParamsFromMap`

driver.go

@@ -42,6 +42,7 @@ import (
 	"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/query"
 	"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"
@@ -228,7 +229,7 @@ func (d *Driver) Table() table.Client {
 }
 
 // Query returns query client
-func (d *Driver) Query() *internalQuery.Client {
+func (d *Driver) Query() query.Client {
 	return d.query.Must()
 }
 

dsn_test.go

@@ -30,7 +30,7 @@ func TestParse(t *testing.T) {
 		return xtable.New(context.Background(), nil, nil, opts...)
 	}
 	newQueryConn := func(opts ...xquery.Option) *xquery.Conn {
-		return xquery.New(context.Background(), nil, nil, opts...)
+		return xquery.New(context.Background(), nil, opts...)
 	}
 	compareConfigs := func(t *testing.T, lhs, rhs *config.Config) {
 		require.Equal(t, lhs.Secure(), rhs.Secure())

internal/query/client.go

@@ -49,6 +49,10 @@ type (
 	}
 )
 
+func (c *Client) Config() *config.Config {
+	return c.config
+}
+
 func fetchScriptResults(ctx context.Context,
 	client Ydb_Query_V1.QueryServiceClient,
 	opID string, opts ...options.FetchScriptOption,
@@ -511,28 +515,28 @@ func (c *Client) DoTx(ctx context.Context, op query.TxOperation, opts ...options
 	return nil
 }
 
-func CreateSession(ctx context.Context, c *Client) (*Session, error) {
+func CreateSession(ctx context.Context, client Ydb_Query_V1.QueryServiceClient, cfg *config.Config) (*Session, error) {
 	s, err := retry.RetryWithResult(ctx, func(ctx context.Context) (*Session, error) {
 		var (
 			createCtx    context.Context
 			cancelCreate context.CancelFunc
 		)
-		if d := c.config.SessionCreateTimeout(); d > 0 {
+		if d := cfg.SessionCreateTimeout(); d > 0 {
 			createCtx, cancelCreate = xcontext.WithTimeout(ctx, d)
 		} else {
 			createCtx, cancelCreate = xcontext.WithCancel(ctx)
 		}
 		defer cancelCreate()
 
-		s, err := createSession(createCtx, c.client,
-			WithDeleteTimeout(c.config.SessionDeleteTimeout()),
-			WithTrace(c.config.Trace()),
+		s, err := createSession(createCtx, client,
+			WithDeleteTimeout(cfg.SessionDeleteTimeout()),
+			WithTrace(cfg.Trace()),
 		)
 		if err != nil {
 			return nil, xerrors.WithStackTrace(err)
 		}
 
-		s.lazyTx = c.config.LazyTx()
+		s.lazyTx = cfg.LazyTx()
 
 		return s, nil
 	})

internal/xsql/conn_helpers.go

@@ -46,7 +46,7 @@ func (c *Conn) Engine() Engine {
 }
 
 func (c *Conn) GetDatabaseName() string {
-	return c.connector.Name()
+	return c.connector.parent.Name()
 }
 
 func (c *Conn) normalizePath(tableName string) string {
@@ -153,11 +153,11 @@ func isSysDir(databaseName, dirAbsPath string) bool {
 func (c *Conn) getTables(ctx context.Context, absPath string, recursive, excludeSysDirs bool) (
 	tables []string, _ error,
 ) {
-	if excludeSysDirs && isSysDir(c.connector.Name(), absPath) {
+	if excludeSysDirs && isSysDir(c.connector.parent.Name(), absPath) {
 		return nil, nil
 	}
 
-	d, err := c.connector.Scheme().ListDirectory(ctx, absPath)
+	d, err := c.connector.parent.Scheme().ListDirectory(ctx, absPath)
 	if err != nil {
 		return nil, xerrors.WithStackTrace(err)
 	}
@@ -189,7 +189,7 @@ func (c *Conn) GetTables(ctx context.Context, folder string, recursive, excludeS
 ) {
 	absPath := c.normalizePath(folder)
 
-	e, err := c.connector.Scheme().DescribePath(ctx, absPath)
+	e, err := c.connector.parent.Scheme().DescribePath(ctx, absPath)
 	if err != nil {
 		return nil, xerrors.WithStackTrace(err)
 	}

internal/xsql/connector.go

@@ -10,16 +10,19 @@ import (
 
 	"github.com/google/uuid"
 	"github.com/jonboulle/clockwork"
+	"github.com/ydb-platform/ydb-go-genproto/Ydb_Query_V1"
 	"google.golang.org/grpc"
 
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/bind"
-	"github.com/ydb-platform/ydb-go-sdk/v3/internal/query"
+	internalQuery "github.com/ydb-platform/ydb-go-sdk/v3/internal/query"
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/query/config"
 	"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"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xsql/xquery"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xsql/xtable"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xsync"
+	"github.com/ydb-platform/ydb-go-sdk/v3/query"
 	"github.com/ydb-platform/ydb-go-sdk/v3/retry/budget"
 	"github.com/ydb-platform/ydb-go-sdk/v3/scheme"
 	"github.com/ydb-platform/ydb-go-sdk/v3/scripting"
@@ -35,8 +38,9 @@ var (
 type (
 	Engine    uint8
 	Connector struct {
-		parent   ydbDriver
-		balancer grpc.ClientConnInterface
+		parent      ydbDriver
+		balancer    grpc.ClientConnInterface
+		queryConfig *config.Config
 
 		processor Engine
 
@@ -58,7 +62,7 @@ type (
 	ydbDriver interface {
 		Name() string
 		Table() table.Client
-		Query() *query.Client
+		Query() query.Client
 		Scripting() scripting.Client
 		Scheme() scheme.Client
 	}
@@ -75,6 +79,10 @@ func (e Engine) String() string {
 	}
 }
 
+func (c *Connector) Parent() ydbDriver {
+	return c.parent
+}
+
 func (c *Connector) RetryBudget() budget.Budget {
 	return c.retryBudget
 }
@@ -95,26 +103,6 @@ func (c *Connector) TraceRetry() *trace.Retry {
 	return c.traceRetry
 }
 
-func (c *Connector) Query() *query.Client {
-	return c.parent.Query()
-}
-
-func (c *Connector) Name() string {
-	return c.parent.Name()
-}
-
-func (c *Connector) Table() table.Client {
-	return c.parent.Table()
-}
-
-func (c *Connector) Scripting() scripting.Client {
-	return c.parent.Scripting()
-}
-
-func (c *Connector) Scheme() scheme.Client {
-	return c.parent.Scheme()
-}
-
 const (
 	QUERY = iota + 1
 	TABLE
@@ -131,7 +119,7 @@ func (c *Connector) Connect(ctx context.Context) (_ driver.Conn, finalErr error)
 
 	switch c.processor {
 	case QUERY:
-		s, err := query.CreateSession(ctx, c.Query())
+		s, err := internalQuery.CreateSession(ctx, Ydb_Query_V1.NewQueryServiceClient(c.balancer), c.queryConfig)
 		defer func() {
 			onDone(s, finalErr)
 		}()
@@ -143,7 +131,7 @@ func (c *Connector) Connect(ctx context.Context) (_ driver.Conn, finalErr error)
 
 		conn := &Conn{
 			processor: QUERY,
-			cc: xquery.New(ctx, c, s, append(
+			cc: xquery.New(ctx, s, append(
 				c.QueryOpts,
 				xquery.WithOnClose(func() {
 					c.conns.Delete(id)
@@ -159,7 +147,7 @@ func (c *Connector) Connect(ctx context.Context) (_ driver.Conn, finalErr error)
 		return conn, nil
 
 	case TABLE:
-		s, err := c.Table().CreateSession(ctx) //nolint:staticcheck
+		s, err := c.parent.Table().CreateSession(ctx) //nolint:staticcheck
 		defer func() {
 			onDone(s, finalErr)
 		}()
@@ -171,7 +159,7 @@ func (c *Connector) Connect(ctx context.Context) (_ driver.Conn, finalErr error)
 
 		conn := &Conn{
 			processor: TABLE,
-			cc: xtable.New(ctx, c, s, append(c.TableOpts,
+			cc: xtable.New(ctx, c.parent.Scripting(), s, append(c.TableOpts,
 				xtable.WithOnClose(func() {
 					c.conns.Delete(id)
 				}))...,
@@ -193,10 +181,6 @@ func (c *Connector) Driver() driver.Driver {
 	return c
 }
 
-func (c *Connector) Parent() ydbDriver {
-	return c.parent
-}
-
 func (c *Connector) Close() error {
 	select {
 	case <-c.done:
@@ -212,10 +196,16 @@ func (c *Connector) Close() error {
 	}
 }
 
-func Open(parent ydbDriver, balancer grpc.ClientConnInterface, opts ...Option) (_ *Connector, err error) {
+func Open(
+	parent ydbDriver,
+	balancer grpc.ClientConnInterface,
+	queryConfig *config.Config,
+	opts ...Option,
+) (_ *Connector, err error) {
 	c := &Connector{
-		parent:   parent,
-		balancer: balancer,
+		parent:      parent,
+		balancer:    balancer,
+		queryConfig: queryConfig,
 		processor: func() Engine {
 			if overQueryService, _ := strconv.ParseBool(os.Getenv("YDB_DATABASE_SQL_OVER_QUERY_SERVICE")); overQueryService {
 				return QUERY

internal/xsql/xquery/conn.go

@@ -28,7 +28,6 @@ type Parent interface {
 
 type Conn struct {
 	ctx     context.Context //nolint:containedctx
-	parent  Parent
 	session *query.Session
 	onClose []func()
 	closed  atomic.Bool
@@ -117,10 +116,9 @@ func (c *Conn) Explain(ctx context.Context, sql string, _ *params.Params) (ast s
 	return ast, plan, nil
 }
 
-func New(ctx context.Context, parent Parent, s *query.Session, opts ...Option) *Conn {
+func New(ctx context.Context, s *query.Session, opts ...Option) *Conn {
 	cc := &Conn{
 		ctx:     ctx,
-		parent:  parent,
 		session: s,
 	}
 

internal/xsql/xtable/conn.go

@@ -27,8 +27,8 @@ type (
 	Conn struct {
 		ctx context.Context //nolint:containedctx
 
-		parent  Parent
-		session table.ClosableSession // Immutable and r/o usage.
+		scriptingClient scripting.Client
+		session         table.ClosableSession // Immutable and r/o usage.
 
 		fakeTxModes []QueryMode
 
@@ -116,10 +116,10 @@ type resultNoRows struct{}
 func (resultNoRows) LastInsertId() (int64, error) { return 0, ErrUnsupported }
 func (resultNoRows) RowsAffected() (int64, error) { return 0, ErrUnsupported }
 
-func New(ctx context.Context, parent Parent, s table.ClosableSession, opts ...Option) *Conn {
+func New(ctx context.Context, scriptingClient scripting.Client, s table.ClosableSession, opts ...Option) *Conn {
 	cc := &Conn{
 		ctx:              ctx,
-		parent:           parent,
+		scriptingClient:  scriptingClient,
 		session:          s,
 		defaultQueryMode: DataQueryMode,
 		defaultTxControl: table.DefaultTxControl(),
@@ -169,7 +169,7 @@ func (c *Conn) executeSchemeQuery(ctx context.Context, sql string) (driver.Resul
 func (c *Conn) executeScriptingQuery(ctx context.Context, sql string, params *params.Params) (
 	driver.Result, error,
 ) {
-	res, err := c.parent.Scripting().StreamExecute(ctx, sql, params)
+	res, err := c.scriptingClient.StreamExecute(ctx, sql, params)
 	if err != nil {
 		return nil, badconn.Map(xerrors.WithStackTrace(err))
 	}
@@ -227,7 +227,7 @@ func (c *Conn) execScanQuery(ctx context.Context, sql string, params *params.Par
 func (c *Conn) execScriptingQuery(ctx context.Context, sql string, params *params.Params) (
 	driver.RowsNextResultSet, error,
 ) {
-	res, err := c.parent.Scripting().StreamExecute(ctx, sql, params)
+	res, err := c.scriptingClient.StreamExecute(ctx, sql, params)
 	if err != nil {
 		return nil, badconn.Map(xerrors.WithStackTrace(err))
 	}

sql.go

@@ -223,7 +223,7 @@ type SQLConnector interface {
 }
 
 func Connector(parent *Driver, opts ...ConnectorOption) (SQLConnector, error) {
-	c, err := xsql.Open(parent, parent.metaBalancer,
+	c, err := xsql.Open(parent, parent.metaBalancer, parent.query.Must().Config(),
 		append(
 			append(
 				parent.databaseSQLOptions,