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

View as plain text