Source file src/database/sql/sql.go

     1  // Copyright 2011 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  // Package sql provides a generic interface around SQL (or SQL-like)
     6  // databases.
     7  //
     8  // The sql package must be used in conjunction with a database driver.
     9  // See https://golang.org/s/sqldrivers for a list of drivers.
    10  //
    11  // Drivers that do not support context cancellation will not return until
    12  // after the query is completed.
    13  //
    14  // For usage examples, see the wiki page at
    15  // https://golang.org/s/sqlwiki.
    16  package sql
    17  
    18  import (
    19  	"context"
    20  	"database/sql/driver"
    21  	"errors"
    22  	"fmt"
    23  	"io"
    24  	"maps"
    25  	"math/rand/v2"
    26  	"reflect"
    27  	"runtime"
    28  	"slices"
    29  	"strconv"
    30  	"sync"
    31  	"sync/atomic"
    32  	"time"
    33  	_ "unsafe"
    34  )
    35  
    36  var driversMu sync.RWMutex
    37  
    38  // drivers should be an internal detail,
    39  // but widely used packages access it using linkname.
    40  // (It is extra wrong that they linkname drivers but not driversMu.)
    41  // Notable members of the hall of shame include:
    42  //   - github.com/instana/go-sensor
    43  //
    44  // Do not remove or change the type signature.
    45  // See go.dev/issue/67401.
    46  //
    47  //go:linkname drivers
    48  var drivers = make(map[string]driver.Driver)
    49  
    50  // nowFunc returns the current time; it's overridden in tests.
    51  var nowFunc = time.Now
    52  
    53  // Register makes a database driver available by the provided name.
    54  // If Register is called twice with the same name or if driver is nil,
    55  // it panics.
    56  func Register(name string, driver driver.Driver) {
    57  	driversMu.Lock()
    58  	defer driversMu.Unlock()
    59  	if driver == nil {
    60  		panic("sql: Register driver is nil")
    61  	}
    62  	if _, dup := drivers[name]; dup {
    63  		panic("sql: Register called twice for driver " + name)
    64  	}
    65  	drivers[name] = driver
    66  }
    67  
    68  func unregisterAllDrivers() {
    69  	driversMu.Lock()
    70  	defer driversMu.Unlock()
    71  	// For tests.
    72  	drivers = make(map[string]driver.Driver)
    73  }
    74  
    75  // Drivers returns a sorted list of the names of the registered drivers.
    76  func Drivers() []string {
    77  	driversMu.RLock()
    78  	defer driversMu.RUnlock()
    79  	return slices.Sorted(maps.Keys(drivers))
    80  }
    81  
    82  // A NamedArg is a named argument. NamedArg values may be used as
    83  // arguments to [DB.Query] or [DB.Exec] and bind to the corresponding named
    84  // parameter in the SQL statement.
    85  //
    86  // For a more concise way to create NamedArg values, see
    87  // the [Named] function.
    88  type NamedArg struct {
    89  	_NamedFieldsRequired struct{}
    90  
    91  	// Name is the name of the parameter placeholder.
    92  	//
    93  	// If empty, the ordinal position in the argument list will be
    94  	// used.
    95  	//
    96  	// Name must omit any symbol prefix.
    97  	Name string
    98  
    99  	// Value is the value of the parameter.
   100  	// It may be assigned the same value types as the query
   101  	// arguments.
   102  	Value any
   103  }
   104  
   105  // Named provides a more concise way to create [NamedArg] values.
   106  //
   107  // Example usage:
   108  //
   109  //	db.ExecContext(ctx, `
   110  //	    delete from Invoice
   111  //	    where
   112  //	        TimeCreated < @end
   113  //	        and TimeCreated >= @start;`,
   114  //	    sql.Named("start", startTime),
   115  //	    sql.Named("end", endTime),
   116  //	)
   117  func Named(name string, value any) NamedArg {
   118  	// This method exists because the go1compat promise
   119  	// doesn't guarantee that structs don't grow more fields,
   120  	// so unkeyed struct literals are a vet error. Thus, we don't
   121  	// want to allow sql.NamedArg{name, value}.
   122  	return NamedArg{Name: name, Value: value}
   123  }
   124  
   125  // IsolationLevel is the transaction isolation level used in [TxOptions].
   126  type IsolationLevel int
   127  
   128  // Various isolation levels that drivers may support in [DB.BeginTx].
   129  // If a driver does not support a given isolation level an error may be returned.
   130  //
   131  // See https://en.wikipedia.org/wiki/Isolation_(database_systems)#Isolation_levels.
   132  const (
   133  	LevelDefault IsolationLevel = iota
   134  	LevelReadUncommitted
   135  	LevelReadCommitted
   136  	LevelWriteCommitted
   137  	LevelRepeatableRead
   138  	LevelSnapshot
   139  	LevelSerializable
   140  	LevelLinearizable
   141  )
   142  
   143  // String returns the name of the transaction isolation level.
   144  func (i IsolationLevel) String() string {
   145  	switch i {
   146  	case LevelDefault:
   147  		return "Default"
   148  	case LevelReadUncommitted:
   149  		return "Read Uncommitted"
   150  	case LevelReadCommitted:
   151  		return "Read Committed"
   152  	case LevelWriteCommitted:
   153  		return "Write Committed"
   154  	case LevelRepeatableRead:
   155  		return "Repeatable Read"
   156  	case LevelSnapshot:
   157  		return "Snapshot"
   158  	case LevelSerializable:
   159  		return "Serializable"
   160  	case LevelLinearizable:
   161  		return "Linearizable"
   162  	default:
   163  		return "IsolationLevel(" + strconv.Itoa(int(i)) + ")"
   164  	}
   165  }
   166  
   167  var _ fmt.Stringer = LevelDefault
   168  
   169  // TxOptions holds the transaction options to be used in [DB.BeginTx].
   170  type TxOptions struct {
   171  	// Isolation is the transaction isolation level.
   172  	// If zero, the driver or database's default level is used.
   173  	Isolation IsolationLevel
   174  	ReadOnly  bool
   175  }
   176  
   177  // RawBytes is a byte slice that holds a reference to memory owned by
   178  // the database itself. After a [Rows.Scan] into a RawBytes, the slice is only
   179  // valid until the next call to [Rows.Next], [Rows.Scan], or [Rows.Close].
   180  type RawBytes []byte
   181  
   182  // NullString represents a string that may be null.
   183  // NullString implements the [Scanner] interface so
   184  // it can be used as a scan destination:
   185  //
   186  //	var s NullString
   187  //	err := db.QueryRow("SELECT name FROM foo WHERE id=?", id).Scan(&s)
   188  //	...
   189  //	if s.Valid {
   190  //	   // use s.String
   191  //	} else {
   192  //	   // NULL value
   193  //	}
   194  type NullString struct {
   195  	String string
   196  	Valid  bool // Valid is true if String is not NULL
   197  }
   198  
   199  // Scan implements the [Scanner] interface.
   200  func (ns *NullString) Scan(value any) error {
   201  	if value == nil {
   202  		ns.String, ns.Valid = "", false
   203  		return nil
   204  	}
   205  	err := convertAssign(&ns.String, value)
   206  	ns.Valid = err == nil
   207  	return err
   208  }
   209  
   210  // Value implements the [driver.Valuer] interface.
   211  func (ns NullString) Value() (driver.Value, error) {
   212  	if !ns.Valid {
   213  		return nil, nil
   214  	}
   215  	return ns.String, nil
   216  }
   217  
   218  // NullInt64 represents an int64 that may be null.
   219  // NullInt64 implements the [Scanner] interface so
   220  // it can be used as a scan destination, similar to [NullString].
   221  type NullInt64 struct {
   222  	Int64 int64
   223  	Valid bool // Valid is true if Int64 is not NULL
   224  }
   225  
   226  // Scan implements the [Scanner] interface.
   227  func (n *NullInt64) Scan(value any) error {
   228  	if value == nil {
   229  		n.Int64, n.Valid = 0, false
   230  		return nil
   231  	}
   232  	err := convertAssign(&n.Int64, value)
   233  	n.Valid = err == nil
   234  	return err
   235  }
   236  
   237  // Value implements the [driver.Valuer] interface.
   238  func (n NullInt64) Value() (driver.Value, error) {
   239  	if !n.Valid {
   240  		return nil, nil
   241  	}
   242  	return n.Int64, nil
   243  }
   244  
   245  // NullInt32 represents an int32 that may be null.
   246  // NullInt32 implements the [Scanner] interface so
   247  // it can be used as a scan destination, similar to [NullString].
   248  type NullInt32 struct {
   249  	Int32 int32
   250  	Valid bool // Valid is true if Int32 is not NULL
   251  }
   252  
   253  // Scan implements the [Scanner] interface.
   254  func (n *NullInt32) Scan(value any) error {
   255  	if value == nil {
   256  		n.Int32, n.Valid = 0, false
   257  		return nil
   258  	}
   259  	err := convertAssign(&n.Int32, value)
   260  	n.Valid = err == nil
   261  	return err
   262  }
   263  
   264  // Value implements the [driver.Valuer] interface.
   265  func (n NullInt32) Value() (driver.Value, error) {
   266  	if !n.Valid {
   267  		return nil, nil
   268  	}
   269  	return int64(n.Int32), nil
   270  }
   271  
   272  // NullInt16 represents an int16 that may be null.
   273  // NullInt16 implements the [Scanner] interface so
   274  // it can be used as a scan destination, similar to [NullString].
   275  type NullInt16 struct {
   276  	Int16 int16
   277  	Valid bool // Valid is true if Int16 is not NULL
   278  }
   279  
   280  // Scan implements the [Scanner] interface.
   281  func (n *NullInt16) Scan(value any) error {
   282  	if value == nil {
   283  		n.Int16, n.Valid = 0, false
   284  		return nil
   285  	}
   286  	err := convertAssign(&n.Int16, value)
   287  	n.Valid = err == nil
   288  	return err
   289  }
   290  
   291  // Value implements the [driver.Valuer] interface.
   292  func (n NullInt16) Value() (driver.Value, error) {
   293  	if !n.Valid {
   294  		return nil, nil
   295  	}
   296  	return int64(n.Int16), nil
   297  }
   298  
   299  // NullByte represents a byte that may be null.
   300  // NullByte implements the [Scanner] interface so
   301  // it can be used as a scan destination, similar to [NullString].
   302  type NullByte struct {
   303  	Byte  byte
   304  	Valid bool // Valid is true if Byte is not NULL
   305  }
   306  
   307  // Scan implements the [Scanner] interface.
   308  func (n *NullByte) Scan(value any) error {
   309  	if value == nil {
   310  		n.Byte, n.Valid = 0, false
   311  		return nil
   312  	}
   313  	err := convertAssign(&n.Byte, value)
   314  	n.Valid = err == nil
   315  	return err
   316  }
   317  
   318  // Value implements the [driver.Valuer] interface.
   319  func (n NullByte) Value() (driver.Value, error) {
   320  	if !n.Valid {
   321  		return nil, nil
   322  	}
   323  	return int64(n.Byte), nil
   324  }
   325  
   326  // NullFloat64 represents a float64 that may be null.
   327  // NullFloat64 implements the [Scanner] interface so
   328  // it can be used as a scan destination, similar to [NullString].
   329  type NullFloat64 struct {
   330  	Float64 float64
   331  	Valid   bool // Valid is true if Float64 is not NULL
   332  }
   333  
   334  // Scan implements the [Scanner] interface.
   335  func (n *NullFloat64) Scan(value any) error {
   336  	if value == nil {
   337  		n.Float64, n.Valid = 0, false
   338  		return nil
   339  	}
   340  	err := convertAssign(&n.Float64, value)
   341  	n.Valid = err == nil
   342  	return err
   343  }
   344  
   345  // Value implements the [driver.Valuer] interface.
   346  func (n NullFloat64) Value() (driver.Value, error) {
   347  	if !n.Valid {
   348  		return nil, nil
   349  	}
   350  	return n.Float64, nil
   351  }
   352  
   353  // NullBool represents a bool that may be null.
   354  // NullBool implements the [Scanner] interface so
   355  // it can be used as a scan destination, similar to [NullString].
   356  type NullBool struct {
   357  	Bool  bool
   358  	Valid bool // Valid is true if Bool is not NULL
   359  }
   360  
   361  // Scan implements the [Scanner] interface.
   362  func (n *NullBool) Scan(value any) error {
   363  	if value == nil {
   364  		n.Bool, n.Valid = false, false
   365  		return nil
   366  	}
   367  	err := convertAssign(&n.Bool, value)
   368  	n.Valid = err == nil
   369  	return err
   370  }
   371  
   372  // Value implements the [driver.Valuer] interface.
   373  func (n NullBool) Value() (driver.Value, error) {
   374  	if !n.Valid {
   375  		return nil, nil
   376  	}
   377  	return n.Bool, nil
   378  }
   379  
   380  // NullTime represents a [time.Time] that may be null.
   381  // NullTime implements the [Scanner] interface so
   382  // it can be used as a scan destination, similar to [NullString].
   383  type NullTime struct {
   384  	Time  time.Time
   385  	Valid bool // Valid is true if Time is not NULL
   386  }
   387  
   388  // Scan implements the [Scanner] interface.
   389  func (n *NullTime) Scan(value any) error {
   390  	if value == nil {
   391  		n.Time, n.Valid = time.Time{}, false
   392  		return nil
   393  	}
   394  	err := convertAssign(&n.Time, value)
   395  	n.Valid = err == nil
   396  	return err
   397  }
   398  
   399  // Value implements the [driver.Valuer] interface.
   400  func (n NullTime) Value() (driver.Value, error) {
   401  	if !n.Valid {
   402  		return nil, nil
   403  	}
   404  	return n.Time, nil
   405  }
   406  
   407  // Null represents a value that may be null.
   408  // Null implements the [Scanner] interface so
   409  // it can be used as a scan destination:
   410  //
   411  //	var s Null[string]
   412  //	err := db.QueryRow("SELECT name FROM foo WHERE id=?", id).Scan(&s)
   413  //	...
   414  //	if s.Valid {
   415  //	   // use s.V
   416  //	} else {
   417  //	   // NULL value
   418  //	}
   419  //
   420  // T should be one of the types accepted by [driver.Value].
   421  type Null[T any] struct {
   422  	V     T
   423  	Valid bool
   424  }
   425  
   426  func (n *Null[T]) Scan(value any) error {
   427  	if value == nil {
   428  		n.V, n.Valid = *new(T), false
   429  		return nil
   430  	}
   431  	err := convertAssign(&n.V, value)
   432  	n.Valid = err == nil
   433  	return err
   434  }
   435  
   436  func (n Null[T]) Value() (driver.Value, error) {
   437  	if !n.Valid {
   438  		return nil, nil
   439  	}
   440  	v := any(n.V)
   441  	// See issue 69728.
   442  	if valuer, ok := v.(driver.Valuer); ok {
   443  		val, err := callValuerValue(valuer)
   444  		if err != nil {
   445  			return val, err
   446  		}
   447  		v = val
   448  	}
   449  	// See issue 69837.
   450  	return driver.DefaultParameterConverter.ConvertValue(v)
   451  }
   452  
   453  // Scanner is an interface used by [Rows.Scan].
   454  type Scanner interface {
   455  	// Scan assigns a value from a database driver.
   456  	//
   457  	// The src value will be of one of the following types:
   458  	//
   459  	//    int64
   460  	//    float64
   461  	//    bool
   462  	//    []byte
   463  	//    string
   464  	//    time.Time
   465  	//    nil - for NULL values
   466  	//
   467  	// An error should be returned if the value cannot be stored
   468  	// without loss of information.
   469  	//
   470  	// Reference types such as []byte are only valid until the next call to Scan
   471  	// and should not be retained. Their underlying memory is owned by the driver.
   472  	// If retention is necessary, copy their values before the next call to Scan.
   473  	Scan(src any) error
   474  }
   475  
   476  // Out may be used to retrieve OUTPUT value parameters from stored procedures.
   477  //
   478  // Not all drivers and databases support OUTPUT value parameters.
   479  //
   480  // Example usage:
   481  //
   482  //	var outArg string
   483  //	_, err := db.ExecContext(ctx, "ProcName", sql.Named("Arg1", sql.Out{Dest: &outArg}))
   484  type Out struct {
   485  	_NamedFieldsRequired struct{}
   486  
   487  	// Dest is a pointer to the value that will be set to the result of the
   488  	// stored procedure's OUTPUT parameter.
   489  	Dest any
   490  
   491  	// In is whether the parameter is an INOUT parameter. If so, the input value to the stored
   492  	// procedure is the dereferenced value of Dest's pointer, which is then replaced with
   493  	// the output value.
   494  	In bool
   495  }
   496  
   497  // ErrNoRows is returned by [Row.Scan] when [DB.QueryRow] doesn't return a
   498  // row. In such a case, QueryRow returns a placeholder [*Row] value that
   499  // defers this error until a Scan.
   500  var ErrNoRows = errors.New("sql: no rows in result set")
   501  
   502  // DB is a database handle representing a pool of zero or more
   503  // underlying connections. It's safe for concurrent use by multiple
   504  // goroutines.
   505  //
   506  // The sql package creates and frees connections automatically; it
   507  // also maintains a free pool of idle connections. If the database has
   508  // a concept of per-connection state, such state can be reliably observed
   509  // within a transaction ([Tx]) or connection ([Conn]). Once [DB.Begin] is called, the
   510  // returned [Tx] is bound to a single connection. Once [Tx.Commit] or
   511  // [Tx.Rollback] is called on the transaction, that transaction's
   512  // connection is returned to [DB]'s idle connection pool. The pool size
   513  // can be controlled with [DB.SetMaxIdleConns].
   514  type DB struct {
   515  	// Total time waited for new connections.
   516  	waitDuration atomic.Int64
   517  
   518  	connector driver.Connector
   519  	// numClosed is an atomic counter which represents a total number of
   520  	// closed connections. Stmt.openStmt checks it before cleaning closed
   521  	// connections in Stmt.css.
   522  	numClosed atomic.Uint64
   523  
   524  	mu           sync.Mutex    // protects following fields
   525  	freeConn     []*driverConn // free connections ordered by returnedAt oldest to newest
   526  	connRequests connRequestSet
   527  	numOpen      int // number of opened and pending open connections
   528  	// Used to signal the need for new connections
   529  	// a goroutine running connectionOpener() reads on this chan and
   530  	// maybeOpenNewConnections sends on the chan (one send per needed connection)
   531  	// It is closed during db.Close(). The close tells the connectionOpener
   532  	// goroutine to exit.
   533  	openerCh          chan struct{}
   534  	closed            bool
   535  	dep               map[finalCloser]depSet
   536  	lastPut           map[*driverConn]string // stacktrace of last conn's put; debug only
   537  	maxIdleCount      int                    // zero means defaultMaxIdleConns; negative means 0
   538  	maxOpen           int                    // <= 0 means unlimited
   539  	maxLifetime       time.Duration          // maximum amount of time a connection may be reused
   540  	maxIdleTime       time.Duration          // maximum amount of time a connection may be idle before being closed
   541  	cleanerCh         chan struct{}
   542  	waitCount         int64 // Total number of connections waited for.
   543  	maxIdleClosed     int64 // Total number of connections closed due to idle count.
   544  	maxIdleTimeClosed int64 // Total number of connections closed due to idle time.
   545  	maxLifetimeClosed int64 // Total number of connections closed due to max connection lifetime limit.
   546  
   547  	stop func() // stop cancels the connection opener.
   548  }
   549  
   550  // connReuseStrategy determines how (*DB).conn returns database connections.
   551  type connReuseStrategy uint8
   552  
   553  const (
   554  	// alwaysNewConn forces a new connection to the database.
   555  	alwaysNewConn connReuseStrategy = iota
   556  	// cachedOrNewConn returns a cached connection, if available, else waits
   557  	// for one to become available (if MaxOpenConns has been reached) or
   558  	// creates a new database connection.
   559  	cachedOrNewConn
   560  )
   561  
   562  // driverConn wraps a driver.Conn with a mutex, to
   563  // be held during all calls into the Conn. (including any calls onto
   564  // interfaces returned via that Conn, such as calls on Tx, Stmt,
   565  // Result, Rows)
   566  type driverConn struct {
   567  	db        *DB
   568  	createdAt time.Time
   569  
   570  	sync.Mutex  // guards following
   571  	ci          driver.Conn
   572  	needReset   bool // The connection session should be reset before use if true.
   573  	closed      bool
   574  	finalClosed bool // ci.Close has been called
   575  	openStmt    map[*driverStmt]bool
   576  
   577  	// guarded by db.mu
   578  	inUse      bool
   579  	dbmuClosed bool      // same as closed, but guarded by db.mu, for removeClosedStmtLocked
   580  	returnedAt time.Time // Time the connection was created or returned.
   581  	onPut      []func()  // code (with db.mu held) run when conn is next returned
   582  }
   583  
   584  func (dc *driverConn) releaseConn(err error) {
   585  	dc.db.putConn(dc, err, true)
   586  }
   587  
   588  func (dc *driverConn) removeOpenStmt(ds *driverStmt) {
   589  	dc.Lock()
   590  	defer dc.Unlock()
   591  	delete(dc.openStmt, ds)
   592  }
   593  
   594  func (dc *driverConn) expired(timeout time.Duration) bool {
   595  	if timeout <= 0 {
   596  		return false
   597  	}
   598  	return dc.createdAt.Add(timeout).Before(nowFunc())
   599  }
   600  
   601  // resetSession checks if the driver connection needs the
   602  // session to be reset and if required, resets it.
   603  func (dc *driverConn) resetSession(ctx context.Context) error {
   604  	dc.Lock()
   605  	defer dc.Unlock()
   606  
   607  	if !dc.needReset {
   608  		return nil
   609  	}
   610  	if cr, ok := dc.ci.(driver.SessionResetter); ok {
   611  		return cr.ResetSession(ctx)
   612  	}
   613  	return nil
   614  }
   615  
   616  // validateConnection checks if the connection is valid and can
   617  // still be used. It also marks the session for reset if required.
   618  func (dc *driverConn) validateConnection(needsReset bool) bool {
   619  	dc.Lock()
   620  	defer dc.Unlock()
   621  
   622  	if needsReset {
   623  		dc.needReset = true
   624  	}
   625  	if cv, ok := dc.ci.(driver.Validator); ok {
   626  		return cv.IsValid()
   627  	}
   628  	return true
   629  }
   630  
   631  // prepareLocked prepares the query on dc. When cg == nil the dc must keep track of
   632  // the prepared statements in a pool.
   633  func (dc *driverConn) prepareLocked(ctx context.Context, cg stmtConnGrabber, query string) (*driverStmt, error) {
   634  	si, err := ctxDriverPrepare(ctx, dc.ci, query)
   635  	if err != nil {
   636  		return nil, err
   637  	}
   638  	ds := &driverStmt{Locker: dc, si: si}
   639  
   640  	// No need to manage open statements if there is a single connection grabber.
   641  	if cg != nil {
   642  		return ds, nil
   643  	}
   644  
   645  	// Track each driverConn's open statements, so we can close them
   646  	// before closing the conn.
   647  	//
   648  	// Wrap all driver.Stmt is *driverStmt to ensure they are only closed once.
   649  	if dc.openStmt == nil {
   650  		dc.openStmt = make(map[*driverStmt]bool)
   651  	}
   652  	dc.openStmt[ds] = true
   653  	return ds, nil
   654  }
   655  
   656  // the dc.db's Mutex is held.
   657  func (dc *driverConn) closeDBLocked() func() error {
   658  	dc.Lock()
   659  	defer dc.Unlock()
   660  	if dc.closed {
   661  		return func() error { return errors.New("sql: duplicate driverConn close") }
   662  	}
   663  	dc.closed = true
   664  	return dc.db.removeDepLocked(dc, dc)
   665  }
   666  
   667  func (dc *driverConn) Close() error {
   668  	dc.Lock()
   669  	if dc.closed {
   670  		dc.Unlock()
   671  		return errors.New("sql: duplicate driverConn close")
   672  	}
   673  	dc.closed = true
   674  	dc.Unlock() // not defer; removeDep finalClose calls may need to lock
   675  
   676  	// And now updates that require holding dc.mu.Lock.
   677  	dc.db.mu.Lock()
   678  	dc.dbmuClosed = true
   679  	fn := dc.db.removeDepLocked(dc, dc)
   680  	dc.db.mu.Unlock()
   681  	return fn()
   682  }
   683  
   684  func (dc *driverConn) finalClose() error {
   685  	var err error
   686  
   687  	// Each *driverStmt has a lock to the dc. Copy the list out of the dc
   688  	// before calling close on each stmt.
   689  	var openStmt []*driverStmt
   690  	withLock(dc, func() {
   691  		openStmt = make([]*driverStmt, 0, len(dc.openStmt))
   692  		for ds := range dc.openStmt {
   693  			openStmt = append(openStmt, ds)
   694  		}
   695  		dc.openStmt = nil
   696  	})
   697  	for _, ds := range openStmt {
   698  		ds.Close()
   699  	}
   700  	withLock(dc, func() {
   701  		dc.finalClosed = true
   702  		err = dc.ci.Close()
   703  		dc.ci = nil
   704  	})
   705  
   706  	dc.db.mu.Lock()
   707  	dc.db.numOpen--
   708  	dc.db.maybeOpenNewConnections()
   709  	dc.db.mu.Unlock()
   710  
   711  	dc.db.numClosed.Add(1)
   712  	return err
   713  }
   714  
   715  // driverStmt associates a driver.Stmt with the
   716  // *driverConn from which it came, so the driverConn's lock can be
   717  // held during calls.
   718  type driverStmt struct {
   719  	sync.Locker // the *driverConn
   720  	si          driver.Stmt
   721  	closed      bool
   722  	closeErr    error // return value of previous Close call
   723  }
   724  
   725  // Close ensures driver.Stmt is only closed once and always returns the same
   726  // result.
   727  func (ds *driverStmt) Close() error {
   728  	ds.Lock()
   729  	defer ds.Unlock()
   730  	if ds.closed {
   731  		return ds.closeErr
   732  	}
   733  	ds.closed = true
   734  	ds.closeErr = ds.si.Close()
   735  	return ds.closeErr
   736  }
   737  
   738  // depSet is a finalCloser's outstanding dependencies
   739  type depSet map[any]bool // set of true bools
   740  
   741  // The finalCloser interface is used by (*DB).addDep and related
   742  // dependency reference counting.
   743  type finalCloser interface {
   744  	// finalClose is called when the reference count of an object
   745  	// goes to zero. (*DB).mu is not held while calling it.
   746  	finalClose() error
   747  }
   748  
   749  // addDep notes that x now depends on dep, and x's finalClose won't be
   750  // called until all of x's dependencies are removed with removeDep.
   751  func (db *DB) addDep(x finalCloser, dep any) {
   752  	db.mu.Lock()
   753  	defer db.mu.Unlock()
   754  	db.addDepLocked(x, dep)
   755  }
   756  
   757  func (db *DB) addDepLocked(x finalCloser, dep any) {
   758  	if db.dep == nil {
   759  		db.dep = make(map[finalCloser]depSet)
   760  	}
   761  	xdep := db.dep[x]
   762  	if xdep == nil {
   763  		xdep = make(depSet)
   764  		db.dep[x] = xdep
   765  	}
   766  	xdep[dep] = true
   767  }
   768  
   769  // removeDep notes that x no longer depends on dep.
   770  // If x still has dependencies, nil is returned.
   771  // If x no longer has any dependencies, its finalClose method will be
   772  // called and its error value will be returned.
   773  func (db *DB) removeDep(x finalCloser, dep any) error {
   774  	db.mu.Lock()
   775  	fn := db.removeDepLocked(x, dep)
   776  	db.mu.Unlock()
   777  	return fn()
   778  }
   779  
   780  func (db *DB) removeDepLocked(x finalCloser, dep any) func() error {
   781  	xdep, ok := db.dep[x]
   782  	if !ok {
   783  		panic(fmt.Sprintf("unpaired removeDep: no deps for %T", x))
   784  	}
   785  
   786  	l0 := len(xdep)
   787  	delete(xdep, dep)
   788  
   789  	switch len(xdep) {
   790  	case l0:
   791  		// Nothing removed. Shouldn't happen.
   792  		panic(fmt.Sprintf("unpaired removeDep: no %T dep on %T", dep, x))
   793  	case 0:
   794  		// No more dependencies.
   795  		delete(db.dep, x)
   796  		return x.finalClose
   797  	default:
   798  		// Dependencies remain.
   799  		return func() error { return nil }
   800  	}
   801  }
   802  
   803  // This is the size of the connectionOpener request chan (DB.openerCh).
   804  // This value should be larger than the maximum typical value
   805  // used for DB.maxOpen. If maxOpen is significantly larger than
   806  // connectionRequestQueueSize then it is possible for ALL calls into the *DB
   807  // to block until the connectionOpener can satisfy the backlog of requests.
   808  var connectionRequestQueueSize = 1000000
   809  
   810  type dsnConnector struct {
   811  	dsn    string
   812  	driver driver.Driver
   813  }
   814  
   815  func (t dsnConnector) Connect(_ context.Context) (driver.Conn, error) {
   816  	return t.driver.Open(t.dsn)
   817  }
   818  
   819  func (t dsnConnector) Driver() driver.Driver {
   820  	return t.driver
   821  }
   822  
   823  // OpenDB opens a database using a [driver.Connector], allowing drivers to
   824  // bypass a string based data source name.
   825  //
   826  // Most users will open a database via a driver-specific connection
   827  // helper function that returns a [*DB]. No database drivers are included
   828  // in the Go standard library. See https://golang.org/s/sqldrivers for
   829  // a list of third-party drivers.
   830  //
   831  // OpenDB may just validate its arguments without creating a connection
   832  // to the database. To verify that the data source name is valid, call
   833  // [DB.Ping].
   834  //
   835  // The returned [DB] is safe for concurrent use by multiple goroutines
   836  // and maintains its own pool of idle connections. Thus, the OpenDB
   837  // function should be called just once. It is rarely necessary to
   838  // close a [DB].
   839  func OpenDB(c driver.Connector) *DB {
   840  	ctx, cancel := context.WithCancel(context.Background())
   841  	db := &DB{
   842  		connector: c,
   843  		openerCh:  make(chan struct{}, connectionRequestQueueSize),
   844  		lastPut:   make(map[*driverConn]string),
   845  		stop:      cancel,
   846  	}
   847  
   848  	go db.connectionOpener(ctx)
   849  
   850  	return db
   851  }
   852  
   853  // Open opens a database specified by its database driver name and a
   854  // driver-specific data source name, usually consisting of at least a
   855  // database name and connection information.
   856  //
   857  // Most users will open a database via a driver-specific connection
   858  // helper function that returns a [*DB]. No database drivers are included
   859  // in the Go standard library. See https://golang.org/s/sqldrivers for
   860  // a list of third-party drivers.
   861  //
   862  // Open may just validate its arguments without creating a connection
   863  // to the database. To verify that the data source name is valid, call
   864  // [DB.Ping].
   865  //
   866  // The returned [DB] is safe for concurrent use by multiple goroutines
   867  // and maintains its own pool of idle connections. Thus, the Open
   868  // function should be called just once. It is rarely necessary to
   869  // close a [DB].
   870  func Open(driverName, dataSourceName string) (*DB, error) {
   871  	driversMu.RLock()
   872  	driveri, ok := drivers[driverName]
   873  	driversMu.RUnlock()
   874  	if !ok {
   875  		return nil, fmt.Errorf("sql: unknown driver %q (forgotten import?)", driverName)
   876  	}
   877  
   878  	if driverCtx, ok := driveri.(driver.DriverContext); ok {
   879  		connector, err := driverCtx.OpenConnector(dataSourceName)
   880  		if err != nil {
   881  			return nil, err
   882  		}
   883  		return OpenDB(connector), nil
   884  	}
   885  
   886  	return OpenDB(dsnConnector{dsn: dataSourceName, driver: driveri}), nil
   887  }
   888  
   889  func (db *DB) pingDC(ctx context.Context, dc *driverConn, release func(error)) error {
   890  	var err error
   891  	if pinger, ok := dc.ci.(driver.Pinger); ok {
   892  		withLock(dc, func() {
   893  			err = pinger.Ping(ctx)
   894  		})
   895  	}
   896  	release(err)
   897  	return err
   898  }
   899  
   900  // PingContext verifies a connection to the database is still alive,
   901  // establishing a connection if necessary.
   902  func (db *DB) PingContext(ctx context.Context) error {
   903  	var dc *driverConn
   904  	var err error
   905  
   906  	err = db.retry(func(strategy connReuseStrategy) error {
   907  		dc, err = db.conn(ctx, strategy)
   908  		return err
   909  	})
   910  
   911  	if err != nil {
   912  		return err
   913  	}
   914  
   915  	return db.pingDC(ctx, dc, dc.releaseConn)
   916  }
   917  
   918  // Ping verifies a connection to the database is still alive,
   919  // establishing a connection if necessary.
   920  //
   921  // Ping uses [context.Background] internally; to specify the context, use
   922  // [DB.PingContext].
   923  func (db *DB) Ping() error {
   924  	return db.PingContext(context.Background())
   925  }
   926  
   927  // Close closes the database and prevents new queries from starting.
   928  // Close then waits for all queries that have started processing on the server
   929  // to finish.
   930  //
   931  // It is rare to Close a [DB], as the [DB] handle is meant to be
   932  // long-lived and shared between many goroutines.
   933  func (db *DB) Close() error {
   934  	db.mu.Lock()
   935  	if db.closed { // Make DB.Close idempotent
   936  		db.mu.Unlock()
   937  		return nil
   938  	}
   939  	if db.cleanerCh != nil {
   940  		close(db.cleanerCh)
   941  	}
   942  	var err error
   943  	fns := make([]func() error, 0, len(db.freeConn))
   944  	for _, dc := range db.freeConn {
   945  		fns = append(fns, dc.closeDBLocked())
   946  	}
   947  	db.freeConn = nil
   948  	db.closed = true
   949  	db.connRequests.CloseAndRemoveAll()
   950  	db.mu.Unlock()
   951  	for _, fn := range fns {
   952  		err1 := fn()
   953  		if err1 != nil {
   954  			err = err1
   955  		}
   956  	}
   957  	db.stop()
   958  	if c, ok := db.connector.(io.Closer); ok {
   959  		err1 := c.Close()
   960  		if err1 != nil {
   961  			err = err1
   962  		}
   963  	}
   964  	return err
   965  }
   966  
   967  const defaultMaxIdleConns = 2
   968  
   969  func (db *DB) maxIdleConnsLocked() int {
   970  	n := db.maxIdleCount
   971  	switch {
   972  	case n == 0:
   973  		// TODO(bradfitz): ask driver, if supported, for its default preference
   974  		return defaultMaxIdleConns
   975  	case n < 0:
   976  		return 0
   977  	default:
   978  		return n
   979  	}
   980  }
   981  
   982  func (db *DB) shortestIdleTimeLocked() time.Duration {
   983  	if db.maxIdleTime <= 0 {
   984  		return db.maxLifetime
   985  	}
   986  	if db.maxLifetime <= 0 {
   987  		return db.maxIdleTime
   988  	}
   989  	return min(db.maxIdleTime, db.maxLifetime)
   990  }
   991  
   992  // SetMaxIdleConns sets the maximum number of connections in the idle
   993  // connection pool.
   994  //
   995  // If MaxOpenConns is greater than 0 but less than the new MaxIdleConns,
   996  // then the new MaxIdleConns will be reduced to match the MaxOpenConns limit.
   997  //
   998  // If n <= 0, no idle connections are retained.
   999  //
  1000  // The default max idle connections is currently 2. This may change in
  1001  // a future release.
  1002  func (db *DB) SetMaxIdleConns(n int) {
  1003  	db.mu.Lock()
  1004  	if n > 0 {
  1005  		db.maxIdleCount = n
  1006  	} else {
  1007  		// No idle connections.
  1008  		db.maxIdleCount = -1
  1009  	}
  1010  	// Make sure maxIdle doesn't exceed maxOpen
  1011  	if db.maxOpen > 0 && db.maxIdleConnsLocked() > db.maxOpen {
  1012  		db.maxIdleCount = db.maxOpen
  1013  	}
  1014  	var closing []*driverConn
  1015  	idleCount := len(db.freeConn)
  1016  	maxIdle := db.maxIdleConnsLocked()
  1017  	if idleCount > maxIdle {
  1018  		closing = db.freeConn[maxIdle:]
  1019  		db.freeConn = db.freeConn[:maxIdle]
  1020  	}
  1021  	db.maxIdleClosed += int64(len(closing))
  1022  	db.mu.Unlock()
  1023  	for _, c := range closing {
  1024  		c.Close()
  1025  	}
  1026  }
  1027  
  1028  // SetMaxOpenConns sets the maximum number of open connections to the database.
  1029  //
  1030  // If MaxIdleConns is greater than 0 and the new MaxOpenConns is less than
  1031  // MaxIdleConns, then MaxIdleConns will be reduced to match the new
  1032  // MaxOpenConns limit.
  1033  //
  1034  // If n <= 0, then there is no limit on the number of open connections.
  1035  // The default is 0 (unlimited).
  1036  func (db *DB) SetMaxOpenConns(n int) {
  1037  	db.mu.Lock()
  1038  	db.maxOpen = n
  1039  	if n < 0 {
  1040  		db.maxOpen = 0
  1041  	}
  1042  	syncMaxIdle := db.maxOpen > 0 && db.maxIdleConnsLocked() > db.maxOpen
  1043  	db.mu.Unlock()
  1044  	if syncMaxIdle {
  1045  		db.SetMaxIdleConns(n)
  1046  	}
  1047  }
  1048  
  1049  // SetConnMaxLifetime sets the maximum amount of time a connection may be reused.
  1050  //
  1051  // Expired connections may be closed lazily before reuse.
  1052  //
  1053  // If d <= 0, connections are not closed due to a connection's age.
  1054  func (db *DB) SetConnMaxLifetime(d time.Duration) {
  1055  	if d < 0 {
  1056  		d = 0
  1057  	}
  1058  	db.mu.Lock()
  1059  	// Wake cleaner up when lifetime is shortened.
  1060  	if d > 0 && d < db.shortestIdleTimeLocked() && db.cleanerCh != nil {
  1061  		select {
  1062  		case db.cleanerCh <- struct{}{}:
  1063  		default:
  1064  		}
  1065  	}
  1066  	db.maxLifetime = d
  1067  	db.startCleanerLocked()
  1068  	db.mu.Unlock()
  1069  }
  1070  
  1071  // SetConnMaxIdleTime sets the maximum amount of time a connection may be idle.
  1072  //
  1073  // Expired connections may be closed lazily before reuse.
  1074  //
  1075  // If d <= 0, connections are not closed due to a connection's idle time.
  1076  func (db *DB) SetConnMaxIdleTime(d time.Duration) {
  1077  	if d < 0 {
  1078  		d = 0
  1079  	}
  1080  	db.mu.Lock()
  1081  	defer db.mu.Unlock()
  1082  
  1083  	// Wake cleaner up when idle time is shortened.
  1084  	if d > 0 && d < db.shortestIdleTimeLocked() && db.cleanerCh != nil {
  1085  		select {
  1086  		case db.cleanerCh <- struct{}{}:
  1087  		default:
  1088  		}
  1089  	}
  1090  	db.maxIdleTime = d
  1091  	db.startCleanerLocked()
  1092  }
  1093  
  1094  // startCleanerLocked starts connectionCleaner if needed.
  1095  func (db *DB) startCleanerLocked() {
  1096  	if (db.maxLifetime > 0 || db.maxIdleTime > 0) && db.numOpen > 0 && db.cleanerCh == nil {
  1097  		db.cleanerCh = make(chan struct{}, 1)
  1098  		go db.connectionCleaner(db.shortestIdleTimeLocked())
  1099  	}
  1100  }
  1101  
  1102  func (db *DB) connectionCleaner(d time.Duration) {
  1103  	const minInterval = time.Second
  1104  
  1105  	if d < minInterval {
  1106  		d = minInterval
  1107  	}
  1108  	t := time.NewTimer(d)
  1109  
  1110  	for {
  1111  		select {
  1112  		case <-t.C:
  1113  		case <-db.cleanerCh: // maxLifetime was changed or db was closed.
  1114  		}
  1115  
  1116  		db.mu.Lock()
  1117  
  1118  		d = db.shortestIdleTimeLocked()
  1119  		if db.closed || db.numOpen == 0 || d <= 0 {
  1120  			db.cleanerCh = nil
  1121  			db.mu.Unlock()
  1122  			return
  1123  		}
  1124  
  1125  		d, closing := db.connectionCleanerRunLocked(d)
  1126  		db.mu.Unlock()
  1127  		for _, c := range closing {
  1128  			c.Close()
  1129  		}
  1130  
  1131  		if d < minInterval {
  1132  			d = minInterval
  1133  		}
  1134  
  1135  		if !t.Stop() {
  1136  			select {
  1137  			case <-t.C:
  1138  			default:
  1139  			}
  1140  		}
  1141  		t.Reset(d)
  1142  	}
  1143  }
  1144  
  1145  // connectionCleanerRunLocked removes connections that should be closed from
  1146  // freeConn and returns them along side an updated duration to the next check
  1147  // if a quicker check is required to ensure connections are checked appropriately.
  1148  func (db *DB) connectionCleanerRunLocked(d time.Duration) (time.Duration, []*driverConn) {
  1149  	var idleClosing int64
  1150  	var closing []*driverConn
  1151  	if db.maxIdleTime > 0 {
  1152  		// As freeConn is ordered by returnedAt process
  1153  		// in reverse order to minimise the work needed.
  1154  		idleSince := nowFunc().Add(-db.maxIdleTime)
  1155  		last := len(db.freeConn) - 1
  1156  		for i := last; i >= 0; i-- {
  1157  			c := db.freeConn[i]
  1158  			if c.returnedAt.Before(idleSince) {
  1159  				i++
  1160  				closing = db.freeConn[:i:i]
  1161  				db.freeConn = db.freeConn[i:]
  1162  				idleClosing = int64(len(closing))
  1163  				db.maxIdleTimeClosed += idleClosing
  1164  				break
  1165  			}
  1166  		}
  1167  
  1168  		if len(db.freeConn) > 0 {
  1169  			c := db.freeConn[0]
  1170  			if d2 := c.returnedAt.Sub(idleSince); d2 < d {
  1171  				// Ensure idle connections are cleaned up as soon as
  1172  				// possible.
  1173  				d = d2
  1174  			}
  1175  		}
  1176  	}
  1177  
  1178  	if db.maxLifetime > 0 {
  1179  		expiredSince := nowFunc().Add(-db.maxLifetime)
  1180  		for i := 0; i < len(db.freeConn); i++ {
  1181  			c := db.freeConn[i]
  1182  			if c.createdAt.Before(expiredSince) {
  1183  				closing = append(closing, c)
  1184  
  1185  				last := len(db.freeConn) - 1
  1186  				// Use slow delete as order is required to ensure
  1187  				// connections are reused least idle time first.
  1188  				copy(db.freeConn[i:], db.freeConn[i+1:])
  1189  				db.freeConn[last] = nil
  1190  				db.freeConn = db.freeConn[:last]
  1191  				i--
  1192  			} else if d2 := c.createdAt.Sub(expiredSince); d2 < d {
  1193  				// Prevent connections sitting the freeConn when they
  1194  				// have expired by updating our next deadline d.
  1195  				d = d2
  1196  			}
  1197  		}
  1198  		db.maxLifetimeClosed += int64(len(closing)) - idleClosing
  1199  	}
  1200  
  1201  	return d, closing
  1202  }
  1203  
  1204  // DBStats contains database statistics.
  1205  type DBStats struct {
  1206  	MaxOpenConnections int // Maximum number of open connections to the database.
  1207  
  1208  	// Pool Status
  1209  	OpenConnections int // The number of established connections both in use and idle.
  1210  	InUse           int // The number of connections currently in use.
  1211  	Idle            int // The number of idle connections.
  1212  
  1213  	// Counters
  1214  	WaitCount         int64         // The total number of connections waited for.
  1215  	WaitDuration      time.Duration // The total time blocked waiting for a new connection.
  1216  	MaxIdleClosed     int64         // The total number of connections closed due to SetMaxIdleConns.
  1217  	MaxIdleTimeClosed int64         // The total number of connections closed due to SetConnMaxIdleTime.
  1218  	MaxLifetimeClosed int64         // The total number of connections closed due to SetConnMaxLifetime.
  1219  }
  1220  
  1221  // Stats returns database statistics.
  1222  func (db *DB) Stats() DBStats {
  1223  	wait := db.waitDuration.Load()
  1224  
  1225  	db.mu.Lock()
  1226  	defer db.mu.Unlock()
  1227  
  1228  	stats := DBStats{
  1229  		MaxOpenConnections: db.maxOpen,
  1230  
  1231  		Idle:            len(db.freeConn),
  1232  		OpenConnections: db.numOpen,
  1233  		InUse:           db.numOpen - len(db.freeConn),
  1234  
  1235  		WaitCount:         db.waitCount,
  1236  		WaitDuration:      time.Duration(wait),
  1237  		MaxIdleClosed:     db.maxIdleClosed,
  1238  		MaxIdleTimeClosed: db.maxIdleTimeClosed,
  1239  		MaxLifetimeClosed: db.maxLifetimeClosed,
  1240  	}
  1241  	return stats
  1242  }
  1243  
  1244  // Assumes db.mu is locked.
  1245  // If there are connRequests and the connection limit hasn't been reached,
  1246  // then tell the connectionOpener to open new connections.
  1247  func (db *DB) maybeOpenNewConnections() {
  1248  	numRequests := db.connRequests.Len()
  1249  	if db.maxOpen > 0 {
  1250  		numCanOpen := db.maxOpen - db.numOpen
  1251  		if numRequests > numCanOpen {
  1252  			numRequests = numCanOpen
  1253  		}
  1254  	}
  1255  	for numRequests > 0 {
  1256  		db.numOpen++ // optimistically
  1257  		numRequests--
  1258  		if db.closed {
  1259  			return
  1260  		}
  1261  		db.openerCh <- struct{}{}
  1262  	}
  1263  }
  1264  
  1265  // Runs in a separate goroutine, opens new connections when requested.
  1266  func (db *DB) connectionOpener(ctx context.Context) {
  1267  	for {
  1268  		select {
  1269  		case <-ctx.Done():
  1270  			return
  1271  		case <-db.openerCh:
  1272  			db.openNewConnection(ctx)
  1273  		}
  1274  	}
  1275  }
  1276  
  1277  // Open one new connection
  1278  func (db *DB) openNewConnection(ctx context.Context) {
  1279  	// maybeOpenNewConnections has already executed db.numOpen++ before it sent
  1280  	// on db.openerCh. This function must execute db.numOpen-- if the
  1281  	// connection fails or is closed before returning.
  1282  	ci, err := db.connector.Connect(ctx)
  1283  	db.mu.Lock()
  1284  	defer db.mu.Unlock()
  1285  	if db.closed {
  1286  		if err == nil {
  1287  			ci.Close()
  1288  		}
  1289  		db.numOpen--
  1290  		return
  1291  	}
  1292  	if err != nil {
  1293  		db.numOpen--
  1294  		db.putConnDBLocked(nil, err)
  1295  		db.maybeOpenNewConnections()
  1296  		return
  1297  	}
  1298  	dc := &driverConn{
  1299  		db:         db,
  1300  		createdAt:  nowFunc(),
  1301  		returnedAt: nowFunc(),
  1302  		ci:         ci,
  1303  	}
  1304  	if db.putConnDBLocked(dc, err) {
  1305  		db.addDepLocked(dc, dc)
  1306  	} else {
  1307  		db.numOpen--
  1308  		ci.Close()
  1309  	}
  1310  }
  1311  
  1312  // connRequest represents one request for a new connection
  1313  // When there are no idle connections available, DB.conn will create
  1314  // a new connRequest and put it on the db.connRequests list.
  1315  type connRequest struct {
  1316  	conn *driverConn
  1317  	err  error
  1318  }
  1319  
  1320  var errDBClosed = errors.New("sql: database is closed")
  1321  
  1322  // conn returns a newly-opened or cached *driverConn.
  1323  func (db *DB) conn(ctx context.Context, strategy connReuseStrategy) (*driverConn, error) {
  1324  	db.mu.Lock()
  1325  	if db.closed {
  1326  		db.mu.Unlock()
  1327  		return nil, errDBClosed
  1328  	}
  1329  	// Check if the context is expired.
  1330  	select {
  1331  	default:
  1332  	case <-ctx.Done():
  1333  		db.mu.Unlock()
  1334  		return nil, ctx.Err()
  1335  	}
  1336  	lifetime := db.maxLifetime
  1337  
  1338  	// Prefer a free connection, if possible.
  1339  	last := len(db.freeConn) - 1
  1340  	if strategy == cachedOrNewConn && last >= 0 {
  1341  		// Reuse the lowest idle time connection so we can close
  1342  		// connections which remain idle as soon as possible.
  1343  		conn := db.freeConn[last]
  1344  		db.freeConn = db.freeConn[:last]
  1345  		conn.inUse = true
  1346  		if conn.expired(lifetime) {
  1347  			db.maxLifetimeClosed++
  1348  			db.mu.Unlock()
  1349  			conn.Close()
  1350  			return nil, driver.ErrBadConn
  1351  		}
  1352  		db.mu.Unlock()
  1353  
  1354  		// Reset the session if required.
  1355  		if err := conn.resetSession(ctx); errors.Is(err, driver.ErrBadConn) {
  1356  			conn.Close()
  1357  			return nil, err
  1358  		}
  1359  
  1360  		return conn, nil
  1361  	}
  1362  
  1363  	// Out of free connections or we were asked not to use one. If we're not
  1364  	// allowed to open any more connections, make a request and wait.
  1365  	if db.maxOpen > 0 && db.numOpen >= db.maxOpen {
  1366  		// Make the connRequest channel. It's buffered so that the
  1367  		// connectionOpener doesn't block while waiting for the req to be read.
  1368  		req := make(chan connRequest, 1)
  1369  		delHandle := db.connRequests.Add(req)
  1370  		db.waitCount++
  1371  		db.mu.Unlock()
  1372  
  1373  		waitStart := nowFunc()
  1374  
  1375  		// Timeout the connection request with the context.
  1376  		select {
  1377  		case <-ctx.Done():
  1378  			// Remove the connection request and ensure no value has been sent
  1379  			// on it after removing.
  1380  			db.mu.Lock()
  1381  			deleted := db.connRequests.Delete(delHandle)
  1382  			db.mu.Unlock()
  1383  
  1384  			db.waitDuration.Add(int64(time.Since(waitStart)))
  1385  
  1386  			// If we failed to delete it, that means either the DB was closed or
  1387  			// something else grabbed it and is about to send on it.
  1388  			if !deleted {
  1389  				// TODO(bradfitz): rather than this best effort select, we
  1390  				// should probably start a goroutine to read from req. This best
  1391  				// effort select existed before the change to check 'deleted'.
  1392  				// But if we know for sure it wasn't deleted and a sender is
  1393  				// outstanding, we should probably block on req (in a new
  1394  				// goroutine) to get the connection back.
  1395  				select {
  1396  				default:
  1397  				case ret, ok := <-req:
  1398  					if ok && ret.conn != nil {
  1399  						db.putConn(ret.conn, ret.err, false)
  1400  					}
  1401  				}
  1402  			}
  1403  			return nil, ctx.Err()
  1404  		case ret, ok := <-req:
  1405  			db.waitDuration.Add(int64(time.Since(waitStart)))
  1406  
  1407  			if !ok {
  1408  				return nil, errDBClosed
  1409  			}
  1410  			// Only check if the connection is expired if the strategy is cachedOrNewConns.
  1411  			// If we require a new connection, just re-use the connection without looking
  1412  			// at the expiry time. If it is expired, it will be checked when it is placed
  1413  			// back into the connection pool.
  1414  			// This prioritizes giving a valid connection to a client over the exact connection
  1415  			// lifetime, which could expire exactly after this point anyway.
  1416  			if strategy == cachedOrNewConn && ret.err == nil && ret.conn.expired(lifetime) {
  1417  				db.mu.Lock()
  1418  				db.maxLifetimeClosed++
  1419  				db.mu.Unlock()
  1420  				ret.conn.Close()
  1421  				return nil, driver.ErrBadConn
  1422  			}
  1423  			if ret.conn == nil {
  1424  				return nil, ret.err
  1425  			}
  1426  
  1427  			// Reset the session if required.
  1428  			if err := ret.conn.resetSession(ctx); errors.Is(err, driver.ErrBadConn) {
  1429  				ret.conn.Close()
  1430  				return nil, err
  1431  			}
  1432  			return ret.conn, ret.err
  1433  		}
  1434  	}
  1435  
  1436  	db.numOpen++ // optimistically
  1437  	db.mu.Unlock()
  1438  	ci, err := db.connector.Connect(ctx)
  1439  	if err != nil {
  1440  		db.mu.Lock()
  1441  		db.numOpen-- // correct for earlier optimism
  1442  		db.maybeOpenNewConnections()
  1443  		db.mu.Unlock()
  1444  		return nil, err
  1445  	}
  1446  	db.mu.Lock()
  1447  	dc := &driverConn{
  1448  		db:         db,
  1449  		createdAt:  nowFunc(),
  1450  		returnedAt: nowFunc(),
  1451  		ci:         ci,
  1452  		inUse:      true,
  1453  	}
  1454  	db.addDepLocked(dc, dc)
  1455  	db.mu.Unlock()
  1456  	return dc, nil
  1457  }
  1458  
  1459  // putConnHook is a hook for testing.
  1460  var putConnHook func(*DB, *driverConn)
  1461  
  1462  // noteUnusedDriverStatement notes that ds is no longer used and should
  1463  // be closed whenever possible (when c is next not in use), unless c is
  1464  // already closed.
  1465  func (db *DB) noteUnusedDriverStatement(c *driverConn, ds *driverStmt) {
  1466  	db.mu.Lock()
  1467  	defer db.mu.Unlock()
  1468  	if c.inUse {
  1469  		c.onPut = append(c.onPut, func() {
  1470  			ds.Close()
  1471  		})
  1472  	} else {
  1473  		c.Lock()
  1474  		fc := c.finalClosed
  1475  		c.Unlock()
  1476  		if !fc {
  1477  			ds.Close()
  1478  		}
  1479  	}
  1480  }
  1481  
  1482  // debugGetPut determines whether getConn & putConn calls' stack traces
  1483  // are returned for more verbose crashes.
  1484  const debugGetPut = false
  1485  
  1486  // putConn adds a connection to the db's free pool.
  1487  // err is optionally the last error that occurred on this connection.
  1488  func (db *DB) putConn(dc *driverConn, err error, resetSession bool) {
  1489  	if !errors.Is(err, driver.ErrBadConn) {
  1490  		if !dc.validateConnection(resetSession) {
  1491  			err = driver.ErrBadConn
  1492  		}
  1493  	}
  1494  	db.mu.Lock()
  1495  	if !dc.inUse {
  1496  		db.mu.Unlock()
  1497  		if debugGetPut {
  1498  			fmt.Printf("putConn(%v) DUPLICATE was: %s\n\nPREVIOUS was: %s", dc, stack(), db.lastPut[dc])
  1499  		}
  1500  		panic("sql: connection returned that was never out")
  1501  	}
  1502  
  1503  	if !errors.Is(err, driver.ErrBadConn) && dc.expired(db.maxLifetime) {
  1504  		db.maxLifetimeClosed++
  1505  		err = driver.ErrBadConn
  1506  	}
  1507  	if debugGetPut {
  1508  		db.lastPut[dc] = stack()
  1509  	}
  1510  	dc.inUse = false
  1511  	dc.returnedAt = nowFunc()
  1512  
  1513  	for _, fn := range dc.onPut {
  1514  		fn()
  1515  	}
  1516  	dc.onPut = nil
  1517  
  1518  	if errors.Is(err, driver.ErrBadConn) {
  1519  		// Don't reuse bad connections.
  1520  		// Since the conn is considered bad and is being discarded, treat it
  1521  		// as closed. Don't decrement the open count here, finalClose will
  1522  		// take care of that.
  1523  		db.maybeOpenNewConnections()
  1524  		db.mu.Unlock()
  1525  		dc.Close()
  1526  		return
  1527  	}
  1528  	if putConnHook != nil {
  1529  		putConnHook(db, dc)
  1530  	}
  1531  	added := db.putConnDBLocked(dc, nil)
  1532  	db.mu.Unlock()
  1533  
  1534  	if !added {
  1535  		dc.Close()
  1536  		return
  1537  	}
  1538  }
  1539  
  1540  // Satisfy a connRequest or put the driverConn in the idle pool and return true
  1541  // or return false.
  1542  // putConnDBLocked will satisfy a connRequest if there is one, or it will
  1543  // return the *driverConn to the freeConn list if err == nil and the idle
  1544  // connection limit will not be exceeded.
  1545  // If err != nil, the value of dc is ignored.
  1546  // If err == nil, then dc must not equal nil.
  1547  // If a connRequest was fulfilled or the *driverConn was placed in the
  1548  // freeConn list, then true is returned, otherwise false is returned.
  1549  func (db *DB) putConnDBLocked(dc *driverConn, err error) bool {
  1550  	if db.closed {
  1551  		return false
  1552  	}
  1553  	if db.maxOpen > 0 && db.numOpen > db.maxOpen {
  1554  		return false
  1555  	}
  1556  	if req, ok := db.connRequests.TakeRandom(); ok {
  1557  		if err == nil {
  1558  			dc.inUse = true
  1559  		}
  1560  		req <- connRequest{
  1561  			conn: dc,
  1562  			err:  err,
  1563  		}
  1564  		return true
  1565  	} else if err == nil && !db.closed {
  1566  		if db.maxIdleConnsLocked() > len(db.freeConn) {
  1567  			db.freeConn = append(db.freeConn, dc)
  1568  			db.startCleanerLocked()
  1569  			return true
  1570  		}
  1571  		db.maxIdleClosed++
  1572  	}
  1573  	return false
  1574  }
  1575  
  1576  // maxBadConnRetries is the number of maximum retries if the driver returns
  1577  // driver.ErrBadConn to signal a broken connection before forcing a new
  1578  // connection to be opened.
  1579  const maxBadConnRetries = 2
  1580  
  1581  func (db *DB) retry(fn func(strategy connReuseStrategy) error) error {
  1582  	for i := int64(0); i < maxBadConnRetries; i++ {
  1583  		err := fn(cachedOrNewConn)
  1584  		// retry if err is driver.ErrBadConn
  1585  		if err == nil || !errors.Is(err, driver.ErrBadConn) {
  1586  			return err
  1587  		}
  1588  	}
  1589  
  1590  	return fn(alwaysNewConn)
  1591  }
  1592  
  1593  // PrepareContext creates a prepared statement for later queries or executions.
  1594  // Multiple queries or executions may be run concurrently from the
  1595  // returned statement.
  1596  // The caller must call the statement's [*Stmt.Close] method
  1597  // when the statement is no longer needed.
  1598  //
  1599  // The provided context is used for the preparation of the statement, not for the
  1600  // execution of the statement.
  1601  func (db *DB) PrepareContext(ctx context.Context, query string) (*Stmt, error) {
  1602  	var stmt *Stmt
  1603  	var err error
  1604  
  1605  	err = db.retry(func(strategy connReuseStrategy) error {
  1606  		stmt, err = db.prepare(ctx, query, strategy)
  1607  		return err
  1608  	})
  1609  
  1610  	return stmt, err
  1611  }
  1612  
  1613  // Prepare creates a prepared statement for later queries or executions.
  1614  // Multiple queries or executions may be run concurrently from the
  1615  // returned statement.
  1616  // The caller must call the statement's [*Stmt.Close] method
  1617  // when the statement is no longer needed.
  1618  //
  1619  // Prepare uses [context.Background] internally; to specify the context, use
  1620  // [DB.PrepareContext].
  1621  func (db *DB) Prepare(query string) (*Stmt, error) {
  1622  	return db.PrepareContext(context.Background(), query)
  1623  }
  1624  
  1625  func (db *DB) prepare(ctx context.Context, query string, strategy connReuseStrategy) (*Stmt, error) {
  1626  	// TODO: check if db.driver supports an optional
  1627  	// driver.Preparer interface and call that instead, if so,
  1628  	// otherwise we make a prepared statement that's bound
  1629  	// to a connection, and to execute this prepared statement
  1630  	// we either need to use this connection (if it's free), else
  1631  	// get a new connection + re-prepare + execute on that one.
  1632  	dc, err := db.conn(ctx, strategy)
  1633  	if err != nil {
  1634  		return nil, err
  1635  	}
  1636  	return db.prepareDC(ctx, dc, dc.releaseConn, nil, query)
  1637  }
  1638  
  1639  // prepareDC prepares a query on the driverConn and calls release before
  1640  // returning. When cg == nil it implies that a connection pool is used, and
  1641  // when cg != nil only a single driver connection is used.
  1642  func (db *DB) prepareDC(ctx context.Context, dc *driverConn, release func(error), cg stmtConnGrabber, query string) (*Stmt, error) {
  1643  	var ds *driverStmt
  1644  	var err error
  1645  	defer func() {
  1646  		release(err)
  1647  	}()
  1648  	withLock(dc, func() {
  1649  		ds, err = dc.prepareLocked(ctx, cg, query)
  1650  	})
  1651  	if err != nil {
  1652  		return nil, err
  1653  	}
  1654  	stmt := &Stmt{
  1655  		db:    db,
  1656  		query: query,
  1657  		cg:    cg,
  1658  		cgds:  ds,
  1659  	}
  1660  
  1661  	// When cg == nil this statement will need to keep track of various
  1662  	// connections they are prepared on and record the stmt dependency on
  1663  	// the DB.
  1664  	if cg == nil {
  1665  		stmt.css = []connStmt{{dc, ds}}
  1666  		stmt.lastNumClosed = db.numClosed.Load()
  1667  		db.addDep(stmt, stmt)
  1668  	}
  1669  	return stmt, nil
  1670  }
  1671  
  1672  // ExecContext executes a query without returning any rows.
  1673  // The args are for any placeholder parameters in the query.
  1674  func (db *DB) ExecContext(ctx context.Context, query string, args ...any) (Result, error) {
  1675  	var res Result
  1676  	var err error
  1677  
  1678  	err = db.retry(func(strategy connReuseStrategy) error {
  1679  		res, err = db.exec(ctx, query, args, strategy)
  1680  		return err
  1681  	})
  1682  
  1683  	return res, err
  1684  }
  1685  
  1686  // Exec executes a query without returning any rows.
  1687  // The args are for any placeholder parameters in the query.
  1688  //
  1689  // Exec uses [context.Background] internally; to specify the context, use
  1690  // [DB.ExecContext].
  1691  func (db *DB) Exec(query string, args ...any) (Result, error) {
  1692  	return db.ExecContext(context.Background(), query, args...)
  1693  }
  1694  
  1695  func (db *DB) exec(ctx context.Context, query string, args []any, strategy connReuseStrategy) (Result, error) {
  1696  	dc, err := db.conn(ctx, strategy)
  1697  	if err != nil {
  1698  		return nil, err
  1699  	}
  1700  	return db.execDC(ctx, dc, dc.releaseConn, query, args)
  1701  }
  1702  
  1703  func (db *DB) execDC(ctx context.Context, dc *driverConn, release func(error), query string, args []any) (res Result, err error) {
  1704  	defer func() {
  1705  		release(err)
  1706  	}()
  1707  	execerCtx, ok := dc.ci.(driver.ExecerContext)
  1708  	var execer driver.Execer
  1709  	if !ok {
  1710  		execer, ok = dc.ci.(driver.Execer)
  1711  	}
  1712  	if ok {
  1713  		var nvdargs []driver.NamedValue
  1714  		var resi driver.Result
  1715  		withLock(dc, func() {
  1716  			nvdargs, err = driverArgsConnLocked(dc.ci, nil, args)
  1717  			if err != nil {
  1718  				return
  1719  			}
  1720  			resi, err = ctxDriverExec(ctx, execerCtx, execer, query, nvdargs)
  1721  		})
  1722  		if err != driver.ErrSkip {
  1723  			if err != nil {
  1724  				return nil, err
  1725  			}
  1726  			return driverResult{dc, resi}, nil
  1727  		}
  1728  	}
  1729  
  1730  	var si driver.Stmt
  1731  	withLock(dc, func() {
  1732  		si, err = ctxDriverPrepare(ctx, dc.ci, query)
  1733  	})
  1734  	if err != nil {
  1735  		return nil, err
  1736  	}
  1737  	ds := &driverStmt{Locker: dc, si: si}
  1738  	defer ds.Close()
  1739  	return resultFromStatement(ctx, dc.ci, ds, args...)
  1740  }
  1741  
  1742  // QueryContext executes a query that returns rows, typically a SELECT.
  1743  // The args are for any placeholder parameters in the query.
  1744  func (db *DB) QueryContext(ctx context.Context, query string, args ...any) (*Rows, error) {
  1745  	var rows *Rows
  1746  	var err error
  1747  
  1748  	err = db.retry(func(strategy connReuseStrategy) error {
  1749  		rows, err = db.query(ctx, query, args, strategy)
  1750  		return err
  1751  	})
  1752  
  1753  	return rows, err
  1754  }
  1755  
  1756  // Query executes a query that returns rows, typically a SELECT.
  1757  // The args are for any placeholder parameters in the query.
  1758  //
  1759  // Query uses [context.Background] internally; to specify the context, use
  1760  // [DB.QueryContext].
  1761  func (db *DB) Query(query string, args ...any) (*Rows, error) {
  1762  	return db.QueryContext(context.Background(), query, args...)
  1763  }
  1764  
  1765  func (db *DB) query(ctx context.Context, query string, args []any, strategy connReuseStrategy) (*Rows, error) {
  1766  	dc, err := db.conn(ctx, strategy)
  1767  	if err != nil {
  1768  		return nil, err
  1769  	}
  1770  
  1771  	return db.queryDC(ctx, nil, dc, dc.releaseConn, query, args)
  1772  }
  1773  
  1774  // queryDC executes a query on the given connection.
  1775  // The connection gets released by the releaseConn function.
  1776  // The ctx context is from a query method and the txctx context is from an
  1777  // optional transaction context.
  1778  func (db *DB) queryDC(ctx, txctx context.Context, dc *driverConn, releaseConn func(error), query string, args []any) (*Rows, error) {
  1779  	queryerCtx, ok := dc.ci.(driver.QueryerContext)
  1780  	var queryer driver.Queryer
  1781  	if !ok {
  1782  		queryer, ok = dc.ci.(driver.Queryer)
  1783  	}
  1784  	if ok {
  1785  		var nvdargs []driver.NamedValue
  1786  		var rowsi driver.Rows
  1787  		var err error
  1788  		withLock(dc, func() {
  1789  			nvdargs, err = driverArgsConnLocked(dc.ci, nil, args)
  1790  			if err != nil {
  1791  				return
  1792  			}
  1793  			rowsi, err = ctxDriverQuery(ctx, queryerCtx, queryer, query, nvdargs)
  1794  		})
  1795  		if err != driver.ErrSkip {
  1796  			if err != nil {
  1797  				releaseConn(err)
  1798  				return nil, err
  1799  			}
  1800  			// Note: ownership of dc passes to the *Rows, to be freed
  1801  			// with releaseConn.
  1802  			rows := &Rows{
  1803  				dc:          dc,
  1804  				releaseConn: releaseConn,
  1805  				rowsi:       rowsi,
  1806  			}
  1807  			rows.initContextClose(ctx, txctx)
  1808  			return rows, nil
  1809  		}
  1810  	}
  1811  
  1812  	var si driver.Stmt
  1813  	var err error
  1814  	withLock(dc, func() {
  1815  		si, err = ctxDriverPrepare(ctx, dc.ci, query)
  1816  	})
  1817  	if err != nil {
  1818  		releaseConn(err)
  1819  		return nil, err
  1820  	}
  1821  
  1822  	ds := &driverStmt{Locker: dc, si: si}
  1823  	rowsi, err := rowsiFromStatement(ctx, dc.ci, ds, args...)
  1824  	if err != nil {
  1825  		ds.Close()
  1826  		releaseConn(err)
  1827  		return nil, err
  1828  	}
  1829  
  1830  	// Note: ownership of ci passes to the *Rows, to be freed
  1831  	// with releaseConn.
  1832  	rows := &Rows{
  1833  		dc:          dc,
  1834  		releaseConn: releaseConn,
  1835  		rowsi:       rowsi,
  1836  		closeStmt:   ds,
  1837  	}
  1838  	rows.initContextClose(ctx, txctx)
  1839  	return rows, nil
  1840  }
  1841  
  1842  // QueryRowContext executes a query that is expected to return at most one row.
  1843  // QueryRowContext always returns a non-nil value. Errors are deferred until
  1844  // [Row]'s Scan method is called.
  1845  // If the query selects no rows, the [*Row.Scan] will return [ErrNoRows].
  1846  // Otherwise, [*Row.Scan] scans the first selected row and discards
  1847  // the rest.
  1848  func (db *DB) QueryRowContext(ctx context.Context, query string, args ...any) *Row {
  1849  	rows, err := db.QueryContext(ctx, query, args...)
  1850  	return &Row{rows: rows, err: err}
  1851  }
  1852  
  1853  // QueryRow executes a query that is expected to return at most one row.
  1854  // QueryRow always returns a non-nil value. Errors are deferred until
  1855  // [Row]'s Scan method is called.
  1856  // If the query selects no rows, the [*Row.Scan] will return [ErrNoRows].
  1857  // Otherwise, [*Row.Scan] scans the first selected row and discards
  1858  // the rest.
  1859  //
  1860  // QueryRow uses [context.Background] internally; to specify the context, use
  1861  // [DB.QueryRowContext].
  1862  func (db *DB) QueryRow(query string, args ...any) *Row {
  1863  	return db.QueryRowContext(context.Background(), query, args...)
  1864  }
  1865  
  1866  // BeginTx starts a transaction.
  1867  //
  1868  // The provided context is used until the transaction is committed or rolled back.
  1869  // If the context is canceled, the sql package will roll back
  1870  // the transaction. [Tx.Commit] will return an error if the context provided to
  1871  // BeginTx is canceled.
  1872  //
  1873  // The provided [TxOptions] is optional and may be nil if defaults should be used.
  1874  // If a non-default isolation level is used that the driver doesn't support,
  1875  // an error will be returned.
  1876  func (db *DB) BeginTx(ctx context.Context, opts *TxOptions) (*Tx, error) {
  1877  	var tx *Tx
  1878  	var err error
  1879  
  1880  	err = db.retry(func(strategy connReuseStrategy) error {
  1881  		tx, err = db.begin(ctx, opts, strategy)
  1882  		return err
  1883  	})
  1884  
  1885  	return tx, err
  1886  }
  1887  
  1888  // Begin starts a transaction. The default isolation level is dependent on
  1889  // the driver.
  1890  //
  1891  // Begin uses [context.Background] internally; to specify the context, use
  1892  // [DB.BeginTx].
  1893  func (db *DB) Begin() (*Tx, error) {
  1894  	return db.BeginTx(context.Background(), nil)
  1895  }
  1896  
  1897  func (db *DB) begin(ctx context.Context, opts *TxOptions, strategy connReuseStrategy) (tx *Tx, err error) {
  1898  	dc, err := db.conn(ctx, strategy)
  1899  	if err != nil {
  1900  		return nil, err
  1901  	}
  1902  	return db.beginDC(ctx, dc, dc.releaseConn, opts)
  1903  }
  1904  
  1905  // beginDC starts a transaction. The provided dc must be valid and ready to use.
  1906  func (db *DB) beginDC(ctx context.Context, dc *driverConn, release func(error), opts *TxOptions) (tx *Tx, err error) {
  1907  	var txi driver.Tx
  1908  	keepConnOnRollback := false
  1909  	withLock(dc, func() {
  1910  		_, hasSessionResetter := dc.ci.(driver.SessionResetter)
  1911  		_, hasConnectionValidator := dc.ci.(driver.Validator)
  1912  		keepConnOnRollback = hasSessionResetter && hasConnectionValidator
  1913  		txi, err = ctxDriverBegin(ctx, opts, dc.ci)
  1914  	})
  1915  	if err != nil {
  1916  		release(err)
  1917  		return nil, err
  1918  	}
  1919  
  1920  	// Schedule the transaction to rollback when the context is canceled.
  1921  	// The cancel function in Tx will be called after done is set to true.
  1922  	ctx, cancel := context.WithCancel(ctx)
  1923  	tx = &Tx{
  1924  		db:                 db,
  1925  		dc:                 dc,
  1926  		releaseConn:        release,
  1927  		txi:                txi,
  1928  		cancel:             cancel,
  1929  		keepConnOnRollback: keepConnOnRollback,
  1930  		ctx:                ctx,
  1931  	}
  1932  	go tx.awaitDone()
  1933  	return tx, nil
  1934  }
  1935  
  1936  // Driver returns the database's underlying driver.
  1937  func (db *DB) Driver() driver.Driver {
  1938  	return db.connector.Driver()
  1939  }
  1940  
  1941  // ErrConnDone is returned by any operation that is performed on a connection
  1942  // that has already been returned to the connection pool.
  1943  var ErrConnDone = errors.New("sql: connection is already closed")
  1944  
  1945  // Conn returns a single connection by either opening a new connection
  1946  // or returning an existing connection from the connection pool. Conn will
  1947  // block until either a connection is returned or ctx is canceled.
  1948  // Queries run on the same Conn will be run in the same database session.
  1949  //
  1950  // Every Conn must be returned to the database pool after use by
  1951  // calling [Conn.Close].
  1952  func (db *DB) Conn(ctx context.Context) (*Conn, error) {
  1953  	var dc *driverConn
  1954  	var err error
  1955  
  1956  	err = db.retry(func(strategy connReuseStrategy) error {
  1957  		dc, err = db.conn(ctx, strategy)
  1958  		return err
  1959  	})
  1960  
  1961  	if err != nil {
  1962  		return nil, err
  1963  	}
  1964  
  1965  	conn := &Conn{
  1966  		db: db,
  1967  		dc: dc,
  1968  	}
  1969  	return conn, nil
  1970  }
  1971  
  1972  type releaseConn func(error)
  1973  
  1974  // Conn represents a single database connection rather than a pool of database
  1975  // connections. Prefer running queries from [DB] unless there is a specific
  1976  // need for a continuous single database connection.
  1977  //
  1978  // A Conn must call [Conn.Close] to return the connection to the database pool
  1979  // and may do so concurrently with a running query.
  1980  //
  1981  // After a call to [Conn.Close], all operations on the
  1982  // connection fail with [ErrConnDone].
  1983  type Conn struct {
  1984  	db *DB
  1985  
  1986  	// closemu prevents the connection from closing while there
  1987  	// is an active query. It is held for read during queries
  1988  	// and exclusively during close.
  1989  	closemu sync.RWMutex
  1990  
  1991  	// dc is owned until close, at which point
  1992  	// it's returned to the connection pool.
  1993  	dc *driverConn
  1994  
  1995  	// done transitions from false to true exactly once, on close.
  1996  	// Once done, all operations fail with ErrConnDone.
  1997  	done atomic.Bool
  1998  
  1999  	releaseConnOnce sync.Once
  2000  	// releaseConnCache is a cache of c.closemuRUnlockCondReleaseConn
  2001  	// to save allocations in a call to grabConn.
  2002  	releaseConnCache releaseConn
  2003  }
  2004  
  2005  // grabConn takes a context to implement stmtConnGrabber
  2006  // but the context is not used.
  2007  func (c *Conn) grabConn(context.Context) (*driverConn, releaseConn, error) {
  2008  	if c.done.Load() {
  2009  		return nil, nil, ErrConnDone
  2010  	}
  2011  	c.releaseConnOnce.Do(func() {
  2012  		c.releaseConnCache = c.closemuRUnlockCondReleaseConn
  2013  	})
  2014  	c.closemu.RLock()
  2015  	return c.dc, c.releaseConnCache, nil
  2016  }
  2017  
  2018  // PingContext verifies the connection to the database is still alive.
  2019  func (c *Conn) PingContext(ctx context.Context) error {
  2020  	dc, release, err := c.grabConn(ctx)
  2021  	if err != nil {
  2022  		return err
  2023  	}
  2024  	return c.db.pingDC(ctx, dc, release)
  2025  }
  2026  
  2027  // ExecContext executes a query without returning any rows.
  2028  // The args are for any placeholder parameters in the query.
  2029  func (c *Conn) ExecContext(ctx context.Context, query string, args ...any) (Result, error) {
  2030  	dc, release, err := c.grabConn(ctx)
  2031  	if err != nil {
  2032  		return nil, err
  2033  	}
  2034  	return c.db.execDC(ctx, dc, release, query, args)
  2035  }
  2036  
  2037  // QueryContext executes a query that returns rows, typically a SELECT.
  2038  // The args are for any placeholder parameters in the query.
  2039  func (c *Conn) QueryContext(ctx context.Context, query string, args ...any) (*Rows, error) {
  2040  	dc, release, err := c.grabConn(ctx)
  2041  	if err != nil {
  2042  		return nil, err
  2043  	}
  2044  	return c.db.queryDC(ctx, nil, dc, release, query, args)
  2045  }
  2046  
  2047  // QueryRowContext executes a query that is expected to return at most one row.
  2048  // QueryRowContext always returns a non-nil value. Errors are deferred until
  2049  // the [*Row.Scan] method is called.
  2050  // If the query selects no rows, the [*Row.Scan] will return [ErrNoRows].
  2051  // Otherwise, the [*Row.Scan] scans the first selected row and discards
  2052  // the rest.
  2053  func (c *Conn) QueryRowContext(ctx context.Context, query string, args ...any) *Row {
  2054  	rows, err := c.QueryContext(ctx, query, args...)
  2055  	return &Row{rows: rows, err: err}
  2056  }
  2057  
  2058  // PrepareContext creates a prepared statement for later queries or executions.
  2059  // Multiple queries or executions may be run concurrently from the
  2060  // returned statement.
  2061  // The caller must call the statement's [*Stmt.Close] method
  2062  // when the statement is no longer needed.
  2063  //
  2064  // The provided context is used for the preparation of the statement, not for the
  2065  // execution of the statement.
  2066  func (c *Conn) PrepareContext(ctx context.Context, query string) (*Stmt, error) {
  2067  	dc, release, err := c.grabConn(ctx)
  2068  	if err != nil {
  2069  		return nil, err
  2070  	}
  2071  	return c.db.prepareDC(ctx, dc, release, c, query)
  2072  }
  2073  
  2074  // Raw executes f exposing the underlying driver connection for the
  2075  // duration of f. The driverConn must not be used outside of f.
  2076  //
  2077  // Once f returns and err is not [driver.ErrBadConn], the [Conn] will continue to be usable
  2078  // until [Conn.Close] is called.
  2079  func (c *Conn) Raw(f func(driverConn any) error) (err error) {
  2080  	var dc *driverConn
  2081  	var release releaseConn
  2082  
  2083  	// grabConn takes a context to implement stmtConnGrabber, but the context is not used.
  2084  	dc, release, err = c.grabConn(nil)
  2085  	if err != nil {
  2086  		return
  2087  	}
  2088  	fPanic := true
  2089  	dc.Mutex.Lock()
  2090  	defer func() {
  2091  		dc.Mutex.Unlock()
  2092  
  2093  		// If f panics fPanic will remain true.
  2094  		// Ensure an error is passed to release so the connection
  2095  		// may be discarded.
  2096  		if fPanic {
  2097  			err = driver.ErrBadConn
  2098  		}
  2099  		release(err)
  2100  	}()
  2101  	err = f(dc.ci)
  2102  	fPanic = false
  2103  
  2104  	return
  2105  }
  2106  
  2107  // BeginTx starts a transaction.
  2108  //
  2109  // The provided context is used until the transaction is committed or rolled back.
  2110  // If the context is canceled, the sql package will roll back
  2111  // the transaction. [Tx.Commit] will return an error if the context provided to
  2112  // BeginTx is canceled.
  2113  //
  2114  // The provided [TxOptions] is optional and may be nil if defaults should be used.
  2115  // If a non-default isolation level is used that the driver doesn't support,
  2116  // an error will be returned.
  2117  func (c *Conn) BeginTx(ctx context.Context, opts *TxOptions) (*Tx, error) {
  2118  	dc, release, err := c.grabConn(ctx)
  2119  	if err != nil {
  2120  		return nil, err
  2121  	}
  2122  	return c.db.beginDC(ctx, dc, release, opts)
  2123  }
  2124  
  2125  // closemuRUnlockCondReleaseConn read unlocks closemu
  2126  // as the sql operation is done with the dc.
  2127  func (c *Conn) closemuRUnlockCondReleaseConn(err error) {
  2128  	c.closemu.RUnlock()
  2129  	if errors.Is(err, driver.ErrBadConn) {
  2130  		c.close(err)
  2131  	}
  2132  }
  2133  
  2134  func (c *Conn) txCtx() context.Context {
  2135  	return nil
  2136  }
  2137  
  2138  func (c *Conn) close(err error) error {
  2139  	if !c.done.CompareAndSwap(false, true) {
  2140  		return ErrConnDone
  2141  	}
  2142  
  2143  	// Lock around releasing the driver connection
  2144  	// to ensure all queries have been stopped before doing so.
  2145  	c.closemu.Lock()
  2146  	defer c.closemu.Unlock()
  2147  
  2148  	c.dc.releaseConn(err)
  2149  	c.dc = nil
  2150  	c.db = nil
  2151  	return err
  2152  }
  2153  
  2154  // Close returns the connection to the connection pool.
  2155  // All operations after a Close will return with [ErrConnDone].
  2156  // Close is safe to call concurrently with other operations and will
  2157  // block until all other operations finish. It may be useful to first
  2158  // cancel any used context and then call close directly after.
  2159  func (c *Conn) Close() error {
  2160  	return c.close(nil)
  2161  }
  2162  
  2163  // Tx is an in-progress database transaction.
  2164  //
  2165  // A transaction must end with a call to [Tx.Commit] or [Tx.Rollback].
  2166  //
  2167  // After a call to [Tx.Commit] or [Tx.Rollback], all operations on the
  2168  // transaction fail with [ErrTxDone].
  2169  //
  2170  // The statements prepared for a transaction by calling
  2171  // the transaction's [Tx.Prepare] or [Tx.Stmt] methods are closed
  2172  // by the call to [Tx.Commit] or [Tx.Rollback].
  2173  type Tx struct {
  2174  	db *DB
  2175  
  2176  	// closemu prevents the transaction from closing while there
  2177  	// is an active query. It is held for read during queries
  2178  	// and exclusively during close.
  2179  	closemu sync.RWMutex
  2180  
  2181  	// dc is owned exclusively until Commit or Rollback, at which point
  2182  	// it's returned with putConn.
  2183  	dc  *driverConn
  2184  	txi driver.Tx
  2185  
  2186  	// releaseConn is called once the Tx is closed to release
  2187  	// any held driverConn back to the pool.
  2188  	releaseConn func(error)
  2189  
  2190  	// done transitions from false to true exactly once, on Commit
  2191  	// or Rollback. once done, all operations fail with
  2192  	// ErrTxDone.
  2193  	done atomic.Bool
  2194  
  2195  	// keepConnOnRollback is true if the driver knows
  2196  	// how to reset the connection's session and if need be discard
  2197  	// the connection.
  2198  	keepConnOnRollback bool
  2199  
  2200  	// All Stmts prepared for this transaction. These will be closed after the
  2201  	// transaction has been committed or rolled back.
  2202  	stmts struct {
  2203  		sync.Mutex
  2204  		v []*Stmt
  2205  	}
  2206  
  2207  	// cancel is called after done transitions from 0 to 1.
  2208  	cancel func()
  2209  
  2210  	// ctx lives for the life of the transaction.
  2211  	ctx context.Context
  2212  }
  2213  
  2214  // awaitDone blocks until the context in Tx is canceled and rolls back
  2215  // the transaction if it's not already done.
  2216  func (tx *Tx) awaitDone() {
  2217  	// Wait for either the transaction to be committed or rolled
  2218  	// back, or for the associated context to be closed.
  2219  	<-tx.ctx.Done()
  2220  
  2221  	// Discard and close the connection used to ensure the
  2222  	// transaction is closed and the resources are released.  This
  2223  	// rollback does nothing if the transaction has already been
  2224  	// committed or rolled back.
  2225  	// Do not discard the connection if the connection knows
  2226  	// how to reset the session.
  2227  	discardConnection := !tx.keepConnOnRollback
  2228  	tx.rollback(discardConnection)
  2229  }
  2230  
  2231  func (tx *Tx) isDone() bool {
  2232  	return tx.done.Load()
  2233  }
  2234  
  2235  // ErrTxDone is returned by any operation that is performed on a transaction
  2236  // that has already been committed or rolled back.
  2237  var ErrTxDone = errors.New("sql: transaction has already been committed or rolled back")
  2238  
  2239  // close returns the connection to the pool and
  2240  // must only be called by Tx.rollback or Tx.Commit while
  2241  // tx is already canceled and won't be executed concurrently.
  2242  func (tx *Tx) close(err error) {
  2243  	tx.releaseConn(err)
  2244  	tx.dc = nil
  2245  	tx.txi = nil
  2246  }
  2247  
  2248  // hookTxGrabConn specifies an optional hook to be called on
  2249  // a successful call to (*Tx).grabConn. For tests.
  2250  var hookTxGrabConn func()
  2251  
  2252  func (tx *Tx) grabConn(ctx context.Context) (*driverConn, releaseConn, error) {
  2253  	select {
  2254  	default:
  2255  	case <-ctx.Done():
  2256  		return nil, nil, ctx.Err()
  2257  	}
  2258  
  2259  	// closemu.RLock must come before the check for isDone to prevent the Tx from
  2260  	// closing while a query is executing.
  2261  	tx.closemu.RLock()
  2262  	if tx.isDone() {
  2263  		tx.closemu.RUnlock()
  2264  		return nil, nil, ErrTxDone
  2265  	}
  2266  	if hookTxGrabConn != nil { // test hook
  2267  		hookTxGrabConn()
  2268  	}
  2269  	return tx.dc, tx.closemuRUnlockRelease, nil
  2270  }
  2271  
  2272  func (tx *Tx) txCtx() context.Context {
  2273  	return tx.ctx
  2274  }
  2275  
  2276  // closemuRUnlockRelease is used as a func(error) method value in
  2277  // [DB.ExecContext] and [DB.QueryContext]. Unlocking in the releaseConn keeps
  2278  // the driver conn from being returned to the connection pool until
  2279  // the Rows has been closed.
  2280  func (tx *Tx) closemuRUnlockRelease(error) {
  2281  	tx.closemu.RUnlock()
  2282  }
  2283  
  2284  // Closes all Stmts prepared for this transaction.
  2285  func (tx *Tx) closePrepared() {
  2286  	tx.stmts.Lock()
  2287  	defer tx.stmts.Unlock()
  2288  	for _, stmt := range tx.stmts.v {
  2289  		stmt.Close()
  2290  	}
  2291  }
  2292  
  2293  // Commit commits the transaction.
  2294  func (tx *Tx) Commit() error {
  2295  	// Check context first to avoid transaction leak.
  2296  	// If put it behind tx.done CompareAndSwap statement, we can't ensure
  2297  	// the consistency between tx.done and the real COMMIT operation.
  2298  	select {
  2299  	default:
  2300  	case <-tx.ctx.Done():
  2301  		if tx.done.Load() {
  2302  			return ErrTxDone
  2303  		}
  2304  		return tx.ctx.Err()
  2305  	}
  2306  	if !tx.done.CompareAndSwap(false, true) {
  2307  		return ErrTxDone
  2308  	}
  2309  
  2310  	// Cancel the Tx to release any active R-closemu locks.
  2311  	// This is safe to do because tx.done has already transitioned
  2312  	// from 0 to 1. Hold the W-closemu lock prior to rollback
  2313  	// to ensure no other connection has an active query.
  2314  	tx.cancel()
  2315  	tx.closemu.Lock()
  2316  	tx.closemu.Unlock()
  2317  
  2318  	var err error
  2319  	withLock(tx.dc, func() {
  2320  		err = tx.txi.Commit()
  2321  	})
  2322  	if !errors.Is(err, driver.ErrBadConn) {
  2323  		tx.closePrepared()
  2324  	}
  2325  	tx.close(err)
  2326  	return err
  2327  }
  2328  
  2329  var rollbackHook func()
  2330  
  2331  // rollback aborts the transaction and optionally forces the pool to discard
  2332  // the connection.
  2333  func (tx *Tx) rollback(discardConn bool) error {
  2334  	if !tx.done.CompareAndSwap(false, true) {
  2335  		return ErrTxDone
  2336  	}
  2337  
  2338  	if rollbackHook != nil {
  2339  		rollbackHook()
  2340  	}
  2341  
  2342  	// Cancel the Tx to release any active R-closemu locks.
  2343  	// This is safe to do because tx.done has already transitioned
  2344  	// from 0 to 1. Hold the W-closemu lock prior to rollback
  2345  	// to ensure no other connection has an active query.
  2346  	tx.cancel()
  2347  	tx.closemu.Lock()
  2348  	tx.closemu.Unlock()
  2349  
  2350  	var err error
  2351  	withLock(tx.dc, func() {
  2352  		err = tx.txi.Rollback()
  2353  	})
  2354  	if !errors.Is(err, driver.ErrBadConn) {
  2355  		tx.closePrepared()
  2356  	}
  2357  	if discardConn {
  2358  		err = driver.ErrBadConn
  2359  	}
  2360  	tx.close(err)
  2361  	return err
  2362  }
  2363  
  2364  // Rollback aborts the transaction.
  2365  func (tx *Tx) Rollback() error {
  2366  	return tx.rollback(false)
  2367  }
  2368  
  2369  // PrepareContext creates a prepared statement for use within a transaction.
  2370  //
  2371  // The returned statement operates within the transaction and will be closed
  2372  // when the transaction has been committed or rolled back.
  2373  //
  2374  // To use an existing prepared statement on this transaction, see [Tx.Stmt].
  2375  //
  2376  // The provided context will be used for the preparation of the context, not
  2377  // for the execution of the returned statement. The returned statement
  2378  // will run in the transaction context.
  2379  func (tx *Tx) PrepareContext(ctx context.Context, query string) (*Stmt, error) {
  2380  	dc, release, err := tx.grabConn(ctx)
  2381  	if err != nil {
  2382  		return nil, err
  2383  	}
  2384  
  2385  	stmt, err := tx.db.prepareDC(ctx, dc, release, tx, query)
  2386  	if err != nil {
  2387  		return nil, err
  2388  	}
  2389  	tx.stmts.Lock()
  2390  	tx.stmts.v = append(tx.stmts.v, stmt)
  2391  	tx.stmts.Unlock()
  2392  	return stmt, nil
  2393  }
  2394  
  2395  // Prepare creates a prepared statement for use within a transaction.
  2396  //
  2397  // The returned statement operates within the transaction and will be closed
  2398  // when the transaction has been committed or rolled back.
  2399  //
  2400  // To use an existing prepared statement on this transaction, see [Tx.Stmt].
  2401  //
  2402  // Prepare uses [context.Background] internally; to specify the context, use
  2403  // [Tx.PrepareContext].
  2404  func (tx *Tx) Prepare(query string) (*Stmt, error) {
  2405  	return tx.PrepareContext(context.Background(), query)
  2406  }
  2407  
  2408  // StmtContext returns a transaction-specific prepared statement from
  2409  // an existing statement.
  2410  //
  2411  // Example:
  2412  //
  2413  //	updateMoney, err := db.Prepare("UPDATE balance SET money=money+? WHERE id=?")
  2414  //	...
  2415  //	tx, err := db.Begin()
  2416  //	...
  2417  //	res, err := tx.StmtContext(ctx, updateMoney).Exec(123.45, 98293203)
  2418  //
  2419  // The provided context is used for the preparation of the statement, not for the
  2420  // execution of the statement.
  2421  //
  2422  // The returned statement operates within the transaction and will be closed
  2423  // when the transaction has been committed or rolled back.
  2424  func (tx *Tx) StmtContext(ctx context.Context, stmt *Stmt) *Stmt {
  2425  	dc, release, err := tx.grabConn(ctx)
  2426  	if err != nil {
  2427  		return &Stmt{stickyErr: err}
  2428  	}
  2429  	defer release(nil)
  2430  
  2431  	if tx.db != stmt.db {
  2432  		return &Stmt{stickyErr: errors.New("sql: Tx.Stmt: statement from different database used")}
  2433  	}
  2434  	var si driver.Stmt
  2435  	var parentStmt *Stmt
  2436  	stmt.mu.Lock()
  2437  	if stmt.closed || stmt.cg != nil {
  2438  		// If the statement has been closed or already belongs to a
  2439  		// transaction, we can't reuse it in this connection.
  2440  		// Since tx.StmtContext should never need to be called with a
  2441  		// Stmt already belonging to tx, we ignore this edge case and
  2442  		// re-prepare the statement in this case. No need to add
  2443  		// code-complexity for this.
  2444  		stmt.mu.Unlock()
  2445  		withLock(dc, func() {
  2446  			si, err = ctxDriverPrepare(ctx, dc.ci, stmt.query)
  2447  		})
  2448  		if err != nil {
  2449  			return &Stmt{stickyErr: err}
  2450  		}
  2451  	} else {
  2452  		stmt.removeClosedStmtLocked()
  2453  		// See if the statement has already been prepared on this connection,
  2454  		// and reuse it if possible.
  2455  		for _, v := range stmt.css {
  2456  			if v.dc == dc {
  2457  				si = v.ds.si
  2458  				break
  2459  			}
  2460  		}
  2461  
  2462  		stmt.mu.Unlock()
  2463  
  2464  		if si == nil {
  2465  			var ds *driverStmt
  2466  			withLock(dc, func() {
  2467  				ds, err = stmt.prepareOnConnLocked(ctx, dc)
  2468  			})
  2469  			if err != nil {
  2470  				return &Stmt{stickyErr: err}
  2471  			}
  2472  			si = ds.si
  2473  		}
  2474  		parentStmt = stmt
  2475  	}
  2476  
  2477  	txs := &Stmt{
  2478  		db: tx.db,
  2479  		cg: tx,
  2480  		cgds: &driverStmt{
  2481  			Locker: dc,
  2482  			si:     si,
  2483  		},
  2484  		parentStmt: parentStmt,
  2485  		query:      stmt.query,
  2486  	}
  2487  	if parentStmt != nil {
  2488  		tx.db.addDep(parentStmt, txs)
  2489  	}
  2490  	tx.stmts.Lock()
  2491  	tx.stmts.v = append(tx.stmts.v, txs)
  2492  	tx.stmts.Unlock()
  2493  	return txs
  2494  }
  2495  
  2496  // Stmt returns a transaction-specific prepared statement from
  2497  // an existing statement.
  2498  //
  2499  // Example:
  2500  //
  2501  //	updateMoney, err := db.Prepare("UPDATE balance SET money=money+? WHERE id=?")
  2502  //	...
  2503  //	tx, err := db.Begin()
  2504  //	...
  2505  //	res, err := tx.Stmt(updateMoney).Exec(123.45, 98293203)
  2506  //
  2507  // The returned statement operates within the transaction and will be closed
  2508  // when the transaction has been committed or rolled back.
  2509  //
  2510  // Stmt uses [context.Background] internally; to specify the context, use
  2511  // [Tx.StmtContext].
  2512  func (tx *Tx) Stmt(stmt *Stmt) *Stmt {
  2513  	return tx.StmtContext(context.Background(), stmt)
  2514  }
  2515  
  2516  // ExecContext executes a query that doesn't return rows.
  2517  // For example: an INSERT and UPDATE.
  2518  func (tx *Tx) ExecContext(ctx context.Context, query string, args ...any) (Result, error) {
  2519  	dc, release, err := tx.grabConn(ctx)
  2520  	if err != nil {
  2521  		return nil, err
  2522  	}
  2523  	return tx.db.execDC(ctx, dc, release, query, args)
  2524  }
  2525  
  2526  // Exec executes a query that doesn't return rows.
  2527  // For example: an INSERT and UPDATE.
  2528  //
  2529  // Exec uses [context.Background] internally; to specify the context, use
  2530  // [Tx.ExecContext].
  2531  func (tx *Tx) Exec(query string, args ...any) (Result, error) {
  2532  	return tx.ExecContext(context.Background(), query, args...)
  2533  }
  2534  
  2535  // QueryContext executes a query that returns rows, typically a SELECT.
  2536  func (tx *Tx) QueryContext(ctx context.Context, query string, args ...any) (*Rows, error) {
  2537  	dc, release, err := tx.grabConn(ctx)
  2538  	if err != nil {
  2539  		return nil, err
  2540  	}
  2541  
  2542  	return tx.db.queryDC(ctx, tx.ctx, dc, release, query, args)
  2543  }
  2544  
  2545  // Query executes a query that returns rows, typically a SELECT.
  2546  //
  2547  // Query uses [context.Background] internally; to specify the context, use
  2548  // [Tx.QueryContext].
  2549  func (tx *Tx) Query(query string, args ...any) (*Rows, error) {
  2550  	return tx.QueryContext(context.Background(), query, args...)
  2551  }
  2552  
  2553  // QueryRowContext executes a query that is expected to return at most one row.
  2554  // QueryRowContext always returns a non-nil value. Errors are deferred until
  2555  // [Row]'s Scan method is called.
  2556  // If the query selects no rows, the [*Row.Scan] will return [ErrNoRows].
  2557  // Otherwise, the [*Row.Scan] scans the first selected row and discards
  2558  // the rest.
  2559  func (tx *Tx) QueryRowContext(ctx context.Context, query string, args ...any) *Row {
  2560  	rows, err := tx.QueryContext(ctx, query, args...)
  2561  	return &Row{rows: rows, err: err}
  2562  }
  2563  
  2564  // QueryRow executes a query that is expected to return at most one row.
  2565  // QueryRow always returns a non-nil value. Errors are deferred until
  2566  // [Row]'s Scan method is called.
  2567  // If the query selects no rows, the [*Row.Scan] will return [ErrNoRows].
  2568  // Otherwise, the [*Row.Scan] scans the first selected row and discards
  2569  // the rest.
  2570  //
  2571  // QueryRow uses [context.Background] internally; to specify the context, use
  2572  // [Tx.QueryRowContext].
  2573  func (tx *Tx) QueryRow(query string, args ...any) *Row {
  2574  	return tx.QueryRowContext(context.Background(), query, args...)
  2575  }
  2576  
  2577  // connStmt is a prepared statement on a particular connection.
  2578  type connStmt struct {
  2579  	dc *driverConn
  2580  	ds *driverStmt
  2581  }
  2582  
  2583  // stmtConnGrabber represents a Tx or Conn that will return the underlying
  2584  // driverConn and release function.
  2585  type stmtConnGrabber interface {
  2586  	// grabConn returns the driverConn and the associated release function
  2587  	// that must be called when the operation completes.
  2588  	grabConn(context.Context) (*driverConn, releaseConn, error)
  2589  
  2590  	// txCtx returns the transaction context if available.
  2591  	// The returned context should be selected on along with
  2592  	// any query context when awaiting a cancel.
  2593  	txCtx() context.Context
  2594  }
  2595  
  2596  var (
  2597  	_ stmtConnGrabber = &Tx{}
  2598  	_ stmtConnGrabber = &Conn{}
  2599  )
  2600  
  2601  // Stmt is a prepared statement.
  2602  // A Stmt is safe for concurrent use by multiple goroutines.
  2603  //
  2604  // If a Stmt is prepared on a [Tx] or [Conn], it will be bound to a single
  2605  // underlying connection forever. If the [Tx] or [Conn] closes, the Stmt will
  2606  // become unusable and all operations will return an error.
  2607  // If a Stmt is prepared on a [DB], it will remain usable for the lifetime of the
  2608  // [DB]. When the Stmt needs to execute on a new underlying connection, it will
  2609  // prepare itself on the new connection automatically.
  2610  type Stmt struct {
  2611  	// Immutable:
  2612  	db        *DB    // where we came from
  2613  	query     string // that created the Stmt
  2614  	stickyErr error  // if non-nil, this error is returned for all operations
  2615  
  2616  	closemu sync.RWMutex // held exclusively during close, for read otherwise.
  2617  
  2618  	// If Stmt is prepared on a Tx or Conn then cg is present and will
  2619  	// only ever grab a connection from cg.
  2620  	// If cg is nil then the Stmt must grab an arbitrary connection
  2621  	// from db and determine if it must prepare the stmt again by
  2622  	// inspecting css.
  2623  	cg   stmtConnGrabber
  2624  	cgds *driverStmt
  2625  
  2626  	// parentStmt is set when a transaction-specific statement
  2627  	// is requested from an identical statement prepared on the same
  2628  	// conn. parentStmt is used to track the dependency of this statement
  2629  	// on its originating ("parent") statement so that parentStmt may
  2630  	// be closed by the user without them having to know whether or not
  2631  	// any transactions are still using it.
  2632  	parentStmt *Stmt
  2633  
  2634  	mu     sync.Mutex // protects the rest of the fields
  2635  	closed bool
  2636  
  2637  	// css is a list of underlying driver statement interfaces
  2638  	// that are valid on particular connections. This is only
  2639  	// used if cg == nil and one is found that has idle
  2640  	// connections. If cg != nil, cgds is always used.
  2641  	css []connStmt
  2642  
  2643  	// lastNumClosed is copied from db.numClosed when Stmt is created
  2644  	// without tx and closed connections in css are removed.
  2645  	lastNumClosed uint64
  2646  }
  2647  
  2648  // ExecContext executes a prepared statement with the given arguments and
  2649  // returns a [Result] summarizing the effect of the statement.
  2650  func (s *Stmt) ExecContext(ctx context.Context, args ...any) (Result, error) {
  2651  	s.closemu.RLock()
  2652  	defer s.closemu.RUnlock()
  2653  
  2654  	var res Result
  2655  	err := s.db.retry(func(strategy connReuseStrategy) error {
  2656  		dc, releaseConn, ds, err := s.connStmt(ctx, strategy)
  2657  		if err != nil {
  2658  			return err
  2659  		}
  2660  
  2661  		res, err = resultFromStatement(ctx, dc.ci, ds, args...)
  2662  		releaseConn(err)
  2663  		return err
  2664  	})
  2665  
  2666  	return res, err
  2667  }
  2668  
  2669  // Exec executes a prepared statement with the given arguments and
  2670  // returns a [Result] summarizing the effect of the statement.
  2671  //
  2672  // Exec uses [context.Background] internally; to specify the context, use
  2673  // [Stmt.ExecContext].
  2674  func (s *Stmt) Exec(args ...any) (Result, error) {
  2675  	return s.ExecContext(context.Background(), args...)
  2676  }
  2677  
  2678  func resultFromStatement(ctx context.Context, ci driver.Conn, ds *driverStmt, args ...any) (Result, error) {
  2679  	ds.Lock()
  2680  	defer ds.Unlock()
  2681  
  2682  	dargs, err := driverArgsConnLocked(ci, ds, args)
  2683  	if err != nil {
  2684  		return nil, err
  2685  	}
  2686  
  2687  	resi, err := ctxDriverStmtExec(ctx, ds.si, dargs)
  2688  	if err != nil {
  2689  		return nil, err
  2690  	}
  2691  	return driverResult{ds.Locker, resi}, nil
  2692  }
  2693  
  2694  // removeClosedStmtLocked removes closed conns in s.css.
  2695  //
  2696  // To avoid lock contention on DB.mu, we do it only when
  2697  // s.db.numClosed - s.lastNum is large enough.
  2698  func (s *Stmt) removeClosedStmtLocked() {
  2699  	t := len(s.css)/2 + 1
  2700  	if t > 10 {
  2701  		t = 10
  2702  	}
  2703  	dbClosed := s.db.numClosed.Load()
  2704  	if dbClosed-s.lastNumClosed < uint64(t) {
  2705  		return
  2706  	}
  2707  
  2708  	s.db.mu.Lock()
  2709  	for i := 0; i < len(s.css); i++ {
  2710  		if s.css[i].dc.dbmuClosed {
  2711  			s.css[i] = s.css[len(s.css)-1]
  2712  			// Zero out the last element (for GC) before shrinking the slice.
  2713  			s.css[len(s.css)-1] = connStmt{}
  2714  			s.css = s.css[:len(s.css)-1]
  2715  			i--
  2716  		}
  2717  	}
  2718  	s.db.mu.Unlock()
  2719  	s.lastNumClosed = dbClosed
  2720  }
  2721  
  2722  // connStmt returns a free driver connection on which to execute the
  2723  // statement, a function to call to release the connection, and a
  2724  // statement bound to that connection.
  2725  func (s *Stmt) connStmt(ctx context.Context, strategy connReuseStrategy) (dc *driverConn, releaseConn func(error), ds *driverStmt, err error) {
  2726  	if err = s.stickyErr; err != nil {
  2727  		return
  2728  	}
  2729  	s.mu.Lock()
  2730  	if s.closed {
  2731  		s.mu.Unlock()
  2732  		err = errors.New("sql: statement is closed")
  2733  		return
  2734  	}
  2735  
  2736  	// In a transaction or connection, we always use the connection that the
  2737  	// stmt was created on.
  2738  	if s.cg != nil {
  2739  		s.mu.Unlock()
  2740  		dc, releaseConn, err = s.cg.grabConn(ctx) // blocks, waiting for the connection.
  2741  		if err != nil {
  2742  			return
  2743  		}
  2744  		return dc, releaseConn, s.cgds, nil
  2745  	}
  2746  
  2747  	s.removeClosedStmtLocked()
  2748  	s.mu.Unlock()
  2749  
  2750  	dc, err = s.db.conn(ctx, strategy)
  2751  	if err != nil {
  2752  		return nil, nil, nil, err
  2753  	}
  2754  
  2755  	s.mu.Lock()
  2756  	for _, v := range s.css {
  2757  		if v.dc == dc {
  2758  			s.mu.Unlock()
  2759  			return dc, dc.releaseConn, v.ds, nil
  2760  		}
  2761  	}
  2762  	s.mu.Unlock()
  2763  
  2764  	// No luck; we need to prepare the statement on this connection
  2765  	withLock(dc, func() {
  2766  		ds, err = s.prepareOnConnLocked(ctx, dc)
  2767  	})
  2768  	if err != nil {
  2769  		dc.releaseConn(err)
  2770  		return nil, nil, nil, err
  2771  	}
  2772  
  2773  	return dc, dc.releaseConn, ds, nil
  2774  }
  2775  
  2776  // prepareOnConnLocked prepares the query in Stmt s on dc and adds it to the list of
  2777  // open connStmt on the statement. It assumes the caller is holding the lock on dc.
  2778  func (s *Stmt) prepareOnConnLocked(ctx context.Context, dc *driverConn) (*driverStmt, error) {
  2779  	si, err := dc.prepareLocked(ctx, s.cg, s.query)
  2780  	if err != nil {
  2781  		return nil, err
  2782  	}
  2783  	cs := connStmt{dc, si}
  2784  	s.mu.Lock()
  2785  	s.css = append(s.css, cs)
  2786  	s.mu.Unlock()
  2787  	return cs.ds, nil
  2788  }
  2789  
  2790  // QueryContext executes a prepared query statement with the given arguments
  2791  // and returns the query results as a [*Rows].
  2792  func (s *Stmt) QueryContext(ctx context.Context, args ...any) (*Rows, error) {
  2793  	s.closemu.RLock()
  2794  	defer s.closemu.RUnlock()
  2795  
  2796  	var rowsi driver.Rows
  2797  	var rows *Rows
  2798  
  2799  	err := s.db.retry(func(strategy connReuseStrategy) error {
  2800  		dc, releaseConn, ds, err := s.connStmt(ctx, strategy)
  2801  		if err != nil {
  2802  			return err
  2803  		}
  2804  
  2805  		rowsi, err = rowsiFromStatement(ctx, dc.ci, ds, args...)
  2806  		if err == nil {
  2807  			// Note: ownership of ci passes to the *Rows, to be freed
  2808  			// with releaseConn.
  2809  			rows = &Rows{
  2810  				dc:    dc,
  2811  				rowsi: rowsi,
  2812  				// releaseConn set below
  2813  			}
  2814  			// addDep must be added before initContextClose or it could attempt
  2815  			// to removeDep before it has been added.
  2816  			s.db.addDep(s, rows)
  2817  
  2818  			// releaseConn must be set before initContextClose or it could
  2819  			// release the connection before it is set.
  2820  			rows.releaseConn = func(err error) {
  2821  				releaseConn(err)
  2822  				s.db.removeDep(s, rows)
  2823  			}
  2824  			var txctx context.Context
  2825  			if s.cg != nil {
  2826  				txctx = s.cg.txCtx()
  2827  			}
  2828  			rows.initContextClose(ctx, txctx)
  2829  			return nil
  2830  		}
  2831  
  2832  		releaseConn(err)
  2833  		return err
  2834  	})
  2835  
  2836  	return rows, err
  2837  }
  2838  
  2839  // Query executes a prepared query statement with the given arguments
  2840  // and returns the query results as a *Rows.
  2841  //
  2842  // Query uses [context.Background] internally; to specify the context, use
  2843  // [Stmt.QueryContext].
  2844  func (s *Stmt) Query(args ...any) (*Rows, error) {
  2845  	return s.QueryContext(context.Background(), args...)
  2846  }
  2847  
  2848  func rowsiFromStatement(ctx context.Context, ci driver.Conn, ds *driverStmt, args ...any) (driver.Rows, error) {
  2849  	ds.Lock()
  2850  	defer ds.Unlock()
  2851  	dargs, err := driverArgsConnLocked(ci, ds, args)
  2852  	if err != nil {
  2853  		return nil, err
  2854  	}
  2855  	return ctxDriverStmtQuery(ctx, ds.si, dargs)
  2856  }
  2857  
  2858  // QueryRowContext executes a prepared query statement with the given arguments.
  2859  // If an error occurs during the execution of the statement, that error will
  2860  // be returned by a call to Scan on the returned [*Row], which is always non-nil.
  2861  // If the query selects no rows, the [*Row.Scan] will return [ErrNoRows].
  2862  // Otherwise, the [*Row.Scan] scans the first selected row and discards
  2863  // the rest.
  2864  func (s *Stmt) QueryRowContext(ctx context.Context, args ...any) *Row {
  2865  	rows, err := s.QueryContext(ctx, args...)
  2866  	if err != nil {
  2867  		return &Row{err: err}
  2868  	}
  2869  	return &Row{rows: rows}
  2870  }
  2871  
  2872  // QueryRow executes a prepared query statement with the given arguments.
  2873  // If an error occurs during the execution of the statement, that error will
  2874  // be returned by a call to Scan on the returned [*Row], which is always non-nil.
  2875  // If the query selects no rows, the [*Row.Scan] will return [ErrNoRows].
  2876  // Otherwise, the [*Row.Scan] scans the first selected row and discards
  2877  // the rest.
  2878  //
  2879  // Example usage:
  2880  //
  2881  //	var name string
  2882  //	err := nameByUseridStmt.QueryRow(id).Scan(&name)
  2883  //
  2884  // QueryRow uses [context.Background] internally; to specify the context, use
  2885  // [Stmt.QueryRowContext].
  2886  func (s *Stmt) QueryRow(args ...any) *Row {
  2887  	return s.QueryRowContext(context.Background(), args...)
  2888  }
  2889  
  2890  // Close closes the statement.
  2891  func (s *Stmt) Close() error {
  2892  	s.closemu.Lock()
  2893  	defer s.closemu.Unlock()
  2894  
  2895  	if s.stickyErr != nil {
  2896  		return s.stickyErr
  2897  	}
  2898  	s.mu.Lock()
  2899  	if s.closed {
  2900  		s.mu.Unlock()
  2901  		return nil
  2902  	}
  2903  	s.closed = true
  2904  	txds := s.cgds
  2905  	s.cgds = nil
  2906  
  2907  	s.mu.Unlock()
  2908  
  2909  	if s.cg == nil {
  2910  		return s.db.removeDep(s, s)
  2911  	}
  2912  
  2913  	if s.parentStmt != nil {
  2914  		// If parentStmt is set, we must not close s.txds since it's stored
  2915  		// in the css array of the parentStmt.
  2916  		return s.db.removeDep(s.parentStmt, s)
  2917  	}
  2918  	return txds.Close()
  2919  }
  2920  
  2921  func (s *Stmt) finalClose() error {
  2922  	s.mu.Lock()
  2923  	defer s.mu.Unlock()
  2924  	if s.css != nil {
  2925  		for _, v := range s.css {
  2926  			s.db.noteUnusedDriverStatement(v.dc, v.ds)
  2927  			v.dc.removeOpenStmt(v.ds)
  2928  		}
  2929  		s.css = nil
  2930  	}
  2931  	return nil
  2932  }
  2933  
  2934  // Rows is the result of a query. Its cursor starts before the first row
  2935  // of the result set. Use [Rows.Next] to advance from row to row.
  2936  type Rows struct {
  2937  	dc          *driverConn // owned; must call releaseConn when closed to release
  2938  	releaseConn func(error)
  2939  	rowsi       driver.Rows
  2940  	cancel      func()      // called when Rows is closed, may be nil.
  2941  	closeStmt   *driverStmt // if non-nil, statement to Close on close
  2942  
  2943  	contextDone atomic.Pointer[error] // error that awaitDone saw; set before close attempt
  2944  
  2945  	// closemu prevents Rows from closing while there
  2946  	// is an active streaming result. It is held for read during non-close operations
  2947  	// and exclusively during close.
  2948  	//
  2949  	// closemu guards lasterr and closed.
  2950  	closemu sync.RWMutex
  2951  	lasterr error // non-nil only if closed is true
  2952  	closed  bool
  2953  
  2954  	// closemuScanHold is whether the previous call to Scan kept closemu RLock'ed
  2955  	// without unlocking it. It does that when the user passes a *RawBytes scan
  2956  	// target. In that case, we need to prevent awaitDone from closing the Rows
  2957  	// while the user's still using the memory. See go.dev/issue/60304.
  2958  	//
  2959  	// It is only used by Scan, Next, and NextResultSet which are expected
  2960  	// not to be called concurrently.
  2961  	closemuScanHold bool
  2962  
  2963  	// hitEOF is whether Next hit the end of the rows without
  2964  	// encountering an error. It's set in Next before
  2965  	// returning. It's only used by Next and Err which are
  2966  	// expected not to be called concurrently.
  2967  	hitEOF bool
  2968  
  2969  	// lastcols is only used in Scan, Next, and NextResultSet which are expected
  2970  	// not to be called concurrently.
  2971  	lastcols []driver.Value
  2972  
  2973  	// raw is a buffer for RawBytes that persists between Scan calls.
  2974  	// This is used when the driver returns a mismatched type that requires
  2975  	// a cloning allocation. For example, if the driver returns a *string and
  2976  	// the user is scanning into a *RawBytes, we need to copy the string.
  2977  	// The raw buffer here lets us reuse the memory for that copy across Scan calls.
  2978  	raw []byte
  2979  }
  2980  
  2981  // lasterrOrErrLocked returns either lasterr or the provided err.
  2982  // rs.closemu must be read-locked.
  2983  func (rs *Rows) lasterrOrErrLocked(err error) error {
  2984  	if rs.lasterr != nil && rs.lasterr != io.EOF {
  2985  		return rs.lasterr
  2986  	}
  2987  	return err
  2988  }
  2989  
  2990  // bypassRowsAwaitDone is only used for testing.
  2991  // If true, it will not close the Rows automatically from the context.
  2992  var bypassRowsAwaitDone = false
  2993  
  2994  func (rs *Rows) initContextClose(ctx, txctx context.Context) {
  2995  	if ctx.Done() == nil && (txctx == nil || txctx.Done() == nil) {
  2996  		return
  2997  	}
  2998  	if bypassRowsAwaitDone {
  2999  		return
  3000  	}
  3001  	closectx, cancel := context.WithCancel(ctx)
  3002  	rs.cancel = cancel
  3003  	go rs.awaitDone(ctx, txctx, closectx)
  3004  }
  3005  
  3006  // awaitDone blocks until ctx, txctx, or closectx is canceled.
  3007  // The ctx is provided from the query context.
  3008  // If the query was issued in a transaction, the transaction's context
  3009  // is also provided in txctx, to ensure Rows is closed if the Tx is closed.
  3010  // The closectx is closed by an explicit call to rs.Close.
  3011  func (rs *Rows) awaitDone(ctx, txctx, closectx context.Context) {
  3012  	var txctxDone <-chan struct{}
  3013  	if txctx != nil {
  3014  		txctxDone = txctx.Done()
  3015  	}
  3016  	select {
  3017  	case <-ctx.Done():
  3018  		err := ctx.Err()
  3019  		rs.contextDone.Store(&err)
  3020  	case <-txctxDone:
  3021  		err := txctx.Err()
  3022  		rs.contextDone.Store(&err)
  3023  	case <-closectx.Done():
  3024  		// rs.cancel was called via Close(); don't store this into contextDone
  3025  		// to ensure Err() is unaffected.
  3026  	}
  3027  	rs.close(ctx.Err())
  3028  }
  3029  
  3030  // Next prepares the next result row for reading with the [Rows.Scan] method. It
  3031  // returns true on success, or false if there is no next result row or an error
  3032  // happened while preparing it. [Rows.Err] should be consulted to distinguish between
  3033  // the two cases.
  3034  //
  3035  // Every call to [Rows.Scan], even the first one, must be preceded by a call to [Rows.Next].
  3036  func (rs *Rows) Next() bool {
  3037  	// If the user's calling Next, they're done with their previous row's Scan
  3038  	// results (any RawBytes memory), so we can release the read lock that would
  3039  	// be preventing awaitDone from calling close.
  3040  	rs.closemuRUnlockIfHeldByScan()
  3041  
  3042  	if rs.contextDone.Load() != nil {
  3043  		return false
  3044  	}
  3045  
  3046  	var doClose, ok bool
  3047  	withLock(rs.closemu.RLocker(), func() {
  3048  		doClose, ok = rs.nextLocked()
  3049  	})
  3050  	if doClose {
  3051  		rs.Close()
  3052  	}
  3053  	if doClose && !ok {
  3054  		rs.hitEOF = true
  3055  	}
  3056  	return ok
  3057  }
  3058  
  3059  func (rs *Rows) nextLocked() (doClose, ok bool) {
  3060  	if rs.closed {
  3061  		return false, false
  3062  	}
  3063  
  3064  	// Lock the driver connection before calling the driver interface
  3065  	// rowsi to prevent a Tx from rolling back the connection at the same time.
  3066  	rs.dc.Lock()
  3067  	defer rs.dc.Unlock()
  3068  
  3069  	if rs.lastcols == nil {
  3070  		rs.lastcols = make([]driver.Value, len(rs.rowsi.Columns()))
  3071  	}
  3072  
  3073  	rs.lasterr = rs.rowsi.Next(rs.lastcols)
  3074  	if rs.lasterr != nil {
  3075  		// Close the connection if there is a driver error.
  3076  		if rs.lasterr != io.EOF {
  3077  			return true, false
  3078  		}
  3079  		nextResultSet, ok := rs.rowsi.(driver.RowsNextResultSet)
  3080  		if !ok {
  3081  			return true, false
  3082  		}
  3083  		// The driver is at the end of the current result set.
  3084  		// Test to see if there is another result set after the current one.
  3085  		// Only close Rows if there is no further result sets to read.
  3086  		if !nextResultSet.HasNextResultSet() {
  3087  			doClose = true
  3088  		}
  3089  		return doClose, false
  3090  	}
  3091  	return false, true
  3092  }
  3093  
  3094  // NextResultSet prepares the next result set for reading. It reports whether
  3095  // there is further result sets, or false if there is no further result set
  3096  // or if there is an error advancing to it. The [Rows.Err] method should be consulted
  3097  // to distinguish between the two cases.
  3098  //
  3099  // After calling NextResultSet, the [Rows.Next] method should always be called before
  3100  // scanning. If there are further result sets they may not have rows in the result
  3101  // set.
  3102  func (rs *Rows) NextResultSet() bool {
  3103  	// If the user's calling NextResultSet, they're done with their previous
  3104  	// row's Scan results (any RawBytes memory), so we can release the read lock
  3105  	// that would be preventing awaitDone from calling close.
  3106  	rs.closemuRUnlockIfHeldByScan()
  3107  
  3108  	var doClose bool
  3109  	defer func() {
  3110  		if doClose {
  3111  			rs.Close()
  3112  		}
  3113  	}()
  3114  	rs.closemu.RLock()
  3115  	defer rs.closemu.RUnlock()
  3116  
  3117  	if rs.closed {
  3118  		return false
  3119  	}
  3120  
  3121  	rs.lastcols = nil
  3122  	nextResultSet, ok := rs.rowsi.(driver.RowsNextResultSet)
  3123  	if !ok {
  3124  		doClose = true
  3125  		return false
  3126  	}
  3127  
  3128  	// Lock the driver connection before calling the driver interface
  3129  	// rowsi to prevent a Tx from rolling back the connection at the same time.
  3130  	rs.dc.Lock()
  3131  	defer rs.dc.Unlock()
  3132  
  3133  	rs.lasterr = nextResultSet.NextResultSet()
  3134  	if rs.lasterr != nil {
  3135  		doClose = true
  3136  		return false
  3137  	}
  3138  	return true
  3139  }
  3140  
  3141  // Err returns the error, if any, that was encountered during iteration.
  3142  // Err may be called after an explicit or implicit [Rows.Close].
  3143  func (rs *Rows) Err() error {
  3144  	// Return any context error that might've happened during row iteration,
  3145  	// but only if we haven't reported the final Next() = false after rows
  3146  	// are done, in which case the user might've canceled their own context
  3147  	// before calling Rows.Err.
  3148  	if !rs.hitEOF {
  3149  		if errp := rs.contextDone.Load(); errp != nil {
  3150  			return *errp
  3151  		}
  3152  	}
  3153  
  3154  	rs.closemu.RLock()
  3155  	defer rs.closemu.RUnlock()
  3156  	return rs.lasterrOrErrLocked(nil)
  3157  }
  3158  
  3159  // rawbuf returns the buffer to append RawBytes values to.
  3160  // This buffer is reused across calls to Rows.Scan.
  3161  //
  3162  // Usage:
  3163  //
  3164  //	rawBytes = rows.setrawbuf(append(rows.rawbuf(), value...))
  3165  func (rs *Rows) rawbuf() []byte {
  3166  	if rs == nil {
  3167  		// convertAssignRows can take a nil *Rows; for simplicity handle it here
  3168  		return nil
  3169  	}
  3170  	return rs.raw
  3171  }
  3172  
  3173  // setrawbuf updates the RawBytes buffer with the result of appending a new value to it.
  3174  // It returns the new value.
  3175  func (rs *Rows) setrawbuf(b []byte) RawBytes {
  3176  	if rs == nil {
  3177  		// convertAssignRows can take a nil *Rows; for simplicity handle it here
  3178  		return RawBytes(b)
  3179  	}
  3180  	off := len(rs.raw)
  3181  	rs.raw = b
  3182  	return RawBytes(rs.raw[off:])
  3183  }
  3184  
  3185  var errRowsClosed = errors.New("sql: Rows are closed")
  3186  var errNoRows = errors.New("sql: no Rows available")
  3187  
  3188  // Columns returns the column names.
  3189  // Columns returns an error if the rows are closed.
  3190  func (rs *Rows) Columns() ([]string, error) {
  3191  	rs.closemu.RLock()
  3192  	defer rs.closemu.RUnlock()
  3193  	if rs.closed {
  3194  		return nil, rs.lasterrOrErrLocked(errRowsClosed)
  3195  	}
  3196  	if rs.rowsi == nil {
  3197  		return nil, rs.lasterrOrErrLocked(errNoRows)
  3198  	}
  3199  	rs.dc.Lock()
  3200  	defer rs.dc.Unlock()
  3201  
  3202  	return rs.rowsi.Columns(), nil
  3203  }
  3204  
  3205  // ColumnTypes returns column information such as column type, length,
  3206  // and nullable. Some information may not be available from some drivers.
  3207  func (rs *Rows) ColumnTypes() ([]*ColumnType, error) {
  3208  	rs.closemu.RLock()
  3209  	defer rs.closemu.RUnlock()
  3210  	if rs.closed {
  3211  		return nil, rs.lasterrOrErrLocked(errRowsClosed)
  3212  	}
  3213  	if rs.rowsi == nil {
  3214  		return nil, rs.lasterrOrErrLocked(errNoRows)
  3215  	}
  3216  	rs.dc.Lock()
  3217  	defer rs.dc.Unlock()
  3218  
  3219  	return rowsColumnInfoSetupConnLocked(rs.rowsi), nil
  3220  }
  3221  
  3222  // ColumnType contains the name and type of a column.
  3223  type ColumnType struct {
  3224  	name string
  3225  
  3226  	hasNullable       bool
  3227  	hasLength         bool
  3228  	hasPrecisionScale bool
  3229  
  3230  	nullable     bool
  3231  	length       int64
  3232  	databaseType string
  3233  	precision    int64
  3234  	scale        int64
  3235  	scanType     reflect.Type
  3236  }
  3237  
  3238  // Name returns the name or alias of the column.
  3239  func (ci *ColumnType) Name() string {
  3240  	return ci.name
  3241  }
  3242  
  3243  // Length returns the column type length for variable length column types such
  3244  // as text and binary field types. If the type length is unbounded the value will
  3245  // be [math.MaxInt64] (any database limits will still apply).
  3246  // If the column type is not variable length, such as an int, or if not supported
  3247  // by the driver ok is false.
  3248  func (ci *ColumnType) Length() (length int64, ok bool) {
  3249  	return ci.length, ci.hasLength
  3250  }
  3251  
  3252  // DecimalSize returns the scale and precision of a decimal type.
  3253  // If not applicable or if not supported ok is false.
  3254  func (ci *ColumnType) DecimalSize() (precision, scale int64, ok bool) {
  3255  	return ci.precision, ci.scale, ci.hasPrecisionScale
  3256  }
  3257  
  3258  // ScanType returns a Go type suitable for scanning into using [Rows.Scan].
  3259  // If a driver does not support this property ScanType will return
  3260  // the type of an empty interface.
  3261  func (ci *ColumnType) ScanType() reflect.Type {
  3262  	return ci.scanType
  3263  }
  3264  
  3265  // Nullable reports whether the column may be null.
  3266  // If a driver does not support this property ok will be false.
  3267  func (ci *ColumnType) Nullable() (nullable, ok bool) {
  3268  	return ci.nullable, ci.hasNullable
  3269  }
  3270  
  3271  // DatabaseTypeName returns the database system name of the column type. If an empty
  3272  // string is returned, then the driver type name is not supported.
  3273  // Consult your driver documentation for a list of driver data types. [ColumnType.Length] specifiers
  3274  // are not included.
  3275  // Common type names include "VARCHAR", "TEXT", "NVARCHAR", "DECIMAL", "BOOL",
  3276  // "INT", and "BIGINT".
  3277  func (ci *ColumnType) DatabaseTypeName() string {
  3278  	return ci.databaseType
  3279  }
  3280  
  3281  func rowsColumnInfoSetupConnLocked(rowsi driver.Rows) []*ColumnType {
  3282  	names := rowsi.Columns()
  3283  
  3284  	list := make([]*ColumnType, len(names))
  3285  	for i := range list {
  3286  		ci := &ColumnType{
  3287  			name: names[i],
  3288  		}
  3289  		list[i] = ci
  3290  
  3291  		if prop, ok := rowsi.(driver.RowsColumnTypeScanType); ok {
  3292  			ci.scanType = prop.ColumnTypeScanType(i)
  3293  		} else {
  3294  			ci.scanType = reflect.TypeFor[any]()
  3295  		}
  3296  		if prop, ok := rowsi.(driver.RowsColumnTypeDatabaseTypeName); ok {
  3297  			ci.databaseType = prop.ColumnTypeDatabaseTypeName(i)
  3298  		}
  3299  		if prop, ok := rowsi.(driver.RowsColumnTypeLength); ok {
  3300  			ci.length, ci.hasLength = prop.ColumnTypeLength(i)
  3301  		}
  3302  		if prop, ok := rowsi.(driver.RowsColumnTypeNullable); ok {
  3303  			ci.nullable, ci.hasNullable = prop.ColumnTypeNullable(i)
  3304  		}
  3305  		if prop, ok := rowsi.(driver.RowsColumnTypePrecisionScale); ok {
  3306  			ci.precision, ci.scale, ci.hasPrecisionScale = prop.ColumnTypePrecisionScale(i)
  3307  		}
  3308  	}
  3309  	return list
  3310  }
  3311  
  3312  // Scan copies the columns in the current row into the values pointed
  3313  // at by dest. The number of values in dest must be the same as the
  3314  // number of columns in [Rows].
  3315  //
  3316  // Scan converts columns read from the database into the following
  3317  // common Go types and special types provided by the sql package:
  3318  //
  3319  //	*string
  3320  //	*[]byte
  3321  //	*int, *int8, *int16, *int32, *int64
  3322  //	*uint, *uint8, *uint16, *uint32, *uint64
  3323  //	*bool
  3324  //	*float32, *float64
  3325  //	*interface{}
  3326  //	*RawBytes
  3327  //	*Rows (cursor value)
  3328  //	any type implementing Scanner (see Scanner docs)
  3329  //
  3330  // In the most simple case, if the type of the value from the source
  3331  // column is an integer, bool or string type T and dest is of type *T,
  3332  // Scan simply assigns the value through the pointer.
  3333  //
  3334  // Scan also converts between string and numeric types, as long as no
  3335  // information would be lost. While Scan stringifies all numbers
  3336  // scanned from numeric database columns into *string, scans into
  3337  // numeric types are checked for overflow. For example, a float64 with
  3338  // value 300 or a string with value "300" can scan into a uint16, but
  3339  // not into a uint8, though float64(255) or "255" can scan into a
  3340  // uint8. One exception is that scans of some float64 numbers to
  3341  // strings may lose information when stringifying. In general, scan
  3342  // floating point columns into *float64.
  3343  //
  3344  // If a dest argument has type *[]byte, Scan saves in that argument a
  3345  // copy of the corresponding data. The copy is owned by the caller and
  3346  // can be modified and held indefinitely. The copy can be avoided by
  3347  // using an argument of type [*RawBytes] instead; see the documentation
  3348  // for [RawBytes] for restrictions on its use.
  3349  //
  3350  // If an argument has type *interface{}, Scan copies the value
  3351  // provided by the underlying driver without conversion. When scanning
  3352  // from a source value of type []byte to *interface{}, a copy of the
  3353  // slice is made and the caller owns the result.
  3354  //
  3355  // Source values of type [time.Time] may be scanned into values of type
  3356  // *time.Time, *interface{}, *string, or *[]byte. When converting to
  3357  // the latter two, [time.RFC3339Nano] is used.
  3358  //
  3359  // Source values of type bool may be scanned into types *bool,
  3360  // *interface{}, *string, *[]byte, or [*RawBytes].
  3361  //
  3362  // For scanning into *bool, the source may be true, false, 1, 0, or
  3363  // string inputs parseable by [strconv.ParseBool].
  3364  //
  3365  // Scan can also convert a cursor returned from a query, such as
  3366  // "select cursor(select * from my_table) from dual", into a
  3367  // [*Rows] value that can itself be scanned from. The parent
  3368  // select query will close any cursor [*Rows] if the parent [*Rows] is closed.
  3369  //
  3370  // If any of the first arguments implementing [Scanner] returns an error,
  3371  // that error will be wrapped in the returned error.
  3372  func (rs *Rows) Scan(dest ...any) error {
  3373  	if rs.closemuScanHold {
  3374  		// This should only be possible if the user calls Scan twice in a row
  3375  		// without calling Next.
  3376  		return fmt.Errorf("sql: Scan called without calling Next (closemuScanHold)")
  3377  	}
  3378  
  3379  	rs.closemu.RLock()
  3380  	rs.raw = rs.raw[:0]
  3381  	err := rs.scanLocked(dest...)
  3382  	if err == nil && scanArgsContainRawBytes(dest) {
  3383  		rs.closemuScanHold = true
  3384  	} else {
  3385  		rs.closemu.RUnlock()
  3386  	}
  3387  	return err
  3388  }
  3389  
  3390  func (rs *Rows) scanLocked(dest ...any) error {
  3391  	if rs.lasterr != nil && rs.lasterr != io.EOF {
  3392  		return rs.lasterr
  3393  	}
  3394  	if rs.closed {
  3395  		return rs.lasterrOrErrLocked(errRowsClosed)
  3396  	}
  3397  
  3398  	if rs.lastcols == nil {
  3399  		return errors.New("sql: Scan called without calling Next")
  3400  	}
  3401  	if len(dest) != len(rs.lastcols) {
  3402  		return fmt.Errorf("sql: expected %d destination arguments in Scan, not %d", len(rs.lastcols), len(dest))
  3403  	}
  3404  
  3405  	for i, sv := range rs.lastcols {
  3406  		err := convertAssignRows(dest[i], sv, rs)
  3407  		if err != nil {
  3408  			return fmt.Errorf(`sql: Scan error on column index %d, name %q: %w`, i, rs.rowsi.Columns()[i], err)
  3409  		}
  3410  	}
  3411  	return nil
  3412  }
  3413  
  3414  // closemuRUnlockIfHeldByScan releases any closemu.RLock held open by a previous
  3415  // call to Scan with *RawBytes.
  3416  func (rs *Rows) closemuRUnlockIfHeldByScan() {
  3417  	if rs.closemuScanHold {
  3418  		rs.closemuScanHold = false
  3419  		rs.closemu.RUnlock()
  3420  	}
  3421  }
  3422  
  3423  func scanArgsContainRawBytes(args []any) bool {
  3424  	for _, a := range args {
  3425  		if _, ok := a.(*RawBytes); ok {
  3426  			return true
  3427  		}
  3428  	}
  3429  	return false
  3430  }
  3431  
  3432  // rowsCloseHook returns a function so tests may install the
  3433  // hook through a test only mutex.
  3434  var rowsCloseHook = func() func(*Rows, *error) { return nil }
  3435  
  3436  // Close closes the [Rows], preventing further enumeration. If [Rows.Next] is called
  3437  // and returns false and there are no further result sets,
  3438  // the [Rows] are closed automatically and it will suffice to check the
  3439  // result of [Rows.Err]. Close is idempotent and does not affect the result of [Rows.Err].
  3440  func (rs *Rows) Close() error {
  3441  	// If the user's calling Close, they're done with their previous row's Scan
  3442  	// results (any RawBytes memory), so we can release the read lock that would
  3443  	// be preventing awaitDone from calling the unexported close before we do so.
  3444  	rs.closemuRUnlockIfHeldByScan()
  3445  
  3446  	return rs.close(nil)
  3447  }
  3448  
  3449  func (rs *Rows) close(err error) error {
  3450  	rs.closemu.Lock()
  3451  	defer rs.closemu.Unlock()
  3452  
  3453  	if rs.closed {
  3454  		return nil
  3455  	}
  3456  	rs.closed = true
  3457  
  3458  	if rs.lasterr == nil {
  3459  		rs.lasterr = err
  3460  	}
  3461  
  3462  	withLock(rs.dc, func() {
  3463  		err = rs.rowsi.Close()
  3464  	})
  3465  	if fn := rowsCloseHook(); fn != nil {
  3466  		fn(rs, &err)
  3467  	}
  3468  	if rs.cancel != nil {
  3469  		rs.cancel()
  3470  	}
  3471  
  3472  	if rs.closeStmt != nil {
  3473  		rs.closeStmt.Close()
  3474  	}
  3475  	rs.releaseConn(err)
  3476  
  3477  	rs.lasterr = rs.lasterrOrErrLocked(err)
  3478  	return err
  3479  }
  3480  
  3481  // Row is the result of calling [DB.QueryRow] to select a single row.
  3482  type Row struct {
  3483  	// One of these two will be non-nil:
  3484  	err  error // deferred error for easy chaining
  3485  	rows *Rows
  3486  }
  3487  
  3488  // Scan copies the columns from the matched row into the values
  3489  // pointed at by dest. See the documentation on [Rows.Scan] for details.
  3490  // If more than one row matches the query,
  3491  // Scan uses the first row and discards the rest. If no row matches
  3492  // the query, Scan returns [ErrNoRows].
  3493  func (r *Row) Scan(dest ...any) error {
  3494  	if r.err != nil {
  3495  		return r.err
  3496  	}
  3497  
  3498  	// TODO(bradfitz): for now we need to defensively clone all
  3499  	// []byte that the driver returned (not permitting
  3500  	// *RawBytes in Rows.Scan), since we're about to close
  3501  	// the Rows in our defer, when we return from this function.
  3502  	// the contract with the driver.Next(...) interface is that it
  3503  	// can return slices into read-only temporary memory that's
  3504  	// only valid until the next Scan/Close. But the TODO is that
  3505  	// for a lot of drivers, this copy will be unnecessary. We
  3506  	// should provide an optional interface for drivers to
  3507  	// implement to say, "don't worry, the []bytes that I return
  3508  	// from Next will not be modified again." (for instance, if
  3509  	// they were obtained from the network anyway) But for now we
  3510  	// don't care.
  3511  	defer r.rows.Close()
  3512  	if scanArgsContainRawBytes(dest) {
  3513  		return errors.New("sql: RawBytes isn't allowed on Row.Scan")
  3514  	}
  3515  
  3516  	if !r.rows.Next() {
  3517  		if err := r.rows.Err(); err != nil {
  3518  			return err
  3519  		}
  3520  		return ErrNoRows
  3521  	}
  3522  	err := r.rows.Scan(dest...)
  3523  	if err != nil {
  3524  		return err
  3525  	}
  3526  	// Make sure the query can be processed to completion with no errors.
  3527  	return r.rows.Close()
  3528  }
  3529  
  3530  // Err provides a way for wrapping packages to check for
  3531  // query errors without calling [Row.Scan].
  3532  // Err returns the error, if any, that was encountered while running the query.
  3533  // If this error is not nil, this error will also be returned from [Row.Scan].
  3534  func (r *Row) Err() error {
  3535  	return r.err
  3536  }
  3537  
  3538  // A Result summarizes an executed SQL command.
  3539  type Result interface {
  3540  	// LastInsertId returns the integer generated by the database
  3541  	// in response to a command. Typically this will be from an
  3542  	// "auto increment" column when inserting a new row. Not all
  3543  	// databases support this feature, and the syntax of such
  3544  	// statements varies.
  3545  	LastInsertId() (int64, error)
  3546  
  3547  	// RowsAffected returns the number of rows affected by an
  3548  	// update, insert, or delete. Not every database or database
  3549  	// driver may support this.
  3550  	RowsAffected() (int64, error)
  3551  }
  3552  
  3553  type driverResult struct {
  3554  	sync.Locker // the *driverConn
  3555  	resi        driver.Result
  3556  }
  3557  
  3558  func (dr driverResult) LastInsertId() (int64, error) {
  3559  	dr.Lock()
  3560  	defer dr.Unlock()
  3561  	return dr.resi.LastInsertId()
  3562  }
  3563  
  3564  func (dr driverResult) RowsAffected() (int64, error) {
  3565  	dr.Lock()
  3566  	defer dr.Unlock()
  3567  	return dr.resi.RowsAffected()
  3568  }
  3569  
  3570  func stack() string {
  3571  	var buf [2 << 10]byte
  3572  	return string(buf[:runtime.Stack(buf[:], false)])
  3573  }
  3574  
  3575  // withLock runs while holding lk.
  3576  func withLock(lk sync.Locker, fn func()) {
  3577  	lk.Lock()
  3578  	defer lk.Unlock() // in case fn panics
  3579  	fn()
  3580  }
  3581  
  3582  // connRequestSet is a set of chan connRequest that's
  3583  // optimized for:
  3584  //
  3585  //   - adding an element
  3586  //   - removing an element (only by the caller who added it)
  3587  //   - taking (get + delete) a random element
  3588  //
  3589  // We previously used a map for this but the take of a random element
  3590  // was expensive, making mapiters. This type avoids a map entirely
  3591  // and just uses a slice.
  3592  type connRequestSet struct {
  3593  	// s are the elements in the set.
  3594  	s []connRequestAndIndex
  3595  }
  3596  
  3597  type connRequestAndIndex struct {
  3598  	// req is the element in the set.
  3599  	req chan connRequest
  3600  
  3601  	// curIdx points to the current location of this element in
  3602  	// connRequestSet.s. It gets set to -1 upon removal.
  3603  	curIdx *int
  3604  }
  3605  
  3606  // CloseAndRemoveAll closes all channels in the set
  3607  // and clears the set.
  3608  func (s *connRequestSet) CloseAndRemoveAll() {
  3609  	for _, v := range s.s {
  3610  		*v.curIdx = -1
  3611  		close(v.req)
  3612  	}
  3613  	s.s = nil
  3614  }
  3615  
  3616  // Len returns the length of the set.
  3617  func (s *connRequestSet) Len() int { return len(s.s) }
  3618  
  3619  // connRequestDelHandle is an opaque handle to delete an
  3620  // item from calling Add.
  3621  type connRequestDelHandle struct {
  3622  	idx *int // pointer to index; or -1 if not in slice
  3623  }
  3624  
  3625  // Add adds v to the set of waiting requests.
  3626  // The returned connRequestDelHandle can be used to remove the item from
  3627  // the set.
  3628  func (s *connRequestSet) Add(v chan connRequest) connRequestDelHandle {
  3629  	idx := len(s.s)
  3630  	// TODO(bradfitz): for simplicity, this always allocates a new int-sized
  3631  	// allocation to store the index. But generally the set will be small and
  3632  	// under a scannable-threshold. As an optimization, we could permit the *int
  3633  	// to be nil when the set is small and should be scanned. This works even if
  3634  	// the set grows over the threshold with delete handles outstanding because
  3635  	// an element can only move to a lower index. So if it starts with a nil
  3636  	// position, it'll always be in a low index and thus scannable. But that
  3637  	// can be done in a follow-up change.
  3638  	idxPtr := &idx
  3639  	s.s = append(s.s, connRequestAndIndex{v, idxPtr})
  3640  	return connRequestDelHandle{idxPtr}
  3641  }
  3642  
  3643  // Delete removes an element from the set.
  3644  //
  3645  // It reports whether the element was deleted. (It can return false if a caller
  3646  // of TakeRandom took it meanwhile, or upon the second call to Delete)
  3647  func (s *connRequestSet) Delete(h connRequestDelHandle) bool {
  3648  	idx := *h.idx
  3649  	if idx < 0 {
  3650  		return false
  3651  	}
  3652  	s.deleteIndex(idx)
  3653  	return true
  3654  }
  3655  
  3656  func (s *connRequestSet) deleteIndex(idx int) {
  3657  	// Mark item as deleted.
  3658  	*(s.s[idx].curIdx) = -1
  3659  	// Copy last element, updating its position
  3660  	// to its new home.
  3661  	if idx < len(s.s)-1 {
  3662  		last := s.s[len(s.s)-1]
  3663  		*last.curIdx = idx
  3664  		s.s[idx] = last
  3665  	}
  3666  	// Zero out last element (for GC) before shrinking the slice.
  3667  	s.s[len(s.s)-1] = connRequestAndIndex{}
  3668  	s.s = s.s[:len(s.s)-1]
  3669  }
  3670  
  3671  // TakeRandom returns and removes a random element from s
  3672  // and reports whether there was one to take. (It returns ok=false
  3673  // if the set is empty.)
  3674  func (s *connRequestSet) TakeRandom() (v chan connRequest, ok bool) {
  3675  	if len(s.s) == 0 {
  3676  		return nil, false
  3677  	}
  3678  	pick := rand.IntN(len(s.s))
  3679  	e := s.s[pick]
  3680  	s.deleteIndex(pick)
  3681  	return e.req, true
  3682  }
  3683  

View as plain text