github.com/ydb-platform/ydb-go-sdk/v3@v3.89.2/sql.go (about)

     1  package ydb
     2  
     3  import (
     4  	"context"
     5  	"database/sql"
     6  	"database/sql/driver"
     7  	"fmt"
     8  
     9  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/bind"
    10  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
    11  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xsql"
    12  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xsync"
    13  	"github.com/ydb-platform/ydb-go-sdk/v3/table"
    14  	"github.com/ydb-platform/ydb-go-sdk/v3/table/options"
    15  	"github.com/ydb-platform/ydb-go-sdk/v3/trace"
    16  )
    17  
    18  var d = &sqlDriver{connectors: make(map[*xsql.Connector]*Driver)} //nolint:gochecknoglobals
    19  
    20  func init() { //nolint:gochecknoinits
    21  	sql.Register("ydb", d)
    22  	sql.Register("ydb/v3", d)
    23  }
    24  
    25  func withConnectorOptions(opts ...ConnectorOption) Option {
    26  	return func(ctx context.Context, d *Driver) error {
    27  		d.databaseSQLOptions = append(d.databaseSQLOptions, opts...)
    28  
    29  		return nil
    30  	}
    31  }
    32  
    33  type sqlDriver struct {
    34  	connectors    map[*xsql.Connector]*Driver
    35  	connectorsMtx xsync.RWMutex
    36  }
    37  
    38  var (
    39  	_ driver.Driver        = &sqlDriver{}
    40  	_ driver.DriverContext = &sqlDriver{}
    41  )
    42  
    43  func (d *sqlDriver) Close() error {
    44  	var connectors map[*xsql.Connector]*Driver
    45  	d.connectorsMtx.WithRLock(func() {
    46  		connectors = d.connectors
    47  	})
    48  	var errs []error
    49  	for c := range connectors {
    50  		if err := c.Close(); err != nil {
    51  			errs = append(errs, err)
    52  		}
    53  	}
    54  	if len(errs) > 0 {
    55  		return xerrors.NewWithIssues("ydb legacy driver close failed", errs...)
    56  	}
    57  
    58  	return nil
    59  }
    60  
    61  // Open returns a new Driver to the ydb.
    62  func (d *sqlDriver) Open(string) (driver.Conn, error) {
    63  	return nil, xsql.ErrUnsupported
    64  }
    65  
    66  func (d *sqlDriver) OpenConnector(dataSourceName string) (driver.Connector, error) {
    67  	db, err := Open(context.Background(), dataSourceName)
    68  	if err != nil {
    69  		return nil, xerrors.WithStackTrace(fmt.Errorf("failed to connect by data source name '%s': %w", dataSourceName, err))
    70  	}
    71  
    72  	return Connector(db, db.databaseSQLOptions...)
    73  }
    74  
    75  func (d *sqlDriver) attach(c *xsql.Connector, parent *Driver) {
    76  	d.connectorsMtx.WithLock(func() {
    77  		d.connectors[c] = parent
    78  	})
    79  }
    80  
    81  func (d *sqlDriver) detach(c *xsql.Connector) {
    82  	d.connectorsMtx.WithLock(func() {
    83  		delete(d.connectors, c)
    84  	})
    85  }
    86  
    87  type QueryMode = xsql.QueryMode
    88  
    89  const (
    90  	DataQueryMode      = xsql.DataQueryMode
    91  	ExplainQueryMode   = xsql.ExplainQueryMode
    92  	ScanQueryMode      = xsql.ScanQueryMode
    93  	SchemeQueryMode    = xsql.SchemeQueryMode
    94  	ScriptingQueryMode = xsql.ScriptingQueryMode
    95  )
    96  
    97  func WithQueryMode(ctx context.Context, mode QueryMode) context.Context {
    98  	return xsql.WithQueryMode(ctx, mode)
    99  }
   100  
   101  func WithTxControl(ctx context.Context, txc *table.TransactionControl) context.Context {
   102  	return xsql.WithTxControl(ctx, txc)
   103  }
   104  
   105  type ConnectorOption = xsql.ConnectorOption
   106  
   107  type QueryBindConnectorOption interface {
   108  	ConnectorOption
   109  	bind.Bind
   110  }
   111  
   112  func WithDefaultQueryMode(mode QueryMode) ConnectorOption {
   113  	return xsql.WithDefaultQueryMode(mode)
   114  }
   115  
   116  func WithFakeTx(mode QueryMode) ConnectorOption {
   117  	return xsql.WithFakeTx(mode)
   118  }
   119  
   120  func WithTablePathPrefix(tablePathPrefix string) QueryBindConnectorOption {
   121  	return xsql.WithTablePathPrefix(tablePathPrefix)
   122  }
   123  
   124  func WithAutoDeclare() QueryBindConnectorOption {
   125  	return xsql.WithQueryBind(bind.AutoDeclare{})
   126  }
   127  
   128  func WithPositionalArgs() QueryBindConnectorOption {
   129  	return xsql.WithQueryBind(bind.PositionalArgs{})
   130  }
   131  
   132  func WithNumericArgs() QueryBindConnectorOption {
   133  	return xsql.WithQueryBind(bind.NumericArgs{})
   134  }
   135  
   136  func WithDefaultTxControl(txControl *table.TransactionControl) ConnectorOption {
   137  	return xsql.WithDefaultTxControl(txControl)
   138  }
   139  
   140  func WithDefaultDataQueryOptions(opts ...options.ExecuteDataQueryOption) ConnectorOption {
   141  	return xsql.WithDefaultDataQueryOptions(opts...)
   142  }
   143  
   144  func WithDefaultScanQueryOptions(opts ...options.ExecuteScanQueryOption) ConnectorOption {
   145  	return xsql.WithDefaultScanQueryOptions(opts...)
   146  }
   147  
   148  func WithDatabaseSQLTrace(
   149  	t trace.DatabaseSQL, //nolint:gocritic
   150  	opts ...trace.DatabaseSQLComposeOption,
   151  ) ConnectorOption {
   152  	return xsql.WithTrace(&t, opts...)
   153  }
   154  
   155  func WithDisableServerBalancer() ConnectorOption {
   156  	return xsql.WithDisableServerBalancer()
   157  }
   158  
   159  type SQLConnector interface {
   160  	driver.Connector
   161  
   162  	Close() error
   163  }
   164  
   165  func Connector(parent *Driver, opts ...ConnectorOption) (SQLConnector, error) {
   166  	c, err := xsql.Open(parent,
   167  		append(
   168  			append(
   169  				parent.databaseSQLOptions,
   170  				opts...,
   171  			),
   172  			xsql.WithOnClose(d.detach),
   173  			xsql.WithTraceRetry(parent.config.TraceRetry()),
   174  			xsql.WithretryBudget(parent.config.RetryBudget()),
   175  		)...,
   176  	)
   177  	if err != nil {
   178  		return nil, xerrors.WithStackTrace(err)
   179  	}
   180  	d.attach(c, parent)
   181  
   182  	return c, nil
   183  }
   184  
   185  func MustConnector(parent *Driver, opts ...ConnectorOption) SQLConnector {
   186  	c, err := Connector(parent, opts...)
   187  	if err != nil {
   188  		panic(err)
   189  	}
   190  
   191  	return c
   192  }