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

View as plain text