mirror of
https://github.com/superseriousbusiness/gotosocial.git
synced 2024-11-27 14:16:39 +00:00
ca729aa4a0
Bumps [github.com/jackc/pgx/v5](https://github.com/jackc/pgx) from 5.6.0 to 5.7.1. - [Changelog](https://github.com/jackc/pgx/blob/master/CHANGELOG.md) - [Commits](https://github.com/jackc/pgx/compare/v5.6.0...v5.7.1) --- updated-dependencies: - dependency-name: github.com/jackc/pgx/v5 dependency-type: direct:production update-type: version-update:semver-minor ... Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
1414 lines
44 KiB
Go
1414 lines
44 KiB
Go
package pgx
|
|
|
|
import (
|
|
"context"
|
|
"crypto/sha256"
|
|
"database/sql"
|
|
"encoding/hex"
|
|
"errors"
|
|
"fmt"
|
|
"strconv"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/jackc/pgx/v5/internal/sanitize"
|
|
"github.com/jackc/pgx/v5/internal/stmtcache"
|
|
"github.com/jackc/pgx/v5/pgconn"
|
|
"github.com/jackc/pgx/v5/pgtype"
|
|
)
|
|
|
|
// ConnConfig contains all the options used to establish a connection. It must be created by ParseConfig and
|
|
// then it can be modified. A manually initialized ConnConfig will cause ConnectConfig to panic.
|
|
type ConnConfig struct {
|
|
pgconn.Config
|
|
|
|
Tracer QueryTracer
|
|
|
|
// Original connection string that was parsed into config.
|
|
connString string
|
|
|
|
// StatementCacheCapacity is maximum size of the statement cache used when executing a query with "cache_statement"
|
|
// query exec mode.
|
|
StatementCacheCapacity int
|
|
|
|
// DescriptionCacheCapacity is the maximum size of the description cache used when executing a query with
|
|
// "cache_describe" query exec mode.
|
|
DescriptionCacheCapacity int
|
|
|
|
// DefaultQueryExecMode controls the default mode for executing queries. By default pgx uses the extended protocol
|
|
// and automatically prepares and caches prepared statements. However, this may be incompatible with proxies such as
|
|
// PGBouncer. In this case it may be preferable to use QueryExecModeExec or QueryExecModeSimpleProtocol. The same
|
|
// functionality can be controlled on a per query basis by passing a QueryExecMode as the first query argument.
|
|
DefaultQueryExecMode QueryExecMode
|
|
|
|
createdByParseConfig bool // Used to enforce created by ParseConfig rule.
|
|
}
|
|
|
|
// ParseConfigOptions contains options that control how a config is built such as getsslpassword.
|
|
type ParseConfigOptions struct {
|
|
pgconn.ParseConfigOptions
|
|
}
|
|
|
|
// Copy returns a deep copy of the config that is safe to use and modify.
|
|
// The only exception is the tls.Config:
|
|
// according to the tls.Config docs it must not be modified after creation.
|
|
func (cc *ConnConfig) Copy() *ConnConfig {
|
|
newConfig := new(ConnConfig)
|
|
*newConfig = *cc
|
|
newConfig.Config = *newConfig.Config.Copy()
|
|
return newConfig
|
|
}
|
|
|
|
// ConnString returns the connection string as parsed by pgx.ParseConfig into pgx.ConnConfig.
|
|
func (cc *ConnConfig) ConnString() string { return cc.connString }
|
|
|
|
// Conn is a PostgreSQL connection handle. It is not safe for concurrent usage. Use a connection pool to manage access
|
|
// to multiple database connections from multiple goroutines.
|
|
type Conn struct {
|
|
pgConn *pgconn.PgConn
|
|
config *ConnConfig // config used when establishing this connection
|
|
preparedStatements map[string]*pgconn.StatementDescription
|
|
statementCache stmtcache.Cache
|
|
descriptionCache stmtcache.Cache
|
|
|
|
queryTracer QueryTracer
|
|
batchTracer BatchTracer
|
|
copyFromTracer CopyFromTracer
|
|
prepareTracer PrepareTracer
|
|
|
|
notifications []*pgconn.Notification
|
|
|
|
doneChan chan struct{}
|
|
closedChan chan error
|
|
|
|
typeMap *pgtype.Map
|
|
|
|
wbuf []byte
|
|
eqb ExtendedQueryBuilder
|
|
}
|
|
|
|
// Identifier a PostgreSQL identifier or name. Identifiers can be composed of
|
|
// multiple parts such as ["schema", "table"] or ["table", "column"].
|
|
type Identifier []string
|
|
|
|
// Sanitize returns a sanitized string safe for SQL interpolation.
|
|
func (ident Identifier) Sanitize() string {
|
|
parts := make([]string, len(ident))
|
|
for i := range ident {
|
|
s := strings.ReplaceAll(ident[i], string([]byte{0}), "")
|
|
parts[i] = `"` + strings.ReplaceAll(s, `"`, `""`) + `"`
|
|
}
|
|
return strings.Join(parts, ".")
|
|
}
|
|
|
|
var (
|
|
// ErrNoRows occurs when rows are expected but none are returned.
|
|
ErrNoRows = newProxyErr(sql.ErrNoRows, "no rows in result set")
|
|
// ErrTooManyRows occurs when more rows than expected are returned.
|
|
ErrTooManyRows = errors.New("too many rows in result set")
|
|
)
|
|
|
|
func newProxyErr(background error, msg string) error {
|
|
return &proxyError{
|
|
msg: msg,
|
|
background: background,
|
|
}
|
|
}
|
|
|
|
type proxyError struct {
|
|
msg string
|
|
background error
|
|
}
|
|
|
|
func (err *proxyError) Error() string { return err.msg }
|
|
|
|
func (err *proxyError) Unwrap() error { return err.background }
|
|
|
|
var (
|
|
errDisabledStatementCache = fmt.Errorf("cannot use QueryExecModeCacheStatement with disabled statement cache")
|
|
errDisabledDescriptionCache = fmt.Errorf("cannot use QueryExecModeCacheDescribe with disabled description cache")
|
|
)
|
|
|
|
// Connect establishes a connection with a PostgreSQL server with a connection string. See
|
|
// pgconn.Connect for details.
|
|
func Connect(ctx context.Context, connString string) (*Conn, error) {
|
|
connConfig, err := ParseConfig(connString)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return connect(ctx, connConfig)
|
|
}
|
|
|
|
// ConnectWithOptions behaves exactly like Connect with the addition of options. At the present options is only used to
|
|
// provide a GetSSLPassword function.
|
|
func ConnectWithOptions(ctx context.Context, connString string, options ParseConfigOptions) (*Conn, error) {
|
|
connConfig, err := ParseConfigWithOptions(connString, options)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return connect(ctx, connConfig)
|
|
}
|
|
|
|
// ConnectConfig establishes a connection with a PostgreSQL server with a configuration struct.
|
|
// connConfig must have been created by ParseConfig.
|
|
func ConnectConfig(ctx context.Context, connConfig *ConnConfig) (*Conn, error) {
|
|
// In general this improves safety. In particular avoid the config.Config.OnNotification mutation from affecting other
|
|
// connections with the same config. See https://github.com/jackc/pgx/issues/618.
|
|
connConfig = connConfig.Copy()
|
|
|
|
return connect(ctx, connConfig)
|
|
}
|
|
|
|
// ParseConfigWithOptions behaves exactly as ParseConfig does with the addition of options. At the present options is
|
|
// only used to provide a GetSSLPassword function.
|
|
func ParseConfigWithOptions(connString string, options ParseConfigOptions) (*ConnConfig, error) {
|
|
config, err := pgconn.ParseConfigWithOptions(connString, options.ParseConfigOptions)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
statementCacheCapacity := 512
|
|
if s, ok := config.RuntimeParams["statement_cache_capacity"]; ok {
|
|
delete(config.RuntimeParams, "statement_cache_capacity")
|
|
n, err := strconv.ParseInt(s, 10, 32)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("cannot parse statement_cache_capacity: %w", err)
|
|
}
|
|
statementCacheCapacity = int(n)
|
|
}
|
|
|
|
descriptionCacheCapacity := 512
|
|
if s, ok := config.RuntimeParams["description_cache_capacity"]; ok {
|
|
delete(config.RuntimeParams, "description_cache_capacity")
|
|
n, err := strconv.ParseInt(s, 10, 32)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("cannot parse description_cache_capacity: %w", err)
|
|
}
|
|
descriptionCacheCapacity = int(n)
|
|
}
|
|
|
|
defaultQueryExecMode := QueryExecModeCacheStatement
|
|
if s, ok := config.RuntimeParams["default_query_exec_mode"]; ok {
|
|
delete(config.RuntimeParams, "default_query_exec_mode")
|
|
switch s {
|
|
case "cache_statement":
|
|
defaultQueryExecMode = QueryExecModeCacheStatement
|
|
case "cache_describe":
|
|
defaultQueryExecMode = QueryExecModeCacheDescribe
|
|
case "describe_exec":
|
|
defaultQueryExecMode = QueryExecModeDescribeExec
|
|
case "exec":
|
|
defaultQueryExecMode = QueryExecModeExec
|
|
case "simple_protocol":
|
|
defaultQueryExecMode = QueryExecModeSimpleProtocol
|
|
default:
|
|
return nil, fmt.Errorf("invalid default_query_exec_mode: %s", s)
|
|
}
|
|
}
|
|
|
|
connConfig := &ConnConfig{
|
|
Config: *config,
|
|
createdByParseConfig: true,
|
|
StatementCacheCapacity: statementCacheCapacity,
|
|
DescriptionCacheCapacity: descriptionCacheCapacity,
|
|
DefaultQueryExecMode: defaultQueryExecMode,
|
|
connString: connString,
|
|
}
|
|
|
|
return connConfig, nil
|
|
}
|
|
|
|
// ParseConfig creates a ConnConfig from a connection string. ParseConfig handles all options that [pgconn.ParseConfig]
|
|
// does. In addition, it accepts the following options:
|
|
//
|
|
// - default_query_exec_mode.
|
|
// Possible values: "cache_statement", "cache_describe", "describe_exec", "exec", and "simple_protocol". See
|
|
// QueryExecMode constant documentation for the meaning of these values. Default: "cache_statement".
|
|
//
|
|
// - statement_cache_capacity.
|
|
// The maximum size of the statement cache used when executing a query with "cache_statement" query exec mode.
|
|
// Default: 512.
|
|
//
|
|
// - description_cache_capacity.
|
|
// The maximum size of the description cache used when executing a query with "cache_describe" query exec mode.
|
|
// Default: 512.
|
|
func ParseConfig(connString string) (*ConnConfig, error) {
|
|
return ParseConfigWithOptions(connString, ParseConfigOptions{})
|
|
}
|
|
|
|
// connect connects to a database. connect takes ownership of config. The caller must not use or access it again.
|
|
func connect(ctx context.Context, config *ConnConfig) (c *Conn, err error) {
|
|
if connectTracer, ok := config.Tracer.(ConnectTracer); ok {
|
|
ctx = connectTracer.TraceConnectStart(ctx, TraceConnectStartData{ConnConfig: config})
|
|
defer func() {
|
|
connectTracer.TraceConnectEnd(ctx, TraceConnectEndData{Conn: c, Err: err})
|
|
}()
|
|
}
|
|
|
|
// Default values are set in ParseConfig. Enforce initial creation by ParseConfig rather than setting defaults from
|
|
// zero values.
|
|
if !config.createdByParseConfig {
|
|
panic("config must be created by ParseConfig")
|
|
}
|
|
|
|
c = &Conn{
|
|
config: config,
|
|
typeMap: pgtype.NewMap(),
|
|
queryTracer: config.Tracer,
|
|
}
|
|
|
|
if t, ok := c.queryTracer.(BatchTracer); ok {
|
|
c.batchTracer = t
|
|
}
|
|
if t, ok := c.queryTracer.(CopyFromTracer); ok {
|
|
c.copyFromTracer = t
|
|
}
|
|
if t, ok := c.queryTracer.(PrepareTracer); ok {
|
|
c.prepareTracer = t
|
|
}
|
|
|
|
// Only install pgx notification system if no other callback handler is present.
|
|
if config.Config.OnNotification == nil {
|
|
config.Config.OnNotification = c.bufferNotifications
|
|
}
|
|
|
|
c.pgConn, err = pgconn.ConnectConfig(ctx, &config.Config)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
c.preparedStatements = make(map[string]*pgconn.StatementDescription)
|
|
c.doneChan = make(chan struct{})
|
|
c.closedChan = make(chan error)
|
|
c.wbuf = make([]byte, 0, 1024)
|
|
|
|
if c.config.StatementCacheCapacity > 0 {
|
|
c.statementCache = stmtcache.NewLRUCache(c.config.StatementCacheCapacity)
|
|
}
|
|
|
|
if c.config.DescriptionCacheCapacity > 0 {
|
|
c.descriptionCache = stmtcache.NewLRUCache(c.config.DescriptionCacheCapacity)
|
|
}
|
|
|
|
return c, nil
|
|
}
|
|
|
|
// Close closes a connection. It is safe to call Close on an already closed
|
|
// connection.
|
|
func (c *Conn) Close(ctx context.Context) error {
|
|
if c.IsClosed() {
|
|
return nil
|
|
}
|
|
|
|
err := c.pgConn.Close(ctx)
|
|
return err
|
|
}
|
|
|
|
// Prepare creates a prepared statement with name and sql. sql can contain placeholders for bound parameters. These
|
|
// placeholders are referenced positionally as $1, $2, etc. name can be used instead of sql with Query, QueryRow, and
|
|
// Exec to execute the statement. It can also be used with Batch.Queue.
|
|
//
|
|
// The underlying PostgreSQL identifier for the prepared statement will be name if name != sql or a digest of sql if
|
|
// name == sql.
|
|
//
|
|
// Prepare is idempotent; i.e. it is safe to call Prepare multiple times with the same name and sql arguments. This
|
|
// allows a code path to Prepare and Query/Exec without concern for if the statement has already been prepared.
|
|
func (c *Conn) Prepare(ctx context.Context, name, sql string) (sd *pgconn.StatementDescription, err error) {
|
|
if c.prepareTracer != nil {
|
|
ctx = c.prepareTracer.TracePrepareStart(ctx, c, TracePrepareStartData{Name: name, SQL: sql})
|
|
}
|
|
|
|
if name != "" {
|
|
var ok bool
|
|
if sd, ok = c.preparedStatements[name]; ok && sd.SQL == sql {
|
|
if c.prepareTracer != nil {
|
|
c.prepareTracer.TracePrepareEnd(ctx, c, TracePrepareEndData{AlreadyPrepared: true})
|
|
}
|
|
return sd, nil
|
|
}
|
|
}
|
|
|
|
if c.prepareTracer != nil {
|
|
defer func() {
|
|
c.prepareTracer.TracePrepareEnd(ctx, c, TracePrepareEndData{Err: err})
|
|
}()
|
|
}
|
|
|
|
var psName, psKey string
|
|
if name == sql {
|
|
digest := sha256.Sum256([]byte(sql))
|
|
psName = "stmt_" + hex.EncodeToString(digest[0:24])
|
|
psKey = sql
|
|
} else {
|
|
psName = name
|
|
psKey = name
|
|
}
|
|
|
|
sd, err = c.pgConn.Prepare(ctx, psName, sql, nil)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
if psKey != "" {
|
|
c.preparedStatements[psKey] = sd
|
|
}
|
|
|
|
return sd, nil
|
|
}
|
|
|
|
// Deallocate releases a prepared statement. Calling Deallocate on a non-existent prepared statement will succeed.
|
|
func (c *Conn) Deallocate(ctx context.Context, name string) error {
|
|
var psName string
|
|
sd := c.preparedStatements[name]
|
|
if sd != nil {
|
|
psName = sd.Name
|
|
} else {
|
|
psName = name
|
|
}
|
|
|
|
err := c.pgConn.Deallocate(ctx, psName)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
if sd != nil {
|
|
delete(c.preparedStatements, name)
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
// DeallocateAll releases all previously prepared statements from the server and client, where it also resets the statement and description cache.
|
|
func (c *Conn) DeallocateAll(ctx context.Context) error {
|
|
c.preparedStatements = map[string]*pgconn.StatementDescription{}
|
|
if c.config.StatementCacheCapacity > 0 {
|
|
c.statementCache = stmtcache.NewLRUCache(c.config.StatementCacheCapacity)
|
|
}
|
|
if c.config.DescriptionCacheCapacity > 0 {
|
|
c.descriptionCache = stmtcache.NewLRUCache(c.config.DescriptionCacheCapacity)
|
|
}
|
|
_, err := c.pgConn.Exec(ctx, "deallocate all").ReadAll()
|
|
return err
|
|
}
|
|
|
|
func (c *Conn) bufferNotifications(_ *pgconn.PgConn, n *pgconn.Notification) {
|
|
c.notifications = append(c.notifications, n)
|
|
}
|
|
|
|
// WaitForNotification waits for a PostgreSQL notification. It wraps the underlying pgconn notification system in a
|
|
// slightly more convenient form.
|
|
func (c *Conn) WaitForNotification(ctx context.Context) (*pgconn.Notification, error) {
|
|
var n *pgconn.Notification
|
|
|
|
// Return already received notification immediately
|
|
if len(c.notifications) > 0 {
|
|
n = c.notifications[0]
|
|
c.notifications = c.notifications[1:]
|
|
return n, nil
|
|
}
|
|
|
|
err := c.pgConn.WaitForNotification(ctx)
|
|
if len(c.notifications) > 0 {
|
|
n = c.notifications[0]
|
|
c.notifications = c.notifications[1:]
|
|
}
|
|
return n, err
|
|
}
|
|
|
|
// IsClosed reports if the connection has been closed.
|
|
func (c *Conn) IsClosed() bool {
|
|
return c.pgConn.IsClosed()
|
|
}
|
|
|
|
func (c *Conn) die(err error) {
|
|
if c.IsClosed() {
|
|
return
|
|
}
|
|
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
cancel() // force immediate hard cancel
|
|
c.pgConn.Close(ctx)
|
|
}
|
|
|
|
func quoteIdentifier(s string) string {
|
|
return `"` + strings.ReplaceAll(s, `"`, `""`) + `"`
|
|
}
|
|
|
|
// Ping delegates to the underlying *pgconn.PgConn.Ping.
|
|
func (c *Conn) Ping(ctx context.Context) error {
|
|
return c.pgConn.Ping(ctx)
|
|
}
|
|
|
|
// PgConn returns the underlying *pgconn.PgConn. This is an escape hatch method that allows lower level access to the
|
|
// PostgreSQL connection than pgx exposes.
|
|
//
|
|
// It is strongly recommended that the connection be idle (no in-progress queries) before the underlying *pgconn.PgConn
|
|
// is used and the connection must be returned to the same state before any *pgx.Conn methods are again used.
|
|
func (c *Conn) PgConn() *pgconn.PgConn { return c.pgConn }
|
|
|
|
// TypeMap returns the connection info used for this connection.
|
|
func (c *Conn) TypeMap() *pgtype.Map { return c.typeMap }
|
|
|
|
// Config returns a copy of config that was used to establish this connection.
|
|
func (c *Conn) Config() *ConnConfig { return c.config.Copy() }
|
|
|
|
// Exec executes sql. sql can be either a prepared statement name or an SQL string. arguments should be referenced
|
|
// positionally from the sql string as $1, $2, etc.
|
|
func (c *Conn) Exec(ctx context.Context, sql string, arguments ...any) (pgconn.CommandTag, error) {
|
|
if c.queryTracer != nil {
|
|
ctx = c.queryTracer.TraceQueryStart(ctx, c, TraceQueryStartData{SQL: sql, Args: arguments})
|
|
}
|
|
|
|
if err := c.deallocateInvalidatedCachedStatements(ctx); err != nil {
|
|
return pgconn.CommandTag{}, err
|
|
}
|
|
|
|
commandTag, err := c.exec(ctx, sql, arguments...)
|
|
|
|
if c.queryTracer != nil {
|
|
c.queryTracer.TraceQueryEnd(ctx, c, TraceQueryEndData{CommandTag: commandTag, Err: err})
|
|
}
|
|
|
|
return commandTag, err
|
|
}
|
|
|
|
func (c *Conn) exec(ctx context.Context, sql string, arguments ...any) (commandTag pgconn.CommandTag, err error) {
|
|
mode := c.config.DefaultQueryExecMode
|
|
var queryRewriter QueryRewriter
|
|
|
|
optionLoop:
|
|
for len(arguments) > 0 {
|
|
switch arg := arguments[0].(type) {
|
|
case QueryExecMode:
|
|
mode = arg
|
|
arguments = arguments[1:]
|
|
case QueryRewriter:
|
|
queryRewriter = arg
|
|
arguments = arguments[1:]
|
|
default:
|
|
break optionLoop
|
|
}
|
|
}
|
|
|
|
if queryRewriter != nil {
|
|
sql, arguments, err = queryRewriter.RewriteQuery(ctx, c, sql, arguments)
|
|
if err != nil {
|
|
return pgconn.CommandTag{}, fmt.Errorf("rewrite query failed: %w", err)
|
|
}
|
|
}
|
|
|
|
// Always use simple protocol when there are no arguments.
|
|
if len(arguments) == 0 {
|
|
mode = QueryExecModeSimpleProtocol
|
|
}
|
|
|
|
if sd, ok := c.preparedStatements[sql]; ok {
|
|
return c.execPrepared(ctx, sd, arguments)
|
|
}
|
|
|
|
switch mode {
|
|
case QueryExecModeCacheStatement:
|
|
if c.statementCache == nil {
|
|
return pgconn.CommandTag{}, errDisabledStatementCache
|
|
}
|
|
sd := c.statementCache.Get(sql)
|
|
if sd == nil {
|
|
sd, err = c.Prepare(ctx, stmtcache.StatementName(sql), sql)
|
|
if err != nil {
|
|
return pgconn.CommandTag{}, err
|
|
}
|
|
c.statementCache.Put(sd)
|
|
}
|
|
|
|
return c.execPrepared(ctx, sd, arguments)
|
|
case QueryExecModeCacheDescribe:
|
|
if c.descriptionCache == nil {
|
|
return pgconn.CommandTag{}, errDisabledDescriptionCache
|
|
}
|
|
sd := c.descriptionCache.Get(sql)
|
|
if sd == nil {
|
|
sd, err = c.Prepare(ctx, "", sql)
|
|
if err != nil {
|
|
return pgconn.CommandTag{}, err
|
|
}
|
|
c.descriptionCache.Put(sd)
|
|
}
|
|
|
|
return c.execParams(ctx, sd, arguments)
|
|
case QueryExecModeDescribeExec:
|
|
sd, err := c.Prepare(ctx, "", sql)
|
|
if err != nil {
|
|
return pgconn.CommandTag{}, err
|
|
}
|
|
return c.execPrepared(ctx, sd, arguments)
|
|
case QueryExecModeExec:
|
|
return c.execSQLParams(ctx, sql, arguments)
|
|
case QueryExecModeSimpleProtocol:
|
|
return c.execSimpleProtocol(ctx, sql, arguments)
|
|
default:
|
|
return pgconn.CommandTag{}, fmt.Errorf("unknown QueryExecMode: %v", mode)
|
|
}
|
|
}
|
|
|
|
func (c *Conn) execSimpleProtocol(ctx context.Context, sql string, arguments []any) (commandTag pgconn.CommandTag, err error) {
|
|
if len(arguments) > 0 {
|
|
sql, err = c.sanitizeForSimpleQuery(sql, arguments...)
|
|
if err != nil {
|
|
return pgconn.CommandTag{}, err
|
|
}
|
|
}
|
|
|
|
mrr := c.pgConn.Exec(ctx, sql)
|
|
for mrr.NextResult() {
|
|
commandTag, _ = mrr.ResultReader().Close()
|
|
}
|
|
err = mrr.Close()
|
|
return commandTag, err
|
|
}
|
|
|
|
func (c *Conn) execParams(ctx context.Context, sd *pgconn.StatementDescription, arguments []any) (pgconn.CommandTag, error) {
|
|
err := c.eqb.Build(c.typeMap, sd, arguments)
|
|
if err != nil {
|
|
return pgconn.CommandTag{}, err
|
|
}
|
|
|
|
result := c.pgConn.ExecParams(ctx, sd.SQL, c.eqb.ParamValues, sd.ParamOIDs, c.eqb.ParamFormats, c.eqb.ResultFormats).Read()
|
|
c.eqb.reset() // Allow c.eqb internal memory to be GC'ed as soon as possible.
|
|
return result.CommandTag, result.Err
|
|
}
|
|
|
|
func (c *Conn) execPrepared(ctx context.Context, sd *pgconn.StatementDescription, arguments []any) (pgconn.CommandTag, error) {
|
|
err := c.eqb.Build(c.typeMap, sd, arguments)
|
|
if err != nil {
|
|
return pgconn.CommandTag{}, err
|
|
}
|
|
|
|
result := c.pgConn.ExecPrepared(ctx, sd.Name, c.eqb.ParamValues, c.eqb.ParamFormats, c.eqb.ResultFormats).Read()
|
|
c.eqb.reset() // Allow c.eqb internal memory to be GC'ed as soon as possible.
|
|
return result.CommandTag, result.Err
|
|
}
|
|
|
|
type unknownArgumentTypeQueryExecModeExecError struct {
|
|
arg any
|
|
}
|
|
|
|
func (e *unknownArgumentTypeQueryExecModeExecError) Error() string {
|
|
return fmt.Sprintf("cannot use unregistered type %T as query argument in QueryExecModeExec", e.arg)
|
|
}
|
|
|
|
func (c *Conn) execSQLParams(ctx context.Context, sql string, args []any) (pgconn.CommandTag, error) {
|
|
err := c.eqb.Build(c.typeMap, nil, args)
|
|
if err != nil {
|
|
return pgconn.CommandTag{}, err
|
|
}
|
|
|
|
result := c.pgConn.ExecParams(ctx, sql, c.eqb.ParamValues, nil, c.eqb.ParamFormats, c.eqb.ResultFormats).Read()
|
|
c.eqb.reset() // Allow c.eqb internal memory to be GC'ed as soon as possible.
|
|
return result.CommandTag, result.Err
|
|
}
|
|
|
|
func (c *Conn) getRows(ctx context.Context, sql string, args []any) *baseRows {
|
|
r := &baseRows{}
|
|
|
|
r.ctx = ctx
|
|
r.queryTracer = c.queryTracer
|
|
r.typeMap = c.typeMap
|
|
r.startTime = time.Now()
|
|
r.sql = sql
|
|
r.args = args
|
|
r.conn = c
|
|
|
|
return r
|
|
}
|
|
|
|
type QueryExecMode int32
|
|
|
|
const (
|
|
_ QueryExecMode = iota
|
|
|
|
// Automatically prepare and cache statements. This uses the extended protocol. Queries are executed in a single round
|
|
// trip after the statement is cached. This is the default. If the database schema is modified or the search_path is
|
|
// changed after a statement is cached then the first execution of a previously cached query may fail. e.g. If the
|
|
// number of columns returned by a "SELECT *" changes or the type of a column is changed.
|
|
QueryExecModeCacheStatement
|
|
|
|
// Cache statement descriptions (i.e. argument and result types) and assume they do not change. This uses the extended
|
|
// protocol. Queries are executed in a single round trip after the description is cached. If the database schema is
|
|
// modified or the search_path is changed after a statement is cached then the first execution of a previously cached
|
|
// query may fail. e.g. If the number of columns returned by a "SELECT *" changes or the type of a column is changed.
|
|
QueryExecModeCacheDescribe
|
|
|
|
// Get the statement description on every execution. This uses the extended protocol. Queries require two round trips
|
|
// to execute. It does not use named prepared statements. But it does use the unnamed prepared statement to get the
|
|
// statement description on the first round trip and then uses it to execute the query on the second round trip. This
|
|
// may cause problems with connection poolers that switch the underlying connection between round trips. It is safe
|
|
// even when the database schema is modified concurrently.
|
|
QueryExecModeDescribeExec
|
|
|
|
// Assume the PostgreSQL query parameter types based on the Go type of the arguments. This uses the extended protocol
|
|
// with text formatted parameters and results. Queries are executed in a single round trip. Type mappings can be
|
|
// registered with pgtype.Map.RegisterDefaultPgType. Queries will be rejected that have arguments that are
|
|
// unregistered or ambiguous. e.g. A map[string]string may have the PostgreSQL type json or hstore. Modes that know
|
|
// the PostgreSQL type can use a map[string]string directly as an argument. This mode cannot.
|
|
QueryExecModeExec
|
|
|
|
// Use the simple protocol. Assume the PostgreSQL query parameter types based on the Go type of the arguments.
|
|
// Queries are executed in a single round trip. Type mappings can be registered with
|
|
// pgtype.Map.RegisterDefaultPgType. Queries will be rejected that have arguments that are unregistered or ambiguous.
|
|
// e.g. A map[string]string may have the PostgreSQL type json or hstore. Modes that know the PostgreSQL type can use
|
|
// a map[string]string directly as an argument. This mode cannot.
|
|
//
|
|
// QueryExecModeSimpleProtocol should have the user application visible behavior as QueryExecModeExec with minor
|
|
// exceptions such as behavior when multiple result returning queries are erroneously sent in a single string.
|
|
//
|
|
// QueryExecModeSimpleProtocol uses client side parameter interpolation. All values are quoted and escaped. Prefer
|
|
// QueryExecModeExec over QueryExecModeSimpleProtocol whenever possible. In general QueryExecModeSimpleProtocol
|
|
// should only be used if connecting to a proxy server, connection pool server, or non-PostgreSQL server that does
|
|
// not support the extended protocol.
|
|
QueryExecModeSimpleProtocol
|
|
)
|
|
|
|
func (m QueryExecMode) String() string {
|
|
switch m {
|
|
case QueryExecModeCacheStatement:
|
|
return "cache statement"
|
|
case QueryExecModeCacheDescribe:
|
|
return "cache describe"
|
|
case QueryExecModeDescribeExec:
|
|
return "describe exec"
|
|
case QueryExecModeExec:
|
|
return "exec"
|
|
case QueryExecModeSimpleProtocol:
|
|
return "simple protocol"
|
|
default:
|
|
return "invalid"
|
|
}
|
|
}
|
|
|
|
// QueryResultFormats controls the result format (text=0, binary=1) of a query by result column position.
|
|
type QueryResultFormats []int16
|
|
|
|
// QueryResultFormatsByOID controls the result format (text=0, binary=1) of a query by the result column OID.
|
|
type QueryResultFormatsByOID map[uint32]int16
|
|
|
|
// QueryRewriter rewrites a query when used as the first arguments to a query method.
|
|
type QueryRewriter interface {
|
|
RewriteQuery(ctx context.Context, conn *Conn, sql string, args []any) (newSQL string, newArgs []any, err error)
|
|
}
|
|
|
|
// Query sends a query to the server and returns a Rows to read the results. Only errors encountered sending the query
|
|
// and initializing Rows will be returned. Err() on the returned Rows must be checked after the Rows is closed to
|
|
// determine if the query executed successfully.
|
|
//
|
|
// The returned Rows must be closed before the connection can be used again. It is safe to attempt to read from the
|
|
// returned Rows even if an error is returned. The error will be the available in rows.Err() after rows are closed. It
|
|
// is allowed to ignore the error returned from Query and handle it in Rows.
|
|
//
|
|
// It is possible for a call of FieldDescriptions on the returned Rows to return nil even if the Query call did not
|
|
// return an error.
|
|
//
|
|
// It is possible for a query to return one or more rows before encountering an error. In most cases the rows should be
|
|
// collected before processing rather than processed while receiving each row. This avoids the possibility of the
|
|
// application processing rows from a query that the server rejected. The CollectRows function is useful here.
|
|
//
|
|
// An implementor of QueryRewriter may be passed as the first element of args. It can rewrite the sql and change or
|
|
// replace args. For example, NamedArgs is QueryRewriter that implements named arguments.
|
|
//
|
|
// For extra control over how the query is executed, the types QueryExecMode, QueryResultFormats, and
|
|
// QueryResultFormatsByOID may be used as the first args to control exactly how the query is executed. This is rarely
|
|
// needed. See the documentation for those types for details.
|
|
func (c *Conn) Query(ctx context.Context, sql string, args ...any) (Rows, error) {
|
|
if c.queryTracer != nil {
|
|
ctx = c.queryTracer.TraceQueryStart(ctx, c, TraceQueryStartData{SQL: sql, Args: args})
|
|
}
|
|
|
|
if err := c.deallocateInvalidatedCachedStatements(ctx); err != nil {
|
|
if c.queryTracer != nil {
|
|
c.queryTracer.TraceQueryEnd(ctx, c, TraceQueryEndData{Err: err})
|
|
}
|
|
return &baseRows{err: err, closed: true}, err
|
|
}
|
|
|
|
var resultFormats QueryResultFormats
|
|
var resultFormatsByOID QueryResultFormatsByOID
|
|
mode := c.config.DefaultQueryExecMode
|
|
var queryRewriter QueryRewriter
|
|
|
|
optionLoop:
|
|
for len(args) > 0 {
|
|
switch arg := args[0].(type) {
|
|
case QueryResultFormats:
|
|
resultFormats = arg
|
|
args = args[1:]
|
|
case QueryResultFormatsByOID:
|
|
resultFormatsByOID = arg
|
|
args = args[1:]
|
|
case QueryExecMode:
|
|
mode = arg
|
|
args = args[1:]
|
|
case QueryRewriter:
|
|
queryRewriter = arg
|
|
args = args[1:]
|
|
default:
|
|
break optionLoop
|
|
}
|
|
}
|
|
|
|
if queryRewriter != nil {
|
|
var err error
|
|
originalSQL := sql
|
|
originalArgs := args
|
|
sql, args, err = queryRewriter.RewriteQuery(ctx, c, sql, args)
|
|
if err != nil {
|
|
rows := c.getRows(ctx, originalSQL, originalArgs)
|
|
err = fmt.Errorf("rewrite query failed: %w", err)
|
|
rows.fatal(err)
|
|
return rows, err
|
|
}
|
|
}
|
|
|
|
// Bypass any statement caching.
|
|
if sql == "" {
|
|
mode = QueryExecModeSimpleProtocol
|
|
}
|
|
|
|
c.eqb.reset()
|
|
rows := c.getRows(ctx, sql, args)
|
|
|
|
var err error
|
|
sd, explicitPreparedStatement := c.preparedStatements[sql]
|
|
if sd != nil || mode == QueryExecModeCacheStatement || mode == QueryExecModeCacheDescribe || mode == QueryExecModeDescribeExec {
|
|
if sd == nil {
|
|
sd, err = c.getStatementDescription(ctx, mode, sql)
|
|
if err != nil {
|
|
rows.fatal(err)
|
|
return rows, err
|
|
}
|
|
}
|
|
|
|
if len(sd.ParamOIDs) != len(args) {
|
|
rows.fatal(fmt.Errorf("expected %d arguments, got %d", len(sd.ParamOIDs), len(args)))
|
|
return rows, rows.err
|
|
}
|
|
|
|
rows.sql = sd.SQL
|
|
|
|
err = c.eqb.Build(c.typeMap, sd, args)
|
|
if err != nil {
|
|
rows.fatal(err)
|
|
return rows, rows.err
|
|
}
|
|
|
|
if resultFormatsByOID != nil {
|
|
resultFormats = make([]int16, len(sd.Fields))
|
|
for i := range resultFormats {
|
|
resultFormats[i] = resultFormatsByOID[uint32(sd.Fields[i].DataTypeOID)]
|
|
}
|
|
}
|
|
|
|
if resultFormats == nil {
|
|
resultFormats = c.eqb.ResultFormats
|
|
}
|
|
|
|
if !explicitPreparedStatement && mode == QueryExecModeCacheDescribe {
|
|
rows.resultReader = c.pgConn.ExecParams(ctx, sql, c.eqb.ParamValues, sd.ParamOIDs, c.eqb.ParamFormats, resultFormats)
|
|
} else {
|
|
rows.resultReader = c.pgConn.ExecPrepared(ctx, sd.Name, c.eqb.ParamValues, c.eqb.ParamFormats, resultFormats)
|
|
}
|
|
} else if mode == QueryExecModeExec {
|
|
err := c.eqb.Build(c.typeMap, nil, args)
|
|
if err != nil {
|
|
rows.fatal(err)
|
|
return rows, rows.err
|
|
}
|
|
|
|
rows.resultReader = c.pgConn.ExecParams(ctx, sql, c.eqb.ParamValues, nil, c.eqb.ParamFormats, c.eqb.ResultFormats)
|
|
} else if mode == QueryExecModeSimpleProtocol {
|
|
sql, err = c.sanitizeForSimpleQuery(sql, args...)
|
|
if err != nil {
|
|
rows.fatal(err)
|
|
return rows, err
|
|
}
|
|
|
|
mrr := c.pgConn.Exec(ctx, sql)
|
|
if mrr.NextResult() {
|
|
rows.resultReader = mrr.ResultReader()
|
|
rows.multiResultReader = mrr
|
|
} else {
|
|
err = mrr.Close()
|
|
rows.fatal(err)
|
|
return rows, err
|
|
}
|
|
|
|
return rows, nil
|
|
} else {
|
|
err = fmt.Errorf("unknown QueryExecMode: %v", mode)
|
|
rows.fatal(err)
|
|
return rows, rows.err
|
|
}
|
|
|
|
c.eqb.reset() // Allow c.eqb internal memory to be GC'ed as soon as possible.
|
|
|
|
return rows, rows.err
|
|
}
|
|
|
|
// getStatementDescription returns the statement description of the sql query
|
|
// according to the given mode.
|
|
//
|
|
// If the mode is one that doesn't require to know the param and result OIDs
|
|
// then nil is returned without error.
|
|
func (c *Conn) getStatementDescription(
|
|
ctx context.Context,
|
|
mode QueryExecMode,
|
|
sql string,
|
|
) (sd *pgconn.StatementDescription, err error) {
|
|
switch mode {
|
|
case QueryExecModeCacheStatement:
|
|
if c.statementCache == nil {
|
|
return nil, errDisabledStatementCache
|
|
}
|
|
sd = c.statementCache.Get(sql)
|
|
if sd == nil {
|
|
sd, err = c.Prepare(ctx, stmtcache.StatementName(sql), sql)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
c.statementCache.Put(sd)
|
|
}
|
|
case QueryExecModeCacheDescribe:
|
|
if c.descriptionCache == nil {
|
|
return nil, errDisabledDescriptionCache
|
|
}
|
|
sd = c.descriptionCache.Get(sql)
|
|
if sd == nil {
|
|
sd, err = c.Prepare(ctx, "", sql)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
c.descriptionCache.Put(sd)
|
|
}
|
|
case QueryExecModeDescribeExec:
|
|
return c.Prepare(ctx, "", sql)
|
|
}
|
|
return sd, err
|
|
}
|
|
|
|
// QueryRow is a convenience wrapper over Query. Any error that occurs while
|
|
// querying is deferred until calling Scan on the returned Row. That Row will
|
|
// error with ErrNoRows if no rows are returned.
|
|
func (c *Conn) QueryRow(ctx context.Context, sql string, args ...any) Row {
|
|
rows, _ := c.Query(ctx, sql, args...)
|
|
return (*connRow)(rows.(*baseRows))
|
|
}
|
|
|
|
// SendBatch sends all queued queries to the server at once. All queries are run in an implicit transaction unless
|
|
// explicit transaction control statements are executed. The returned BatchResults must be closed before the connection
|
|
// is used again.
|
|
func (c *Conn) SendBatch(ctx context.Context, b *Batch) (br BatchResults) {
|
|
if c.batchTracer != nil {
|
|
ctx = c.batchTracer.TraceBatchStart(ctx, c, TraceBatchStartData{Batch: b})
|
|
defer func() {
|
|
err := br.(interface{ earlyError() error }).earlyError()
|
|
if err != nil {
|
|
c.batchTracer.TraceBatchEnd(ctx, c, TraceBatchEndData{Err: err})
|
|
}
|
|
}()
|
|
}
|
|
|
|
if err := c.deallocateInvalidatedCachedStatements(ctx); err != nil {
|
|
return &batchResults{ctx: ctx, conn: c, err: err}
|
|
}
|
|
|
|
for _, bi := range b.QueuedQueries {
|
|
var queryRewriter QueryRewriter
|
|
sql := bi.SQL
|
|
arguments := bi.Arguments
|
|
|
|
optionLoop:
|
|
for len(arguments) > 0 {
|
|
// Update Batch.Queue function comment when additional options are implemented
|
|
switch arg := arguments[0].(type) {
|
|
case QueryRewriter:
|
|
queryRewriter = arg
|
|
arguments = arguments[1:]
|
|
default:
|
|
break optionLoop
|
|
}
|
|
}
|
|
|
|
if queryRewriter != nil {
|
|
var err error
|
|
sql, arguments, err = queryRewriter.RewriteQuery(ctx, c, sql, arguments)
|
|
if err != nil {
|
|
return &batchResults{ctx: ctx, conn: c, err: fmt.Errorf("rewrite query failed: %w", err)}
|
|
}
|
|
}
|
|
|
|
bi.SQL = sql
|
|
bi.Arguments = arguments
|
|
}
|
|
|
|
// TODO: changing mode per batch? Update Batch.Queue function comment when implemented
|
|
mode := c.config.DefaultQueryExecMode
|
|
if mode == QueryExecModeSimpleProtocol {
|
|
return c.sendBatchQueryExecModeSimpleProtocol(ctx, b)
|
|
}
|
|
|
|
// All other modes use extended protocol and thus can use prepared statements.
|
|
for _, bi := range b.QueuedQueries {
|
|
if sd, ok := c.preparedStatements[bi.SQL]; ok {
|
|
bi.sd = sd
|
|
}
|
|
}
|
|
|
|
switch mode {
|
|
case QueryExecModeExec:
|
|
return c.sendBatchQueryExecModeExec(ctx, b)
|
|
case QueryExecModeCacheStatement:
|
|
return c.sendBatchQueryExecModeCacheStatement(ctx, b)
|
|
case QueryExecModeCacheDescribe:
|
|
return c.sendBatchQueryExecModeCacheDescribe(ctx, b)
|
|
case QueryExecModeDescribeExec:
|
|
return c.sendBatchQueryExecModeDescribeExec(ctx, b)
|
|
default:
|
|
panic("unknown QueryExecMode")
|
|
}
|
|
}
|
|
|
|
func (c *Conn) sendBatchQueryExecModeSimpleProtocol(ctx context.Context, b *Batch) *batchResults {
|
|
var sb strings.Builder
|
|
for i, bi := range b.QueuedQueries {
|
|
if i > 0 {
|
|
sb.WriteByte(';')
|
|
}
|
|
sql, err := c.sanitizeForSimpleQuery(bi.SQL, bi.Arguments...)
|
|
if err != nil {
|
|
return &batchResults{ctx: ctx, conn: c, err: err}
|
|
}
|
|
sb.WriteString(sql)
|
|
}
|
|
mrr := c.pgConn.Exec(ctx, sb.String())
|
|
return &batchResults{
|
|
ctx: ctx,
|
|
conn: c,
|
|
mrr: mrr,
|
|
b: b,
|
|
qqIdx: 0,
|
|
}
|
|
}
|
|
|
|
func (c *Conn) sendBatchQueryExecModeExec(ctx context.Context, b *Batch) *batchResults {
|
|
batch := &pgconn.Batch{}
|
|
|
|
for _, bi := range b.QueuedQueries {
|
|
sd := bi.sd
|
|
if sd != nil {
|
|
err := c.eqb.Build(c.typeMap, sd, bi.Arguments)
|
|
if err != nil {
|
|
return &batchResults{ctx: ctx, conn: c, err: err}
|
|
}
|
|
|
|
batch.ExecPrepared(sd.Name, c.eqb.ParamValues, c.eqb.ParamFormats, c.eqb.ResultFormats)
|
|
} else {
|
|
err := c.eqb.Build(c.typeMap, nil, bi.Arguments)
|
|
if err != nil {
|
|
return &batchResults{ctx: ctx, conn: c, err: err}
|
|
}
|
|
batch.ExecParams(bi.SQL, c.eqb.ParamValues, nil, c.eqb.ParamFormats, c.eqb.ResultFormats)
|
|
}
|
|
}
|
|
|
|
c.eqb.reset() // Allow c.eqb internal memory to be GC'ed as soon as possible.
|
|
|
|
mrr := c.pgConn.ExecBatch(ctx, batch)
|
|
|
|
return &batchResults{
|
|
ctx: ctx,
|
|
conn: c,
|
|
mrr: mrr,
|
|
b: b,
|
|
qqIdx: 0,
|
|
}
|
|
}
|
|
|
|
func (c *Conn) sendBatchQueryExecModeCacheStatement(ctx context.Context, b *Batch) (pbr *pipelineBatchResults) {
|
|
if c.statementCache == nil {
|
|
return &pipelineBatchResults{ctx: ctx, conn: c, err: errDisabledStatementCache, closed: true}
|
|
}
|
|
|
|
distinctNewQueries := []*pgconn.StatementDescription{}
|
|
distinctNewQueriesIdxMap := make(map[string]int)
|
|
|
|
for _, bi := range b.QueuedQueries {
|
|
if bi.sd == nil {
|
|
sd := c.statementCache.Get(bi.SQL)
|
|
if sd != nil {
|
|
bi.sd = sd
|
|
} else {
|
|
if idx, present := distinctNewQueriesIdxMap[bi.SQL]; present {
|
|
bi.sd = distinctNewQueries[idx]
|
|
} else {
|
|
sd = &pgconn.StatementDescription{
|
|
Name: stmtcache.StatementName(bi.SQL),
|
|
SQL: bi.SQL,
|
|
}
|
|
distinctNewQueriesIdxMap[sd.SQL] = len(distinctNewQueries)
|
|
distinctNewQueries = append(distinctNewQueries, sd)
|
|
bi.sd = sd
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
return c.sendBatchExtendedWithDescription(ctx, b, distinctNewQueries, c.statementCache)
|
|
}
|
|
|
|
func (c *Conn) sendBatchQueryExecModeCacheDescribe(ctx context.Context, b *Batch) (pbr *pipelineBatchResults) {
|
|
if c.descriptionCache == nil {
|
|
return &pipelineBatchResults{ctx: ctx, conn: c, err: errDisabledDescriptionCache, closed: true}
|
|
}
|
|
|
|
distinctNewQueries := []*pgconn.StatementDescription{}
|
|
distinctNewQueriesIdxMap := make(map[string]int)
|
|
|
|
for _, bi := range b.QueuedQueries {
|
|
if bi.sd == nil {
|
|
sd := c.descriptionCache.Get(bi.SQL)
|
|
if sd != nil {
|
|
bi.sd = sd
|
|
} else {
|
|
if idx, present := distinctNewQueriesIdxMap[bi.SQL]; present {
|
|
bi.sd = distinctNewQueries[idx]
|
|
} else {
|
|
sd = &pgconn.StatementDescription{
|
|
SQL: bi.SQL,
|
|
}
|
|
distinctNewQueriesIdxMap[sd.SQL] = len(distinctNewQueries)
|
|
distinctNewQueries = append(distinctNewQueries, sd)
|
|
bi.sd = sd
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
return c.sendBatchExtendedWithDescription(ctx, b, distinctNewQueries, c.descriptionCache)
|
|
}
|
|
|
|
func (c *Conn) sendBatchQueryExecModeDescribeExec(ctx context.Context, b *Batch) (pbr *pipelineBatchResults) {
|
|
distinctNewQueries := []*pgconn.StatementDescription{}
|
|
distinctNewQueriesIdxMap := make(map[string]int)
|
|
|
|
for _, bi := range b.QueuedQueries {
|
|
if bi.sd == nil {
|
|
if idx, present := distinctNewQueriesIdxMap[bi.SQL]; present {
|
|
bi.sd = distinctNewQueries[idx]
|
|
} else {
|
|
sd := &pgconn.StatementDescription{
|
|
SQL: bi.SQL,
|
|
}
|
|
distinctNewQueriesIdxMap[sd.SQL] = len(distinctNewQueries)
|
|
distinctNewQueries = append(distinctNewQueries, sd)
|
|
bi.sd = sd
|
|
}
|
|
}
|
|
}
|
|
|
|
return c.sendBatchExtendedWithDescription(ctx, b, distinctNewQueries, nil)
|
|
}
|
|
|
|
func (c *Conn) sendBatchExtendedWithDescription(ctx context.Context, b *Batch, distinctNewQueries []*pgconn.StatementDescription, sdCache stmtcache.Cache) (pbr *pipelineBatchResults) {
|
|
pipeline := c.pgConn.StartPipeline(ctx)
|
|
defer func() {
|
|
if pbr != nil && pbr.err != nil {
|
|
pipeline.Close()
|
|
}
|
|
}()
|
|
|
|
// Prepare any needed queries
|
|
if len(distinctNewQueries) > 0 {
|
|
for _, sd := range distinctNewQueries {
|
|
pipeline.SendPrepare(sd.Name, sd.SQL, nil)
|
|
}
|
|
|
|
err := pipeline.Sync()
|
|
if err != nil {
|
|
return &pipelineBatchResults{ctx: ctx, conn: c, err: err, closed: true}
|
|
}
|
|
|
|
for _, sd := range distinctNewQueries {
|
|
results, err := pipeline.GetResults()
|
|
if err != nil {
|
|
return &pipelineBatchResults{ctx: ctx, conn: c, err: err, closed: true}
|
|
}
|
|
|
|
resultSD, ok := results.(*pgconn.StatementDescription)
|
|
if !ok {
|
|
return &pipelineBatchResults{ctx: ctx, conn: c, err: fmt.Errorf("expected statement description, got %T", results), closed: true}
|
|
}
|
|
|
|
// Fill in the previously empty / pending statement descriptions.
|
|
sd.ParamOIDs = resultSD.ParamOIDs
|
|
sd.Fields = resultSD.Fields
|
|
}
|
|
|
|
results, err := pipeline.GetResults()
|
|
if err != nil {
|
|
return &pipelineBatchResults{ctx: ctx, conn: c, err: err, closed: true}
|
|
}
|
|
|
|
_, ok := results.(*pgconn.PipelineSync)
|
|
if !ok {
|
|
return &pipelineBatchResults{ctx: ctx, conn: c, err: fmt.Errorf("expected sync, got %T", results), closed: true}
|
|
}
|
|
}
|
|
|
|
// Put all statements into the cache. It's fine if it overflows because HandleInvalidated will clean them up later.
|
|
if sdCache != nil {
|
|
for _, sd := range distinctNewQueries {
|
|
sdCache.Put(sd)
|
|
}
|
|
}
|
|
|
|
// Queue the queries.
|
|
for _, bi := range b.QueuedQueries {
|
|
err := c.eqb.Build(c.typeMap, bi.sd, bi.Arguments)
|
|
if err != nil {
|
|
// we wrap the error so we the user can understand which query failed inside the batch
|
|
err = fmt.Errorf("error building query %s: %w", bi.SQL, err)
|
|
return &pipelineBatchResults{ctx: ctx, conn: c, err: err, closed: true}
|
|
}
|
|
|
|
if bi.sd.Name == "" {
|
|
pipeline.SendQueryParams(bi.sd.SQL, c.eqb.ParamValues, bi.sd.ParamOIDs, c.eqb.ParamFormats, c.eqb.ResultFormats)
|
|
} else {
|
|
pipeline.SendQueryPrepared(bi.sd.Name, c.eqb.ParamValues, c.eqb.ParamFormats, c.eqb.ResultFormats)
|
|
}
|
|
}
|
|
|
|
err := pipeline.Sync()
|
|
if err != nil {
|
|
return &pipelineBatchResults{ctx: ctx, conn: c, err: err, closed: true}
|
|
}
|
|
|
|
return &pipelineBatchResults{
|
|
ctx: ctx,
|
|
conn: c,
|
|
pipeline: pipeline,
|
|
b: b,
|
|
}
|
|
}
|
|
|
|
func (c *Conn) sanitizeForSimpleQuery(sql string, args ...any) (string, error) {
|
|
if c.pgConn.ParameterStatus("standard_conforming_strings") != "on" {
|
|
return "", errors.New("simple protocol queries must be run with standard_conforming_strings=on")
|
|
}
|
|
|
|
if c.pgConn.ParameterStatus("client_encoding") != "UTF8" {
|
|
return "", errors.New("simple protocol queries must be run with client_encoding=UTF8")
|
|
}
|
|
|
|
var err error
|
|
valueArgs := make([]any, len(args))
|
|
for i, a := range args {
|
|
valueArgs[i], err = convertSimpleArgument(c.typeMap, a)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
}
|
|
|
|
return sanitize.SanitizeSQL(sql, valueArgs...)
|
|
}
|
|
|
|
// LoadType inspects the database for typeName and produces a pgtype.Type suitable for registration. typeName must be
|
|
// the name of a type where the underlying type(s) is already understood by pgx. It is for derived types. In particular,
|
|
// typeName must be one of the following:
|
|
// - An array type name of a type that is already registered. e.g. "_foo" when "foo" is registered.
|
|
// - A composite type name where all field types are already registered.
|
|
// - A domain type name where the base type is already registered.
|
|
// - An enum type name.
|
|
// - A range type name where the element type is already registered.
|
|
// - A multirange type name where the element type is already registered.
|
|
func (c *Conn) LoadType(ctx context.Context, typeName string) (*pgtype.Type, error) {
|
|
var oid uint32
|
|
|
|
err := c.QueryRow(ctx, "select $1::text::regtype::oid;", typeName).Scan(&oid)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
var typtype string
|
|
var typbasetype uint32
|
|
|
|
err = c.QueryRow(ctx, "select typtype::text, typbasetype from pg_type where oid=$1", oid).Scan(&typtype, &typbasetype)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
switch typtype {
|
|
case "b": // array
|
|
elementOID, err := c.getArrayElementOID(ctx, oid)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
dt, ok := c.TypeMap().TypeForOID(elementOID)
|
|
if !ok {
|
|
return nil, errors.New("array element OID not registered")
|
|
}
|
|
|
|
return &pgtype.Type{Name: typeName, OID: oid, Codec: &pgtype.ArrayCodec{ElementType: dt}}, nil
|
|
case "c": // composite
|
|
fields, err := c.getCompositeFields(ctx, oid)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return &pgtype.Type{Name: typeName, OID: oid, Codec: &pgtype.CompositeCodec{Fields: fields}}, nil
|
|
case "d": // domain
|
|
dt, ok := c.TypeMap().TypeForOID(typbasetype)
|
|
if !ok {
|
|
return nil, errors.New("domain base type OID not registered")
|
|
}
|
|
|
|
return &pgtype.Type{Name: typeName, OID: oid, Codec: dt.Codec}, nil
|
|
case "e": // enum
|
|
return &pgtype.Type{Name: typeName, OID: oid, Codec: &pgtype.EnumCodec{}}, nil
|
|
case "r": // range
|
|
elementOID, err := c.getRangeElementOID(ctx, oid)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
dt, ok := c.TypeMap().TypeForOID(elementOID)
|
|
if !ok {
|
|
return nil, errors.New("range element OID not registered")
|
|
}
|
|
|
|
return &pgtype.Type{Name: typeName, OID: oid, Codec: &pgtype.RangeCodec{ElementType: dt}}, nil
|
|
case "m": // multirange
|
|
elementOID, err := c.getMultiRangeElementOID(ctx, oid)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
dt, ok := c.TypeMap().TypeForOID(elementOID)
|
|
if !ok {
|
|
return nil, errors.New("multirange element OID not registered")
|
|
}
|
|
|
|
return &pgtype.Type{Name: typeName, OID: oid, Codec: &pgtype.MultirangeCodec{ElementType: dt}}, nil
|
|
default:
|
|
return &pgtype.Type{}, errors.New("unknown typtype")
|
|
}
|
|
}
|
|
|
|
func (c *Conn) getArrayElementOID(ctx context.Context, oid uint32) (uint32, error) {
|
|
var typelem uint32
|
|
|
|
err := c.QueryRow(ctx, "select typelem from pg_type where oid=$1", oid).Scan(&typelem)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
|
|
return typelem, nil
|
|
}
|
|
|
|
func (c *Conn) getRangeElementOID(ctx context.Context, oid uint32) (uint32, error) {
|
|
var typelem uint32
|
|
|
|
err := c.QueryRow(ctx, "select rngsubtype from pg_range where rngtypid=$1", oid).Scan(&typelem)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
|
|
return typelem, nil
|
|
}
|
|
|
|
func (c *Conn) getMultiRangeElementOID(ctx context.Context, oid uint32) (uint32, error) {
|
|
var typelem uint32
|
|
|
|
err := c.QueryRow(ctx, "select rngtypid from pg_range where rngmultitypid=$1", oid).Scan(&typelem)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
|
|
return typelem, nil
|
|
}
|
|
|
|
func (c *Conn) getCompositeFields(ctx context.Context, oid uint32) ([]pgtype.CompositeCodecField, error) {
|
|
var typrelid uint32
|
|
|
|
err := c.QueryRow(ctx, "select typrelid from pg_type where oid=$1", oid).Scan(&typrelid)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
var fields []pgtype.CompositeCodecField
|
|
var fieldName string
|
|
var fieldOID uint32
|
|
rows, _ := c.Query(ctx, `select attname, atttypid
|
|
from pg_attribute
|
|
where attrelid=$1
|
|
and not attisdropped
|
|
and attnum > 0
|
|
order by attnum`,
|
|
typrelid,
|
|
)
|
|
_, err = ForEachRow(rows, []any{&fieldName, &fieldOID}, func() error {
|
|
dt, ok := c.TypeMap().TypeForOID(fieldOID)
|
|
if !ok {
|
|
return fmt.Errorf("unknown composite type field OID: %v", fieldOID)
|
|
}
|
|
fields = append(fields, pgtype.CompositeCodecField{Name: fieldName, Type: dt})
|
|
return nil
|
|
})
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return fields, nil
|
|
}
|
|
|
|
func (c *Conn) deallocateInvalidatedCachedStatements(ctx context.Context) error {
|
|
if txStatus := c.pgConn.TxStatus(); txStatus != 'I' && txStatus != 'T' {
|
|
return nil
|
|
}
|
|
|
|
if c.descriptionCache != nil {
|
|
c.descriptionCache.RemoveInvalidated()
|
|
}
|
|
|
|
var invalidatedStatements []*pgconn.StatementDescription
|
|
if c.statementCache != nil {
|
|
invalidatedStatements = c.statementCache.GetInvalidated()
|
|
}
|
|
|
|
if len(invalidatedStatements) == 0 {
|
|
return nil
|
|
}
|
|
|
|
pipeline := c.pgConn.StartPipeline(ctx)
|
|
defer pipeline.Close()
|
|
|
|
for _, sd := range invalidatedStatements {
|
|
pipeline.SendDeallocate(sd.Name)
|
|
}
|
|
|
|
err := pipeline.Sync()
|
|
if err != nil {
|
|
return fmt.Errorf("failed to deallocate cached statement(s): %w", err)
|
|
}
|
|
|
|
err = pipeline.Close()
|
|
if err != nil {
|
|
return fmt.Errorf("failed to deallocate cached statement(s): %w", err)
|
|
}
|
|
|
|
c.statementCache.RemoveInvalidated()
|
|
for _, sd := range invalidatedStatements {
|
|
delete(c.preparedStatements, sd.Name)
|
|
}
|
|
|
|
return nil
|
|
}
|