mirror of
https://github.com/superseriousbusiness/gotosocial.git
synced 2024-12-05 18:12:46 +00:00
3ab6214449
Bumps [github.com/jackc/pgx/v5](https://github.com/jackc/pgx) from 5.5.0 to 5.5.1. - [Changelog](https://github.com/jackc/pgx/blob/master/CHANGELOG.md) - [Commits](https://github.com/jackc/pgx/compare/v5.5.0...v5.5.1) --- updated-dependencies: - dependency-name: github.com/jackc/pgx/v5 dependency-type: direct:production update-type: version-update:semver-patch ... Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
277 lines
7.1 KiB
Go
277 lines
7.1 KiB
Go
package pgx
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"fmt"
|
|
"io"
|
|
|
|
"github.com/jackc/pgx/v5/internal/pgio"
|
|
"github.com/jackc/pgx/v5/pgconn"
|
|
)
|
|
|
|
// CopyFromRows returns a CopyFromSource interface over the provided rows slice
|
|
// making it usable by *Conn.CopyFrom.
|
|
func CopyFromRows(rows [][]any) CopyFromSource {
|
|
return ©FromRows{rows: rows, idx: -1}
|
|
}
|
|
|
|
type copyFromRows struct {
|
|
rows [][]any
|
|
idx int
|
|
}
|
|
|
|
func (ctr *copyFromRows) Next() bool {
|
|
ctr.idx++
|
|
return ctr.idx < len(ctr.rows)
|
|
}
|
|
|
|
func (ctr *copyFromRows) Values() ([]any, error) {
|
|
return ctr.rows[ctr.idx], nil
|
|
}
|
|
|
|
func (ctr *copyFromRows) Err() error {
|
|
return nil
|
|
}
|
|
|
|
// CopyFromSlice returns a CopyFromSource interface over a dynamic func
|
|
// making it usable by *Conn.CopyFrom.
|
|
func CopyFromSlice(length int, next func(int) ([]any, error)) CopyFromSource {
|
|
return ©FromSlice{next: next, idx: -1, len: length}
|
|
}
|
|
|
|
type copyFromSlice struct {
|
|
next func(int) ([]any, error)
|
|
idx int
|
|
len int
|
|
err error
|
|
}
|
|
|
|
func (cts *copyFromSlice) Next() bool {
|
|
cts.idx++
|
|
return cts.idx < cts.len
|
|
}
|
|
|
|
func (cts *copyFromSlice) Values() ([]any, error) {
|
|
values, err := cts.next(cts.idx)
|
|
if err != nil {
|
|
cts.err = err
|
|
}
|
|
return values, err
|
|
}
|
|
|
|
func (cts *copyFromSlice) Err() error {
|
|
return cts.err
|
|
}
|
|
|
|
// CopyFromFunc returns a CopyFromSource interface that relies on nxtf for values.
|
|
// nxtf returns rows until it either signals an 'end of data' by returning row=nil and err=nil,
|
|
// or it returns an error. If nxtf returns an error, the copy is aborted.
|
|
func CopyFromFunc(nxtf func() (row []any, err error)) CopyFromSource {
|
|
return ©FromFunc{next: nxtf}
|
|
}
|
|
|
|
type copyFromFunc struct {
|
|
next func() ([]any, error)
|
|
valueRow []any
|
|
err error
|
|
}
|
|
|
|
func (g *copyFromFunc) Next() bool {
|
|
g.valueRow, g.err = g.next()
|
|
// only return true if valueRow exists and no error
|
|
return g.valueRow != nil && g.err == nil
|
|
}
|
|
|
|
func (g *copyFromFunc) Values() ([]any, error) {
|
|
return g.valueRow, g.err
|
|
}
|
|
|
|
func (g *copyFromFunc) Err() error {
|
|
return g.err
|
|
}
|
|
|
|
// CopyFromSource is the interface used by *Conn.CopyFrom as the source for copy data.
|
|
type CopyFromSource interface {
|
|
// Next returns true if there is another row and makes the next row data
|
|
// available to Values(). When there are no more rows available or an error
|
|
// has occurred it returns false.
|
|
Next() bool
|
|
|
|
// Values returns the values for the current row.
|
|
Values() ([]any, error)
|
|
|
|
// Err returns any error that has been encountered by the CopyFromSource. If
|
|
// this is not nil *Conn.CopyFrom will abort the copy.
|
|
Err() error
|
|
}
|
|
|
|
type copyFrom struct {
|
|
conn *Conn
|
|
tableName Identifier
|
|
columnNames []string
|
|
rowSrc CopyFromSource
|
|
readerErrChan chan error
|
|
mode QueryExecMode
|
|
}
|
|
|
|
func (ct *copyFrom) run(ctx context.Context) (int64, error) {
|
|
if ct.conn.copyFromTracer != nil {
|
|
ctx = ct.conn.copyFromTracer.TraceCopyFromStart(ctx, ct.conn, TraceCopyFromStartData{
|
|
TableName: ct.tableName,
|
|
ColumnNames: ct.columnNames,
|
|
})
|
|
}
|
|
|
|
quotedTableName := ct.tableName.Sanitize()
|
|
cbuf := &bytes.Buffer{}
|
|
for i, cn := range ct.columnNames {
|
|
if i != 0 {
|
|
cbuf.WriteString(", ")
|
|
}
|
|
cbuf.WriteString(quoteIdentifier(cn))
|
|
}
|
|
quotedColumnNames := cbuf.String()
|
|
|
|
var sd *pgconn.StatementDescription
|
|
switch ct.mode {
|
|
case QueryExecModeExec, QueryExecModeSimpleProtocol:
|
|
// These modes don't support the binary format. Before the inclusion of the
|
|
// QueryExecModes, Conn.Prepare was called on every COPY operation to get
|
|
// the OIDs. These prepared statements were not cached.
|
|
//
|
|
// Since that's the same behavior provided by QueryExecModeDescribeExec,
|
|
// we'll default to that mode.
|
|
ct.mode = QueryExecModeDescribeExec
|
|
fallthrough
|
|
case QueryExecModeCacheStatement, QueryExecModeCacheDescribe, QueryExecModeDescribeExec:
|
|
var err error
|
|
sd, err = ct.conn.getStatementDescription(
|
|
ctx,
|
|
ct.mode,
|
|
fmt.Sprintf("select %s from %s", quotedColumnNames, quotedTableName),
|
|
)
|
|
if err != nil {
|
|
return 0, fmt.Errorf("statement description failed: %w", err)
|
|
}
|
|
default:
|
|
return 0, fmt.Errorf("unknown QueryExecMode: %v", ct.mode)
|
|
}
|
|
|
|
r, w := io.Pipe()
|
|
doneChan := make(chan struct{})
|
|
|
|
go func() {
|
|
defer close(doneChan)
|
|
|
|
// Purposely NOT using defer w.Close(). See https://github.com/golang/go/issues/24283.
|
|
buf := ct.conn.wbuf
|
|
|
|
buf = append(buf, "PGCOPY\n\377\r\n\000"...)
|
|
buf = pgio.AppendInt32(buf, 0)
|
|
buf = pgio.AppendInt32(buf, 0)
|
|
|
|
moreRows := true
|
|
for moreRows {
|
|
var err error
|
|
moreRows, buf, err = ct.buildCopyBuf(buf, sd)
|
|
if err != nil {
|
|
w.CloseWithError(err)
|
|
return
|
|
}
|
|
|
|
if ct.rowSrc.Err() != nil {
|
|
w.CloseWithError(ct.rowSrc.Err())
|
|
return
|
|
}
|
|
|
|
if len(buf) > 0 {
|
|
_, err = w.Write(buf)
|
|
if err != nil {
|
|
w.Close()
|
|
return
|
|
}
|
|
}
|
|
|
|
buf = buf[:0]
|
|
}
|
|
|
|
w.Close()
|
|
}()
|
|
|
|
commandTag, err := ct.conn.pgConn.CopyFrom(ctx, r, fmt.Sprintf("copy %s ( %s ) from stdin binary;", quotedTableName, quotedColumnNames))
|
|
|
|
r.Close()
|
|
<-doneChan
|
|
|
|
if ct.conn.copyFromTracer != nil {
|
|
ct.conn.copyFromTracer.TraceCopyFromEnd(ctx, ct.conn, TraceCopyFromEndData{
|
|
CommandTag: commandTag,
|
|
Err: err,
|
|
})
|
|
}
|
|
|
|
return commandTag.RowsAffected(), err
|
|
}
|
|
|
|
func (ct *copyFrom) buildCopyBuf(buf []byte, sd *pgconn.StatementDescription) (bool, []byte, error) {
|
|
const sendBufSize = 65536 - 5 // The packet has a 5-byte header
|
|
lastBufLen := 0
|
|
largestRowLen := 0
|
|
|
|
for ct.rowSrc.Next() {
|
|
lastBufLen = len(buf)
|
|
|
|
values, err := ct.rowSrc.Values()
|
|
if err != nil {
|
|
return false, nil, err
|
|
}
|
|
if len(values) != len(ct.columnNames) {
|
|
return false, nil, fmt.Errorf("expected %d values, got %d values", len(ct.columnNames), len(values))
|
|
}
|
|
|
|
buf = pgio.AppendInt16(buf, int16(len(ct.columnNames)))
|
|
for i, val := range values {
|
|
buf, err = encodeCopyValue(ct.conn.typeMap, buf, sd.Fields[i].DataTypeOID, val)
|
|
if err != nil {
|
|
return false, nil, err
|
|
}
|
|
}
|
|
|
|
rowLen := len(buf) - lastBufLen
|
|
if rowLen > largestRowLen {
|
|
largestRowLen = rowLen
|
|
}
|
|
|
|
// Try not to overflow size of the buffer PgConn.CopyFrom will be reading into. If that happens then the nature of
|
|
// io.Pipe means that the next Read will be short. This can lead to pathological send sizes such as 65531, 13, 65531
|
|
// 13, 65531, 13, 65531, 13.
|
|
if len(buf) > sendBufSize-largestRowLen {
|
|
return true, buf, nil
|
|
}
|
|
}
|
|
|
|
return false, buf, nil
|
|
}
|
|
|
|
// CopyFrom uses the PostgreSQL copy protocol to perform bulk data insertion. It returns the number of rows copied and
|
|
// an error.
|
|
//
|
|
// CopyFrom requires all values use the binary format. A pgtype.Type that supports the binary format must be registered
|
|
// for the type of each column. Almost all types implemented by pgx support the binary format.
|
|
//
|
|
// Even though enum types appear to be strings they still must be registered to use with CopyFrom. This can be done with
|
|
// Conn.LoadType and pgtype.Map.RegisterType.
|
|
func (c *Conn) CopyFrom(ctx context.Context, tableName Identifier, columnNames []string, rowSrc CopyFromSource) (int64, error) {
|
|
ct := ©From{
|
|
conn: c,
|
|
tableName: tableName,
|
|
columnNames: columnNames,
|
|
rowSrc: rowSrc,
|
|
readerErrChan: make(chan error),
|
|
mode: c.config.DefaultQueryExecMode,
|
|
}
|
|
|
|
return ct.run(ctx)
|
|
}
|