Draft: Query service implementations for database/sql interfaces

.github/workflows/tests.yml

@@ -133,6 +133,7 @@ jobs:
       YDB_CONNECTION_STRING_SECURE: grpcs://localhost:2135/local
       YDB_SSL_ROOT_CERTIFICATES_FILE: /tmp/ydb_certs/ca.pem
       YDB_SESSIONS_SHUTDOWN_URLS: http://localhost:8765/actors/kqp_proxy?force_shutdown=all
+      YDB_DATABASE_SQL_OVER_QUERY_SERVICE: 1
       HIDE_APPLICATION_OUTPUT: 1
     steps:
       - name: Checkout code

dsn.go

@@ -12,6 +12,7 @@ import (
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/connector"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/dsn"
 	tableSql "github.com/ydb-platform/ydb-go-sdk/v3/internal/table/conn"
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xcontext"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
 )
 
@@ -60,21 +61,21 @@ func parseConnectionString(dataSourceName string) (opts []Option, _ error) {
 		opts = append(opts, WithBalancer(balancers.FromConfig(balancer)))
 	}
 	if queryMode := info.Params.Get("go_query_mode"); queryMode != "" {
-		mode := tableSql.QueryModeFromString(queryMode)
+		mode := xcontext.QueryModeFromString(queryMode)
 		if mode == tableSql.UnknownQueryMode {
 			return nil, xerrors.WithStackTrace(fmt.Errorf("unknown query mode: %s", queryMode))
 		}
 		opts = append(opts, withConnectorOptions(connector.WithDefaultQueryMode(mode)))
 	} else if queryMode := info.Params.Get("query_mode"); queryMode != "" {
-		mode := tableSql.QueryModeFromString(queryMode)
+		mode := xcontext.QueryModeFromString(queryMode)
 		if mode == tableSql.UnknownQueryMode {
 			return nil, xerrors.WithStackTrace(fmt.Errorf("unknown query mode: %s", queryMode))
 		}
 		opts = append(opts, withConnectorOptions(connector.WithDefaultQueryMode(mode)))
 	}
 	if fakeTx := info.Params.Get("go_fake_tx"); fakeTx != "" {
 		for _, queryMode := range strings.Split(fakeTx, ",") {
-			mode := tableSql.QueryModeFromString(queryMode)
+			mode := xcontext.QueryModeFromString(queryMode)
 			if mode == tableSql.UnknownQueryMode {
 				return nil, xerrors.WithStackTrace(fmt.Errorf("unknown query mode: %s", queryMode))
 			}

internal/connector/connector.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"database/sql/driver"
 	"io"
+	"os"
 	"time"
 
 	"github.com/google/uuid"
@@ -184,9 +185,15 @@ func (c *Connector) Close() error {
 
 func Open(parent ydbDriver, balancer grpc.ClientConnInterface, opts ...Option) (_ *Connector, err error) {
 	c := &Connector{
-		parent:         parent,
-		balancer:       balancer,
-		queryProcessor: TABLE_SERVICE,
+		parent:   parent,
+		balancer: balancer,
+		queryProcessor: func() queryProcessor {
+			if v, has := os.LookupEnv("YDB_DATABASE_SQL_OVER_QUERY_SERVICE"); has && v != "" {
+				return QUERY_SERVICE
+			}
+
+			return TABLE_SERVICE
+		}(),
 		clock:          clockwork.NewRealClock(),
 		done:           make(chan struct{}),
 		trace:          &trace.DatabaseSQL{},

internal/query/conn/conn.go

@@ -2,132 +2,230 @@ package conn
 
 import (
 	"context"
+	"database/sql"
 	"database/sql/driver"
 	"sync/atomic"
-	"time"
 
 	"github.com/jonboulle/clockwork"
 
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/bind"
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/params"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/query"
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/query/options"
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/query/session"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/stack"
-	"github.com/ydb-platform/ydb-go-sdk/v3/internal/table/conn/badconn"
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/stats"
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/tx"
 	"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/retry/budget"
 	"github.com/ydb-platform/ydb-go-sdk/v3/trace"
 )
 
-type (
-	Parent interface {
-		Query() *query.Client
-		Trace() *trace.DatabaseSQL
-		TraceRetry() *trace.Retry
-		RetryBudget() budget.Budget
-		Bindings() bind.Bindings
-		Clock() clockwork.Clock
-	}
-	currentTx interface {
-		Rollback() error
-	}
-	Conn struct {
-		ctx     context.Context //nolint:containedctx
-		parent  Parent
-		session *query.Session
-		onClose []func()
-		closed  atomic.Bool
-		currentTx
-	}
-)
+type resultNoRows struct{}
 
-func (c *Conn) ID() string {
-	return c.session.ID()
-}
+func (resultNoRows) LastInsertId() (int64, error) { return 0, ErrUnsupported }
+func (resultNoRows) RowsAffected() (int64, error) { return 0, ErrUnsupported }
 
-func (c *Conn) IsValid() bool {
-	panic("implement me")
-}
+var _ driver.Result = resultNoRows{}
 
-func (c *Conn) CheckNamedValue(value *driver.NamedValue) error {
-	panic("implement me")
+type Parent interface {
+	Query() *query.Client
+	Trace() *trace.DatabaseSQL
+	TraceRetry() *trace.Retry
+	RetryBudget() budget.Budget
+	Bindings() bind.Bindings
+	Clock() clockwork.Clock
 }
 
-func (c *Conn) Ping(ctx context.Context) error {
-	panic("implement me")
+type currentTx interface {
+	tx.Identifier
+	driver.Tx
+	driver.ExecerContext
+	driver.QueryerContext
+	driver.ConnPrepareContext
+	Rollback() error
 }
 
-func (c *Conn) PrepareContext(ctx context.Context, query string) (driver.Stmt, error) {
-	panic("implement me")
+type Conn struct {
+	currentTx
+	ctx       context.Context //nolint:containedctx
+	parent    Parent
+	session   *query.Session
+	onClose   []func()
+	closed    atomic.Bool
+	lastUsage atomic.Int64
 }
 
-func (c *Conn) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx, error) {
-	panic("implement me")
-}
+func New(ctx context.Context, parent Parent, s *query.Session, opts ...Option) *Conn {
+	cc := &Conn{
+		ctx:     ctx,
+		parent:  parent,
+		session: s,
+	}
+
+	for _, opt := range opts {
+		if opt != nil {
+			opt(cc)
+		}
+	}
 
-func (c *Conn) ExecContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Result, error) {
-	panic("implement me")
+	return cc
 }
 
-func (c *Conn) QueryContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Rows, error) {
-	panic("implement me")
+func (c *Conn) isReady() bool {
+	return c.session.Status() == session.StatusIdle.String()
 }
 
-func (c *Conn) Prepare(query string) (driver.Stmt, error) {
-	panic("implement me")
+func (c *Conn) normalize(q string, args ...driver.NamedValue) (query string, _ params.Parameters, _ error) {
+	queryArgs := make([]any, len(args))
+	for i := range args {
+		queryArgs[i] = args[i]
+	}
+
+	return c.parent.Bindings().RewriteQuery(q, queryArgs...)
 }
 
-func (c *Conn) Close() (finalErr error) {
-	if !c.closed.CompareAndSwap(false, true) {
-		return badconn.Map(xerrors.WithStackTrace(errConnClosedEarly))
+func (c *Conn) beginTx(ctx context.Context, txOptions driver.TxOptions) (tx currentTx, finalErr error) {
+	onDone := trace.DatabaseSQLOnConnBegin(c.parent.Trace(), &ctx,
+		stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/v3/internal/query/conn.(*Conn).beginTx"),
+	)
+	defer func() {
+		onDone(tx, finalErr)
+	}()
+
+	if c.currentTx != nil {
+		return nil, xerrors.WithStackTrace(xerrors.AlreadyHasTx(c.currentTx.ID()))
+	}
+
+	tx, err := beginTx(ctx, c, txOptions)
+	if err != nil {
+		return nil, xerrors.WithStackTrace(err)
 	}
 
+	c.currentTx = tx
+
+	return tx, nil
+}
+
+func (c *Conn) execContext(
+	ctx context.Context,
+	query string,
+	args []driver.NamedValue,
+) (_ driver.Result, finalErr error) {
 	defer func() {
-		for _, onClose := range c.onClose {
-			onClose()
-		}
+		c.lastUsage.Store(c.parent.Clock().Now().Unix())
 	}()
 
-	var (
-		ctx    = c.ctx
-		onDone = trace.DatabaseSQLOnConnClose(
-			c.parent.Trace(), &ctx,
-			stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/v3/internal/query/conn.(*Conn).Close"),
-		)
+	if !c.isReady() {
+		return nil, xerrors.WithStackTrace(errNotReadyConn)
+	}
+
+	if c.currentTx != nil {
+		return c.currentTx.ExecContext(ctx, query, args)
+	}
+
+	onDone := trace.DatabaseSQLOnConnExec(c.parent.Trace(), &ctx,
+		stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/v3/internal/query/conn.(*Conn).execContext"),
+		query, xcontext.UnknownQueryMode.String(), xcontext.IsIdempotent(ctx), c.parent.Clock().Since(c.LastUsage()),
 	)
 	defer func() {
 		onDone(finalErr)
 	}()
-	if c.currentTx != nil {
-		_ = c.currentTx.Rollback()
+
+	normalizedQuery, params, err := c.normalize(query, args...)
+	if err != nil {
+		return nil, xerrors.WithStackTrace(err)
 	}
-	err := c.session.Close(xcontext.ValueOnly(ctx))
+
+	err = c.session.Exec(ctx, normalizedQuery, options.WithParameters(&params))
 	if err != nil {
-		return badconn.Map(xerrors.WithStackTrace(err))
+		return nil, xerrors.WithStackTrace(err)
 	}
 
-	return nil
+	return resultNoRows{}, nil
 }
 
-func (c *Conn) Begin() (driver.Tx, error) {
-	panic("implement me")
-}
+func (c *Conn) queryContext(ctx context.Context, queryString string, args []driver.NamedValue) (
+	_ driver.Rows, finalErr error,
+) {
+	defer func() {
+		c.lastUsage.Store(c.parent.Clock().Now().Unix())
+	}()
+
+	if !c.isReady() {
+		return nil, xerrors.WithStackTrace(errNotReadyConn)
+	}
+
+	if c.currentTx != nil {
+		return c.currentTx.QueryContext(ctx, queryString, args)
+	}
+
+	onDone := trace.DatabaseSQLOnConnQuery(c.parent.Trace(), &ctx,
+		stack.FunctionID("github.com/ydb-platform/ydb-go-sdk/v3/internal/query/conn.(*Conn).queryContext"),
+		queryString, xcontext.UnknownQueryMode.String(), xcontext.IsIdempotent(ctx), c.parent.Clock().Since(c.LastUsage()),
+	)
+
+	defer func() {
+		onDone(finalErr)
+	}()
+
+	normalizedQuery, parameters, err := c.normalize(queryString, args...)
+	if err != nil {
+		return nil, xerrors.WithStackTrace(err)
+	}
 
-func (c *Conn) LastUsage() time.Time {
-	panic("implement me")
+	queryMode := xcontext.QueryModeFromContext(ctx, xcontext.UnknownQueryMode)
+
+	if queryMode == xcontext.ExplainQueryMode {
+		return c.queryContextExplain(ctx, normalizedQuery, parameters)
+	}
+
+	return c.queryContextOther(ctx, normalizedQuery, parameters)
 }
 
-func New(ctx context.Context, parent Parent, s *query.Session, opts ...Option) *Conn {
-	cc := &Conn{
-		ctx:     ctx,
-		parent:  parent,
-		session: s,
+func (c *Conn) queryContextOther(
+	ctx context.Context,
+	queryString string,
+	parameters params.Parameters,
+) (driver.Rows, error) {
+	res, err := c.session.Query(
+		ctx, queryString,
+		options.WithParameters(&parameters),
+	)
+	if err != nil {
+		return nil, xerrors.WithStackTrace(err)
 	}
 
-	for _, opt := range opts {
-		if opt != nil {
-			opt(cc)
-		}
+	return &rows{
+		conn:   c,
+		result: res,
+	}, nil
+}
+
+func (c *Conn) queryContextExplain(
+	ctx context.Context,
+	queryString string,
+	parameters params.Parameters,
+) (driver.Rows, error) {
+	var ast, plan string
+	_, err := c.session.Query(
+		ctx, queryString,
+		options.WithParameters(&parameters),
+		options.WithExecMode(options.ExecModeExplain),
+		options.WithStatsMode(options.StatsModeNone, func(stats stats.QueryStats) {
+			ast = stats.QueryAST()
+			plan = stats.QueryPlan()
+		}),
+	)
+	if err != nil {
+		return nil, xerrors.WithStackTrace(err)
 	}
 
-	return cc
+	return &single{
+		values: []sql.NamedArg{
+			sql.Named("AST", ast),
+			sql.Named("Plan", plan),
+		},
+	}, nil
 }