Skip to content
Snippets Groups Projects
sql.go 103 KiB
Newer Older
  • Learn to ignore specific revisions
  • 
    // ExecContext executes a query without returning any rows.
    // The args are for any placeholder parameters in the query.
    
    func (c *Conn) ExecContext(ctx context.Context, query string, args ...any) (Result, error) {
    
    	return c.db.execDC(ctx, dc, release, query, args)
    
    }
    
    // QueryContext executes a query that returns rows, typically a SELECT.
    // The args are for any placeholder parameters in the query.
    
    func (c *Conn) QueryContext(ctx context.Context, query string, args ...any) (*Rows, error) {
    
    	return c.db.queryDC(ctx, nil, dc, release, query, args)
    
    }
    
    // QueryRowContext executes a query that is expected to return at most one row.
    // QueryRowContext always returns a non-nil value. Errors are deferred until
    
    // the [*Row.Scan] method is called.
    // If the query selects no rows, the [*Row.Scan] will return [ErrNoRows].
    // Otherwise, the [*Row.Scan] scans the first selected row and discards
    
    func (c *Conn) QueryRowContext(ctx context.Context, query string, args ...any) *Row {
    
    	rows, err := c.QueryContext(ctx, query, args...)
    	return &Row{rows: rows, err: err}
    }
    
    // PrepareContext creates a prepared statement for later queries or executions.
    // Multiple queries or executions may be run concurrently from the
    // returned statement.
    
    // The caller must call the statement's [*Stmt.Close] method
    
    // when the statement is no longer needed.
    //
    // The provided context is used for the preparation of the statement, not for the
    // execution of the statement.
    func (c *Conn) PrepareContext(ctx context.Context, query string) (*Stmt, error) {
    
    	return c.db.prepareDC(ctx, dc, release, c, query)
    
    // Raw executes f exposing the underlying driver connection for the
    // duration of f. The driverConn must not be used outside of f.
    //
    
    // Once f returns and err is not [driver.ErrBadConn], the [Conn] will continue to be usable
    
    // until [Conn.Close] is called.
    
    func (c *Conn) Raw(f func(driverConn any) error) (err error) {
    
    	var dc *driverConn
    	var release releaseConn
    
    	// grabConn takes a context to implement stmtConnGrabber, but the context is not used.
    	dc, release, err = c.grabConn(nil)
    	if err != nil {
    		return
    	}
    	fPanic := true
    	dc.Mutex.Lock()
    	defer func() {
    		dc.Mutex.Unlock()
    
    		// If f panics fPanic will remain true.
    		// Ensure an error is passed to release so the connection
    		// may be discarded.
    		if fPanic {
    			err = driver.ErrBadConn
    		}
    		release(err)
    	}()
    	err = f(dc.ci)
    	fPanic = false
    
    	return
    }
    
    
    // BeginTx starts a transaction.
    //
    // The provided context is used until the transaction is committed or rolled back.
    // If the context is canceled, the sql package will roll back
    
    // the transaction. [Tx.Commit] will return an error if the context provided to
    
    // The provided [TxOptions] is optional and may be nil if defaults should be used.
    
    // If a non-default isolation level is used that the driver doesn't support,
    // an error will be returned.
    func (c *Conn) BeginTx(ctx context.Context, opts *TxOptions) (*Tx, error) {
    
    	return c.db.beginDC(ctx, dc, release, opts)
    
    }
    
    // closemuRUnlockCondReleaseConn read unlocks closemu
    // as the sql operation is done with the dc.
    func (c *Conn) closemuRUnlockCondReleaseConn(err error) {
    	c.closemu.RUnlock()
    
    	if errors.Is(err, driver.ErrBadConn) {
    
    func (c *Conn) close(err error) error {
    
    	if !c.done.CompareAndSwap(false, true) {
    
    		return ErrConnDone
    	}
    
    	// Lock around releasing the driver connection
    	// to ensure all queries have been stopped before doing so.
    	c.closemu.Lock()
    	defer c.closemu.Unlock()
    
    	c.dc.releaseConn(err)
    	c.dc = nil
    	c.db = nil
    	return err
    }
    
    // Close returns the connection to the connection pool.
    
    // All operations after a Close will return with [ErrConnDone].
    
    // Close is safe to call concurrently with other operations and will
    // block until all other operations finish. It may be useful to first
    // cancel any used context and then call close directly after.
    func (c *Conn) Close() error {
    	return c.close(nil)
    }
    
    
    // Tx is an in-progress database transaction.
    
    // A transaction must end with a call to [Tx.Commit] or [Tx.Rollback].
    
    // After a call to [Tx.Commit] or [Tx.Rollback], all operations on the
    // transaction fail with [ErrTxDone].
    
    //
    // The statements prepared for a transaction by calling
    
    // the transaction's [Tx.Prepare] or [Tx.Stmt] methods are closed
    // by the call to [Tx.Commit] or [Tx.Rollback].
    
    type Tx struct {
    
    	// closemu prevents the transaction from closing while there
    	// is an active query. It is held for read during queries
    	// and exclusively during close.
    	closemu sync.RWMutex
    
    
    	// dc is owned exclusively until Commit or Rollback, at which point
    
    	// it's returned with putConn.
    
    	// releaseConn is called once the Tx is closed to release
    	// any held driverConn back to the pool.
    	releaseConn func(error)
    
    
    	// done transitions from false to true exactly once, on Commit
    
    	// or Rollback. once done, all operations fail with
    
    	done atomic.Bool
    
    	// keepConnOnRollback is true if the driver knows
    	// how to reset the connection's session and if need be discard
    	// the connection.
    	keepConnOnRollback bool
    
    
    	// All Stmts prepared for this transaction. These will be closed after the
    
    	// transaction has been committed or rolled back.
    	stmts struct {
    		sync.Mutex
    		v []*Stmt
    	}
    
    	// cancel is called after done transitions from 0 to 1.
    
    	cancel func()
    
    
    	// ctx lives for the life of the transaction.
    	ctx context.Context
    }
    
    
    // awaitDone blocks until the context in Tx is canceled and rolls back
    // the transaction if it's not already done.
    func (tx *Tx) awaitDone() {
    	// Wait for either the transaction to be committed or rolled
    	// back, or for the associated context to be closed.
    	<-tx.ctx.Done()
    
    	// Discard and close the connection used to ensure the
    	// transaction is closed and the resources are released.  This
    	// rollback does nothing if the transaction has already been
    	// committed or rolled back.
    
    	// Do not discard the connection if the connection knows
    	// how to reset the session.
    	discardConnection := !tx.keepConnOnRollback
    	tx.rollback(discardConnection)
    
    	return tx.done.Load()
    
    // ErrTxDone is returned by any operation that is performed on a transaction
    // that has already been committed or rolled back.
    
    var ErrTxDone = errors.New("sql: transaction has already been committed or rolled back")
    
    // close returns the connection to the pool and
    
    // must only be called by Tx.rollback or Tx.Commit while
    
    // tx is already canceled and won't be executed concurrently.
    func (tx *Tx) close(err error) {
    
    	tx.releaseConn(err)
    
    // hookTxGrabConn specifies an optional hook to be called on
    // a successful call to (*Tx).grabConn. For tests.
    var hookTxGrabConn func()
    
    
    func (tx *Tx) grabConn(ctx context.Context) (*driverConn, releaseConn, error) {
    
    Ainar Garipov's avatar
    Ainar Garipov committed
    	// closemu.RLock must come before the check for isDone to prevent the Tx from
    
    	// closing while a query is executing.
    	tx.closemu.RLock()
    
    		tx.closemu.RUnlock()
    		return nil, nil, ErrTxDone
    
    	if hookTxGrabConn != nil { // test hook
    		hookTxGrabConn()
    	}
    
    	return tx.dc, tx.closemuRUnlockRelease, nil
    }
    
    func (tx *Tx) txCtx() context.Context {
    	return tx.ctx
    
    // closemuRUnlockRelease is used as a func(error) method value in
    
    // [DB.ExecContext] and [DB.QueryContext]. Unlocking in the releaseConn keeps
    
    // the driver conn from being returned to the connection pool until
    // the Rows has been closed.
    func (tx *Tx) closemuRUnlockRelease(error) {
    	tx.closemu.RUnlock()
    }
    
    
    // Closes all Stmts prepared for this transaction.
    func (tx *Tx) closePrepared() {
    	tx.stmts.Lock()
    
    // Commit commits the transaction.
    
    func (tx *Tx) Commit() error {
    
    	// Check context first to avoid transaction leak.
    
    	// If put it behind tx.done CompareAndSwap statement, we can't ensure
    
    	// the consistency between tx.done and the real COMMIT operation.
    
    		if tx.done.Load() {
    
    			return ErrTxDone
    		}
    
    	if !tx.done.CompareAndSwap(false, true) {
    
    		return ErrTxDone
    	}
    
    
    	// Cancel the Tx to release any active R-closemu locks.
    	// This is safe to do because tx.done has already transitioned
    	// from 0 to 1. Hold the W-closemu lock prior to rollback
    	// to ensure no other connection has an active query.
    	tx.cancel()
    	tx.closemu.Lock()
    
    	var err error
    	withLock(tx.dc, func() {
    		err = tx.txi.Commit()
    	})
    
    	if !errors.Is(err, driver.ErrBadConn) {
    
    // rollback aborts the transaction and optionally forces the pool to discard
    // the connection.
    func (tx *Tx) rollback(discardConn bool) error {
    
    	if !tx.done.CompareAndSwap(false, true) {
    
    
    	if rollbackHook != nil {
    		rollbackHook()
    	}
    
    	// Cancel the Tx to release any active R-closemu locks.
    	// This is safe to do because tx.done has already transitioned
    	// from 0 to 1. Hold the W-closemu lock prior to rollback
    	// to ensure no other connection has an active query.
    	tx.cancel()
    	tx.closemu.Lock()
    
    	var err error
    	withLock(tx.dc, func() {
    		err = tx.txi.Rollback()
    	})
    
    	if !errors.Is(err, driver.ErrBadConn) {
    
    // Rollback aborts the transaction.
    func (tx *Tx) Rollback() error {
    	return tx.rollback(false)
    }
    
    
    // PrepareContext creates a prepared statement for use within a transaction.
    
    // The returned statement operates within the transaction and will be closed
    // when the transaction has been committed or rolled back.
    
    // To use an existing prepared statement on this transaction, see [Tx.Stmt].
    
    //
    // The provided context will be used for the preparation of the context, not
    
    // for the execution of the returned statement. The returned statement
    // will run in the transaction context.
    func (tx *Tx) PrepareContext(ctx context.Context, query string) (*Stmt, error) {
    
    	stmt, err := tx.db.prepareDC(ctx, dc, release, tx, query)
    
    	if err != nil {
    		return nil, err
    	}
    
    	tx.stmts.Lock()
    	tx.stmts.v = append(tx.stmts.v, stmt)
    	tx.stmts.Unlock()
    
    // Prepare creates a prepared statement for use within a transaction.
    //
    
    // The returned statement operates within the transaction and will be closed
    // when the transaction has been committed or rolled back.
    
    // To use an existing prepared statement on this transaction, see [Tx.Stmt].
    
    // Prepare uses [context.Background] internally; to specify the context, use
    // [Tx.PrepareContext].
    
    func (tx *Tx) Prepare(query string) (*Stmt, error) {
    
    	return tx.PrepareContext(context.Background(), query)
    
    }
    
    // StmtContext returns a transaction-specific prepared statement from
    
    Russ Cox's avatar
    Russ Cox committed
    //
    //	updateMoney, err := db.Prepare("UPDATE balance SET money=money+? WHERE id=?")
    //	...
    //	tx, err := db.Begin()
    //	...
    //	res, err := tx.StmtContext(ctx, updateMoney).Exec(123.45, 98293203)
    
    // The provided context is used for the preparation of the statement, not for the
    // execution of the statement.
    //
    
    // The returned statement operates within the transaction and will be closed
    // when the transaction has been committed or rolled back.
    
    func (tx *Tx) StmtContext(ctx context.Context, stmt *Stmt) *Stmt {
    
    	dc, release, err := tx.grabConn(ctx)
    	if err != nil {
    		return &Stmt{stickyErr: err}
    	}
    	defer release(nil)
    
    	if tx.db != stmt.db {
    		return &Stmt{stickyErr: errors.New("sql: Tx.Stmt: statement from different database used")}
    	}
    
    	var parentStmt *Stmt
    	stmt.mu.Lock()
    
    		// If the statement has been closed or already belongs to a
    		// transaction, we can't reuse it in this connection.
    		// Since tx.StmtContext should never need to be called with a
    		// Stmt already belonging to tx, we ignore this edge case and
    		// re-prepare the statement in this case. No need to add
    		// code-complexity for this.
    		stmt.mu.Unlock()
    		withLock(dc, func() {
    			si, err = ctxDriverPrepare(ctx, dc.ci, stmt.query)
    		})
    		if err != nil {
    			return &Stmt{stickyErr: err}
    		}
    	} else {
    		stmt.removeClosedStmtLocked()
    		// See if the statement has already been prepared on this connection,
    		// and reuse it if possible.
    		for _, v := range stmt.css {
    			if v.dc == dc {
    				si = v.ds.si
    				break
    			}
    		}
    
    		stmt.mu.Unlock()
    
    		if si == nil {
    
    			withLock(dc, func() {
    				ds, err = stmt.prepareOnConnLocked(ctx, dc)
    			})
    
    			if err != nil {
    				return &Stmt{stickyErr: err}
    			}
    
    		parentStmt: parentStmt,
    		query:      stmt.query,
    	}
    	if parentStmt != nil {
    		tx.db.addDep(parentStmt, txs)
    
    	tx.stmts.Lock()
    	tx.stmts.v = append(tx.stmts.v, txs)
    	tx.stmts.Unlock()
    	return txs
    
    // Stmt returns a transaction-specific prepared statement from
    // an existing statement.
    //
    // Example:
    
    Russ Cox's avatar
    Russ Cox committed
    //
    //	updateMoney, err := db.Prepare("UPDATE balance SET money=money+? WHERE id=?")
    //	...
    //	tx, err := db.Begin()
    //	...
    //	res, err := tx.Stmt(updateMoney).Exec(123.45, 98293203)
    
    // The returned statement operates within the transaction and will be closed
    // when the transaction has been committed or rolled back.
    
    // Stmt uses [context.Background] internally; to specify the context, use
    // [Tx.StmtContext].
    
    func (tx *Tx) Stmt(stmt *Stmt) *Stmt {
    
    	return tx.StmtContext(context.Background(), stmt)
    
    }
    
    // ExecContext executes a query that doesn't return rows.
    
    // For example: an INSERT and UPDATE.
    
    func (tx *Tx) ExecContext(ctx context.Context, query string, args ...any) (Result, error) {
    
    	if err != nil {
    		return nil, err
    	}
    
    	return tx.db.execDC(ctx, dc, release, query, args)
    
    // Exec executes a query that doesn't return rows.
    // For example: an INSERT and UPDATE.
    
    // Exec uses [context.Background] internally; to specify the context, use
    // [Tx.ExecContext].
    
    func (tx *Tx) Exec(query string, args ...any) (Result, error) {
    
    	return tx.ExecContext(context.Background(), query, args...)
    
    }
    
    // QueryContext executes a query that returns rows, typically a SELECT.
    
    func (tx *Tx) QueryContext(ctx context.Context, query string, args ...any) (*Rows, error) {
    
    	if err != nil {
    		return nil, err
    
    	return tx.db.queryDC(ctx, tx.ctx, dc, release, query, args)
    
    }
    
    // Query executes a query that returns rows, typically a SELECT.
    
    // Query uses [context.Background] internally; to specify the context, use
    // [Tx.QueryContext].
    
    func (tx *Tx) Query(query string, args ...any) (*Rows, error) {
    
    	return tx.QueryContext(context.Background(), query, args...)
    
    }
    
    // QueryRowContext executes a query that is expected to return at most one row.
    // QueryRowContext always returns a non-nil value. Errors are deferred until
    
    // [Row]'s Scan method is called.
    
    // If the query selects no rows, the [*Row.Scan] will return [ErrNoRows].
    // Otherwise, the [*Row.Scan] scans the first selected row and discards
    
    func (tx *Tx) QueryRowContext(ctx context.Context, query string, args ...any) *Row {
    
    	rows, err := tx.QueryContext(ctx, query, args...)
    	return &Row{rows: rows, err: err}
    
    }
    
    // QueryRow executes a query that is expected to return at most one row.
    
    Evan Shaw's avatar
    Evan Shaw committed
    // QueryRow always returns a non-nil value. Errors are deferred until
    
    // [Row]'s Scan method is called.
    
    // If the query selects no rows, the [*Row.Scan] will return [ErrNoRows].
    // Otherwise, the [*Row.Scan] scans the first selected row and discards
    
    // QueryRow uses [context.Background] internally; to specify the context, use
    // [Tx.QueryRowContext].
    
    func (tx *Tx) QueryRow(query string, args ...any) *Row {
    
    	return tx.QueryRowContext(context.Background(), query, args...)
    
    }
    
    // connStmt is a prepared statement on a particular connection.
    type connStmt struct {
    
    // stmtConnGrabber represents a Tx or Conn that will return the underlying
    // driverConn and release function.
    type stmtConnGrabber interface {
    	// grabConn returns the driverConn and the associated release function
    	// that must be called when the operation completes.
    	grabConn(context.Context) (*driverConn, releaseConn, error)
    
    	// txCtx returns the transaction context if available.
    	// The returned context should be selected on along with
    	// any query context when awaiting a cancel.
    	txCtx() context.Context
    }
    
    var (
    	_ stmtConnGrabber = &Tx{}
    	_ stmtConnGrabber = &Conn{}
    )
    
    
    // Stmt is a prepared statement.
    // A Stmt is safe for concurrent use by multiple goroutines.
    
    // If a Stmt is prepared on a [Tx] or [Conn], it will be bound to a single
    // underlying connection forever. If the [Tx] or [Conn] closes, the Stmt will
    
    // become unusable and all operations will return an error.
    
    // If a Stmt is prepared on a [DB], it will remain usable for the lifetime of the
    // [DB]. When the Stmt needs to execute on a new underlying connection, it will
    
    // prepare itself on the new connection automatically.
    
    type Stmt struct {
    	// Immutable:
    
    	db        *DB    // where we came from
    	query     string // that created the Stmt
    	stickyErr error  // if non-nil, this error is returned for all operations
    
    	closemu sync.RWMutex // held exclusively during close, for read otherwise.
    
    
    	// If Stmt is prepared on a Tx or Conn then cg is present and will
    	// only ever grab a connection from cg.
    	// If cg is nil then the Stmt must grab an arbitrary connection
    	// from db and determine if it must prepare the stmt again by
    	// inspecting css.
    	cg   stmtConnGrabber
    	cgds *driverStmt
    
    	// parentStmt is set when a transaction-specific statement
    	// is requested from an identical statement prepared on the same
    	// conn. parentStmt is used to track the dependency of this statement
    	// on its originating ("parent") statement so that parentStmt may
    	// be closed by the user without them having to know whether or not
    	// any transactions are still using it.
    	parentStmt *Stmt
    
    
    	mu     sync.Mutex // protects the rest of the fields
    
    	// css is a list of underlying driver statement interfaces
    
    	// that are valid on particular connections. This is only
    
    	// used if cg == nil and one is found that has idle
    	// connections. If cg != nil, cgds is always used.
    
    
    	// lastNumClosed is copied from db.numClosed when Stmt is created
    	// without tx and closed connections in css are removed.
    	lastNumClosed uint64
    
    // ExecContext executes a prepared statement with the given arguments and
    
    // returns a [Result] summarizing the effect of the statement.
    
    func (s *Stmt) ExecContext(ctx context.Context, args ...any) (Result, error) {
    
    	s.closemu.RLock()
    	defer s.closemu.RUnlock()
    
    	err := s.db.retry(func(strategy connReuseStrategy) error {
    
    		dc, releaseConn, ds, err := s.connStmt(ctx, strategy)
    
    		res, err = resultFromStatement(ctx, dc.ci, ds, args...)
    
    // Exec executes a prepared statement with the given arguments and
    
    // returns a [Result] summarizing the effect of the statement.
    
    // Exec uses [context.Background] internally; to specify the context, use
    // [Stmt.ExecContext].
    
    func (s *Stmt) Exec(args ...any) (Result, error) {
    
    	return s.ExecContext(context.Background(), args...)
    }
    
    
    func resultFromStatement(ctx context.Context, ci driver.Conn, ds *driverStmt, args ...any) (Result, error) {
    
    	ds.Lock()
    	defer ds.Unlock()
    
    	dargs, err := driverArgsConnLocked(ci, ds, args)
    
    	resi, err := ctxDriverStmtExec(ctx, ds.si, dargs)
    
    	if err != nil {
    		return nil, err
    	}
    
    	return driverResult{ds.Locker, resi}, nil
    
    // removeClosedStmtLocked removes closed conns in s.css.
    //
    // To avoid lock contention on DB.mu, we do it only when
    // s.db.numClosed - s.lastNum is large enough.
    func (s *Stmt) removeClosedStmtLocked() {
    	t := len(s.css)/2 + 1
    	if t > 10 {
    		t = 10
    	}
    
    	dbClosed := s.db.numClosed.Load()
    
    	if dbClosed-s.lastNumClosed < uint64(t) {
    		return
    	}
    
    	s.db.mu.Lock()
    	for i := 0; i < len(s.css); i++ {
    		if s.css[i].dc.dbmuClosed {
    			s.css[i] = s.css[len(s.css)-1]
    
    			// Zero out the last element (for GC) before shrinking the slice.
    			s.css[len(s.css)-1] = connStmt{}
    
    			s.css = s.css[:len(s.css)-1]
    			i--
    		}
    	}
    	s.db.mu.Unlock()
    	s.lastNumClosed = dbClosed
    }
    
    
    // connStmt returns a free driver connection on which to execute the
    // statement, a function to call to release the connection, and a
    // statement bound to that connection.
    
    func (s *Stmt) connStmt(ctx context.Context, strategy connReuseStrategy) (dc *driverConn, releaseConn func(error), ds *driverStmt, err error) {
    
    	if err = s.stickyErr; err != nil {
    		return
    
    	s.mu.Lock()
    	if s.closed {
    
    Brad Fitzpatrick's avatar
    Brad Fitzpatrick committed
    		err = errors.New("sql: statement is closed")
    
    	// In a transaction or connection, we always use the connection that the
    
    	// stmt was created on.
    
    		dc, releaseConn, err = s.cg.grabConn(ctx) // blocks, waiting for the connection.
    
    	s.removeClosedStmtLocked()
    
    	if err != nil {
    		return nil, nil, nil, err
    	}
    
    	s.mu.Lock()
    	for _, v := range s.css {
    		if v.dc == dc {
    			s.mu.Unlock()
    
    			return dc, dc.releaseConn, v.ds, nil
    
    	s.mu.Unlock()
    
    	// No luck; we need to prepare the statement on this connection
    
    		ds, err = s.prepareOnConnLocked(ctx, dc)
    
    		dc.releaseConn(err)
    
    
    	return dc, dc.releaseConn, ds, nil
    }
    
    // prepareOnConnLocked prepares the query in Stmt s on dc and adds it to the list of
    // open connStmt on the statement. It assumes the caller is holding the lock on dc.
    func (s *Stmt) prepareOnConnLocked(ctx context.Context, dc *driverConn) (*driverStmt, error) {
    
    	si, err := dc.prepareLocked(ctx, s.cg, s.query)
    
    	if err != nil {
    		return nil, err
    	}
    	cs := connStmt{dc, si}
    
    	s.mu.Lock()
    	s.css = append(s.css, cs)
    	s.mu.Unlock()
    
    // QueryContext executes a prepared query statement with the given arguments
    
    // and returns the query results as a [*Rows].
    
    func (s *Stmt) QueryContext(ctx context.Context, args ...any) (*Rows, error) {
    
    	s.closemu.RLock()
    	defer s.closemu.RUnlock()
    
    
    	var rows *Rows
    
    	err := s.db.retry(func(strategy connReuseStrategy) error {
    
    		dc, releaseConn, ds, err := s.connStmt(ctx, strategy)
    
    		rowsi, err = rowsiFromStatement(ctx, dc.ci, ds, args...)
    
    		if err == nil {
    			// Note: ownership of ci passes to the *Rows, to be freed
    			// with releaseConn.
    
    				dc:    dc,
    				rowsi: rowsi,
    				// releaseConn set below
    			}
    
    			// addDep must be added before initContextClose or it could attempt
    			// to removeDep before it has been added.
    
    
    			// releaseConn must be set before initContextClose or it could
    			// release the connection before it is set.
    
    			rows.releaseConn = func(err error) {
    				releaseConn(err)
    				s.db.removeDep(s, rows)
    			}
    
    			var txctx context.Context
    
    			}
    			rows.initContextClose(ctx, txctx)
    
    // Query executes a prepared query statement with the given arguments
    // and returns the query results as a *Rows.
    
    // Query uses [context.Background] internally; to specify the context, use
    // [Stmt.QueryContext].
    
    func (s *Stmt) Query(args ...any) (*Rows, error) {
    
    	return s.QueryContext(context.Background(), args...)
    }
    
    
    func rowsiFromStatement(ctx context.Context, ci driver.Conn, ds *driverStmt, args ...any) (driver.Rows, error) {
    
    	dargs, err := driverArgsConnLocked(ci, ds, args)
    	if err != nil {
    		return nil, err
    	}
    
    	return ctxDriverStmtQuery(ctx, ds.si, dargs)
    
    // QueryRowContext executes a prepared query statement with the given arguments.
    
    // If an error occurs during the execution of the statement, that error will
    
    // be returned by a call to Scan on the returned [*Row], which is always non-nil.
    // If the query selects no rows, the [*Row.Scan] will return [ErrNoRows].
    // Otherwise, the [*Row.Scan] scans the first selected row and discards
    
    func (s *Stmt) QueryRowContext(ctx context.Context, args ...any) *Row {
    
    	rows, err := s.QueryContext(ctx, args...)
    
    	if err != nil {
    		return &Row{err: err}
    	}
    	return &Row{rows: rows}
    }
    
    
    // QueryRow executes a prepared query statement with the given arguments.
    // If an error occurs during the execution of the statement, that error will
    
    // be returned by a call to Scan on the returned [*Row], which is always non-nil.
    // If the query selects no rows, the [*Row.Scan] will return [ErrNoRows].
    // Otherwise, the [*Row.Scan] scans the first selected row and discards
    
    // the rest.
    //
    // Example usage:
    //
    
    Russ Cox's avatar
    Russ Cox committed
    //	var name string
    //	err := nameByUseridStmt.QueryRow(id).Scan(&name)
    
    // QueryRow uses [context.Background] internally; to specify the context, use
    // [Stmt.QueryRowContext].
    
    func (s *Stmt) QueryRow(args ...any) *Row {
    
    	return s.QueryRowContext(context.Background(), args...)
    }
    
    
    // Close closes the statement.
    
    func (s *Stmt) Close() error {
    
    	s.closemu.Lock()
    	defer s.closemu.Unlock()
    
    
    	if s.stickyErr != nil {
    		return s.stickyErr
    	}
    
    	s.mu.Lock()
    	if s.closed {
    
    		s.mu.Unlock()
    
    		return nil
    	}
    	s.closed = true
    
    	if s.parentStmt != nil {
    		// If parentStmt is set, we must not close s.txds since it's stored
    		// in the css array of the parentStmt.
    		return s.db.removeDep(s.parentStmt, s)
    	}
    
    }
    
    func (s *Stmt) finalClose() error {
    
    	s.mu.Lock()
    	defer s.mu.Unlock()
    	if s.css != nil {
    		for _, v := range s.css {
    
    			s.db.noteUnusedDriverStatement(v.dc, v.ds)
    			v.dc.removeOpenStmt(v.ds)
    
    		}
    		s.css = nil
    
    	}
    	return nil
    }
    
    // Rows is the result of a query. Its cursor starts before the first row
    
    // of the result set. Use [Rows.Next] to advance from row to row.
    
    type Rows struct {
    
    	dc          *driverConn // owned; must call releaseConn when closed to release
    
    	cancel      func()      // called when Rows is closed, may be nil.
    	closeStmt   *driverStmt // if non-nil, statement to Close on close
    
    	contextDone atomic.Pointer[error] // error that awaitDone saw; set before close attempt
    
    
    	// closemu prevents Rows from closing while there
    	// is an active streaming result. It is held for read during non-close operations
    	// and exclusively during close.
    	//
    	// closemu guards lasterr and closed.
    	closemu sync.RWMutex
    	closed  bool
    	lasterr error // non-nil only if closed is true
    
    	// lastcols is only used in Scan, Next, and NextResultSet which are expected
    
    Kevin Burke's avatar
    Kevin Burke committed
    	// not to be called concurrently.
    
    	// raw is a buffer for RawBytes that persists between Scan calls.
    	// This is used when the driver returns a mismatched type that requires
    	// a cloning allocation. For example, if the driver returns a *string and
    	// the user is scanning into a *RawBytes, we need to copy the string.
    	// The raw buffer here lets us reuse the memory for that copy across Scan calls.
    	raw []byte
    
    
    	// closemuScanHold is whether the previous call to Scan kept closemu RLock'ed
    	// without unlocking it. It does that when the user passes a *RawBytes scan
    	// target. In that case, we need to prevent awaitDone from closing the Rows
    	// while the user's still using the memory. See go.dev/issue/60304.
    	//
    	// It is only used by Scan, Next, and NextResultSet which are expected
    	// not to be called concurrently.
    	closemuScanHold bool
    
    
    	// hitEOF is whether Next hit the end of the rows without
    	// encountering an error. It's set in Next before
    	// returning. It's only used by Next and Err which are
    	// expected not to be called concurrently.
    	hitEOF bool
    
    // lasterrOrErrLocked returns either lasterr or the provided err.
    // rs.closemu must be read-locked.
    func (rs *Rows) lasterrOrErrLocked(err error) error {
    	if rs.lasterr != nil && rs.lasterr != io.EOF {
    		return rs.lasterr
    	}
    	return err
    }
    
    
    // bypassRowsAwaitDone is only used for testing.
    // If true, it will not close the Rows automatically from the context.
    var bypassRowsAwaitDone = false
    
    
    func (rs *Rows) initContextClose(ctx, txctx context.Context) {
    
    	if ctx.Done() == nil && (txctx == nil || txctx.Done() == nil) {
    		return
    	}
    
    	closectx, cancel := context.WithCancel(ctx)
    	rs.cancel = cancel
    	go rs.awaitDone(ctx, txctx, closectx)
    
    // awaitDone blocks until ctx, txctx, or closectx is canceled.
    
    // The ctx is provided from the query context.
    
    // If the query was issued in a transaction, the transaction's context
    
    // is also provided in txctx, to ensure Rows is closed if the Tx is closed.
    // The closectx is closed by an explicit call to rs.Close.
    func (rs *Rows) awaitDone(ctx, txctx, closectx context.Context) {
    
    	var txctxDone <-chan struct{}
    	if txctx != nil {
    		txctxDone = txctx.Done()
    	}
    	select {
    	case <-ctx.Done():
    
    		err := ctx.Err()
    		rs.contextDone.Store(&err)
    
    		err := txctx.Err()
    		rs.contextDone.Store(&err)
    
    	case <-closectx.Done():
    		// rs.cancel was called via Close(); don't store this into contextDone
    		// to ensure Err() is unaffected.