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