sql.go - master - githubmirror/ydb-go-sdk

sql.go
322 строки9 KB
Aleksey Myasnikov
database/sql: add prefetch query result parts option (#2205)
18 июн 2026, 14:32
18 июн 2026, 14:32b4754dd
package ydb

import (
	"context"
	"database/sql"
	"database/sql/driver"
	"fmt"

	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb_Issue"

	"github.com/ydb-platform/ydb-go-sdk/v3/config"
	"github.com/ydb-platform/ydb-go-sdk/v3/internal/bind"
	"github.com/ydb-platform/ydb-go-sdk/v3/internal/meta"
	"github.com/ydb-platform/ydb-go-sdk/v3/internal/secret"
	"github.com/ydb-platform/ydb-go-sdk/v3/internal/tx"
	"github.com/ydb-platform/ydb-go-sdk/v3/internal/version"
	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xsql"
	"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/table"
	"github.com/ydb-platform/ydb-go-sdk/v3/table/options"
	"github.com/ydb-platform/ydb-go-sdk/v3/trace"
)

var d = &sqlDriver{} //nolint:gochecknoglobals

func init() { //nolint:gochecknoinits
	sql.Register("ydb", d)
	sql.Register("ydb/v3", d)
}

func withConnectorOptions(opts ...ConnectorOption) Option {
	return func(ctx context.Context, d *Driver) error {
		d.databaseSQLOptions = append(d.databaseSQLOptions, opts...)

		return nil
	}
}

type sqlDriver struct{}

var (
	_ driver.Driver        = &sqlDriver{}
	_ driver.DriverContext = &sqlDriver{}
)

// Open returns a new Driver to the ydb.
func (d *sqlDriver) Open(string) (driver.Conn, error) {
	return nil, xsql.ErrUnsupported
}

func (d *sqlDriver) OpenConnector(dataSourceName string) (driver.Connector, error) {
	db, err := Open(context.Background(), dataSourceName,
		With(config.WithBuildInfo(
			"database/sql",
			version.Version,
		)),
	)
	if err != nil {
		return nil, xerrors.WithStackTrace(fmt.Errorf(
			"failed to connect by data source name '%s': %w",
			secret.DSN(dataSourceName), err,
		))
	}

	c, err := Connector(db, append(db.databaseSQLOptions,
		xsql.WithOnClose(func(connector *xsql.Connector) {
			_ = db.Close(context.Background())
		}),
	)...)
	if err != nil {
		return nil, xerrors.WithStackTrace(fmt.Errorf("failed to create connector: %w", err))
	}

	return c, nil
}

type QueryMode int

const (
	unknownQueryMode = QueryMode(iota)
	DataQueryMode
	ExplainQueryMode
	ScanQueryMode
	SchemeQueryMode
	ScriptingQueryMode
	QueryExecuteQueryMode
)

// WithQueryMode set query mode for legacy database/sql driver
//
// For actual database/sql driver works over query service client and no needs query mode
func WithQueryMode(ctx context.Context, mode QueryMode) context.Context {
	switch mode {
	case ExplainQueryMode:
		return xsql.WithExplain(ctx)
	case DataQueryMode:
		return xtable.WithQueryMode(ctx, xtable.DataQueryMode)
	case ScanQueryMode:
		return xtable.WithQueryMode(ctx, xtable.ScanQueryMode)
	case SchemeQueryMode:
		return xtable.WithQueryMode(ctx, xtable.SchemeQueryMode)
	case ScriptingQueryMode:
		return xtable.WithQueryMode(ctx, xtable.ScriptingQueryMode)
	default:
		return ctx
	}
}

// WithTxControl modifies context for explicit define transaction control for a single query execute
//
// Allowed the table.TransactionControl and the query.TransactionControl
// table.TransactionControl and query.TransactionControl are the type aliases to internal tx.Control
func WithTxControl(ctx context.Context, txControl *tx.Control) context.Context {
	return tx.WithTxControl(ctx, txControl)
}

// WithCommitTxContext modifies context to request commit along with the query execution.
// When used inside a database/sql transaction, the next ExecContext or QueryContext call
// will commit the transaction together with the query in a single RPC call.
//
// This is an optimization to reduce latency by combining Execute + Commit into one call.
// After commit via context, the subsequent tx.Commit() or tx.Rollback() will be a no-op.
//
// Important: when using QueryContext with WithCommitTxContext, the server commits the
// transaction during query execution. However, on the client side the transaction is marked
// as completed only after the result rows are fully consumed (rows.Next() until false,
// then rows.Close()). If tx.Commit() is called before rows are consumed, an extra
// CommitTx RPC will be sent to the server (which is harmless but redundant).
//
// Example:
//
//	tx, _ := db.BeginTx(ctx, nil)
//	tx.ExecContext(ydb.WithCommitTxContext(ctx), "INSERT INTO ...")
//	// Transaction is already committed, tx.Commit() will be a no-op
//
// Experimental: https://github.com/ydb-platform/ydb-go-sdk/blob/master/VERSIONING.md#experimental
func WithCommitTxContext(ctx context.Context) context.Context {
	return tx.WithCommitTx(ctx)
}

// WithIssuesHandler sets a callback which is invoked for query
// execution result to receive YDB issue messages.
//
// Supported only for QueryService.
func WithIssuesHandler(ctx context.Context, callback func([]*Ydb_Issue.IssueMessage)) context.Context {
	return xquery.WithIssuesHandler(ctx, callback)
}

type ConnectorOption = xsql.Option

type QueryBindConnectorOption interface {
	ConnectorOption
	bind.Bind
}

func modeToMode(mode QueryMode) xtable.QueryMode {
	switch mode {
	case ScanQueryMode:
		return xtable.ScanQueryMode
	case SchemeQueryMode:
		return xtable.SchemeQueryMode
	case ScriptingQueryMode:
		return xtable.ScriptingQueryMode
	default:
		return xtable.DataQueryMode
	}
}

func WithDefaultQueryMode(mode QueryMode) ConnectorOption {
	if mode == QueryExecuteQueryMode {
		return xsql.WithQueryService(true)
	}

	return xsql.WithTableOptions(
		xtable.WithDefaultQueryMode(modeToMode(mode)),
	)
}

// WithQueryService is an experimental flag for create database/sql driver over query service client
//
// By default database/sql driver works over table service client
// Default will be changed to `WithQueryService` after March 2025
func WithQueryService(b bool) ConnectorOption {
	return xsql.WithQueryService(b)
}

func WithFakeTx(modes ...QueryMode) ConnectorOption {
	opts := make([]ConnectorOption, 0, len(modes))

	for _, mode := range modes {
		switch mode {
		case DataQueryMode:
			opts = append(opts,
				xsql.WithTableOptions(xtable.WithFakeTxModes(
					xtable.DataQueryMode,
				)),
			)
		case ScanQueryMode:
			opts = append(opts,
				xsql.WithTableOptions(xtable.WithFakeTxModes(
					xtable.ScanQueryMode,
				)),
			)
		case SchemeQueryMode:
			opts = append(opts,
				xsql.WithTableOptions(xtable.WithFakeTxModes(
					xtable.SchemeQueryMode,
				)),
			)
		case ScriptingQueryMode:
			opts = append(opts,
				xsql.WithTableOptions(xtable.WithFakeTxModes(
					xtable.ScriptingQueryMode,
				)),
				xsql.WithQueryOptions(xquery.WithFakeTx()),
			)
		case QueryExecuteQueryMode:
			opts = append(opts,
				xsql.WithQueryOptions(xquery.WithFakeTx()),
			)
		default:
		}
	}

	return xsql.Merge(opts...)
}

func WithTablePathPrefix(tablePathPrefix string) QueryBindConnectorOption {
	return xsql.WithTablePathPrefix(tablePathPrefix)
}

func WithAutoDeclare() QueryBindConnectorOption {
	return xsql.WithQueryBind(bind.AutoDeclare{})
}

func WithPositionalArgs() QueryBindConnectorOption {
	return xsql.WithQueryBind(bind.PositionalArgs{})
}

func WithWideTimeTypes(b bool) QueryBindConnectorOption {
	return xsql.WithQueryBind(bind.WideTimeTypes(b))
}

func WithNumericArgs() QueryBindConnectorOption {
	return xsql.WithQueryBind(bind.NumericArgs{})
}

func WithDefaultTxControl(txControl *table.TransactionControl) ConnectorOption {
	return xsql.WithTableOptions(xtable.WithDefaultTxControl(txControl))
}

func WithDefaultDataQueryOptions(opts ...options.ExecuteDataQueryOption) ConnectorOption {
	return xsql.WithTableOptions(xtable.WithDataOpts(opts...))
}

func WithDefaultScanQueryOptions(opts ...options.ExecuteScanQueryOption) ConnectorOption {
	return xsql.WithTableOptions(xtable.WithScanOpts(opts...))
}

func WithDatabaseSQLTrace(
	t trace.DatabaseSQL, //nolint:gocritic
) ConnectorOption {
	return xsql.WithTrace(&t)
}

func WithDisableServerBalancer() ConnectorOption {
	return xsql.WithDisableServerBalancer()
}

// WithPrefetchQueryResultParts enables prefetching of ExecuteQuery response parts
// for all queries executed through database/sql over Query Service.
//
// The value is passed to query.WithResponsePartPrefetch under the hood.
// Zero disables prefetch (the default).
func WithPrefetchQueryResultParts(parts int) ConnectorOption {
	return xsql.WithQueryOptions(xquery.WithResponsePartPrefetch(parts))
}

type SQLConnector interface {
	driver.Connector

	Close() error
}

func Connector(parent *Driver, opts ...ConnectorOption) (SQLConnector, error) {
	meta.WithBuildInfo(
		"database/sql",
		version.Version,
	)(parent.config.Meta())

	c, err := xsql.Open(parent, parent.metaBalancer, parent.query.Must().Config(),
		append(
			append(
				append(
					[]ConnectorOption{
						xsql.WithComposePanicCallback(parent.config.PanicCallback()),
					},
					parent.databaseSQLOptions...,
				),
				opts...,
			),
			xsql.WithTraceRetry(parent.config.TraceRetry()),
			xsql.WithRetryBudget(parent.config.RetryBudget()),
		)...,
	)
	if err != nil {
		return nil, xerrors.WithStackTrace(err)
	}

	return c, nil
}

func MustConnector(parent *Driver, opts ...ConnectorOption) SQLConnector {
	c, err := Connector(parent, opts...)
	if err != nil {
		panic(err)
	}

	return c
}