mirror of https://github.com/jackc/pgx.git
268 lines
7.2 KiB
Go
268 lines
7.2 KiB
Go
package pgx_test
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"os"
|
|
"testing"
|
|
|
|
"github.com/jackc/pgx/v5"
|
|
"github.com/jackc/pgx/v5/pgconn"
|
|
"github.com/jackc/pgx/v5/pgxtest"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
type testRowScanner struct {
|
|
name string
|
|
age int32
|
|
}
|
|
|
|
func (rs *testRowScanner) ScanRow(rows pgx.Rows) error {
|
|
return rows.Scan(&rs.name, &rs.age)
|
|
}
|
|
|
|
func TestRowScanner(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
defaultConnTestRunner.RunTest(context.Background(), t, func(ctx context.Context, t testing.TB, conn *pgx.Conn) {
|
|
var s testRowScanner
|
|
err := conn.QueryRow(ctx, "select 'Adam' as name, 72 as height").Scan(&s)
|
|
require.NoError(t, err)
|
|
require.Equal(t, "Adam", s.name)
|
|
require.Equal(t, int32(72), s.age)
|
|
})
|
|
}
|
|
|
|
func TestForEachRow(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
pgxtest.RunWithQueryExecModes(context.Background(), t, defaultConnTestRunner, nil, func(ctx context.Context, t testing.TB, conn *pgx.Conn) {
|
|
var actualResults []any
|
|
|
|
rows, _ := conn.Query(
|
|
context.Background(),
|
|
"select n, n * 2 from generate_series(1, $1) n",
|
|
3,
|
|
)
|
|
var a, b int
|
|
ct, err := pgx.ForEachRow(rows, []any{&a, &b}, func() error {
|
|
actualResults = append(actualResults, []any{a, b})
|
|
return nil
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
expectedResults := []any{
|
|
[]any{1, 2},
|
|
[]any{2, 4},
|
|
[]any{3, 6},
|
|
}
|
|
require.Equal(t, expectedResults, actualResults)
|
|
require.EqualValues(t, 3, ct.RowsAffected())
|
|
})
|
|
}
|
|
|
|
func TestForEachRowScanError(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
pgxtest.RunWithQueryExecModes(context.Background(), t, defaultConnTestRunner, nil, func(ctx context.Context, t testing.TB, conn *pgx.Conn) {
|
|
var actualResults []any
|
|
|
|
rows, _ := conn.Query(
|
|
context.Background(),
|
|
"select 'foo', 'bar' from generate_series(1, $1) n",
|
|
3,
|
|
)
|
|
var a, b int
|
|
ct, err := pgx.ForEachRow(rows, []any{&a, &b}, func() error {
|
|
actualResults = append(actualResults, []any{a, b})
|
|
return nil
|
|
})
|
|
require.EqualError(t, err, "can't scan into dest[0]: cannot scan OID 25 in text format into *int")
|
|
require.Equal(t, pgconn.CommandTag{}, ct)
|
|
})
|
|
}
|
|
|
|
func TestForEachRowAbort(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
pgxtest.RunWithQueryExecModes(context.Background(), t, defaultConnTestRunner, nil, func(ctx context.Context, t testing.TB, conn *pgx.Conn) {
|
|
rows, _ := conn.Query(
|
|
context.Background(),
|
|
"select n, n * 2 from generate_series(1, $1) n",
|
|
3,
|
|
)
|
|
var a, b int
|
|
ct, err := pgx.ForEachRow(rows, []any{&a, &b}, func() error {
|
|
return errors.New("abort")
|
|
})
|
|
require.EqualError(t, err, "abort")
|
|
require.Equal(t, pgconn.CommandTag{}, ct)
|
|
})
|
|
}
|
|
|
|
func ExampleForEachRow() {
|
|
conn, err := pgx.Connect(context.Background(), os.Getenv("PGX_TEST_DATABASE"))
|
|
if err != nil {
|
|
fmt.Printf("Unable to establish connection: %v", err)
|
|
return
|
|
}
|
|
|
|
rows, _ := conn.Query(
|
|
context.Background(),
|
|
"select n, n * 2 from generate_series(1, $1) n",
|
|
3,
|
|
)
|
|
var a, b int
|
|
_, err = pgx.ForEachRow(rows, []any{&a, &b}, func() error {
|
|
fmt.Printf("%v, %v\n", a, b)
|
|
return nil
|
|
})
|
|
if err != nil {
|
|
fmt.Printf("ForEachRow error: %v", err)
|
|
return
|
|
}
|
|
|
|
// Output:
|
|
// 1, 2
|
|
// 2, 4
|
|
// 3, 6
|
|
}
|
|
|
|
func TestCollectRows(t *testing.T) {
|
|
defaultConnTestRunner.RunTest(context.Background(), t, func(ctx context.Context, t testing.TB, conn *pgx.Conn) {
|
|
rows, _ := conn.Query(ctx, `select n from generate_series(0, 99) n`)
|
|
numbers, err := pgx.CollectRows(rows, func(row pgx.CollectableRow) (int32, error) {
|
|
var n int32
|
|
err := row.Scan(&n)
|
|
return n, err
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
assert.Len(t, numbers, 100)
|
|
for i := range numbers {
|
|
assert.Equal(t, int32(i), numbers[i])
|
|
}
|
|
})
|
|
}
|
|
|
|
func TestCollectOneRow(t *testing.T) {
|
|
defaultConnTestRunner.RunTest(context.Background(), t, func(ctx context.Context, t testing.TB, conn *pgx.Conn) {
|
|
rows, _ := conn.Query(ctx, `select 42`)
|
|
n, err := pgx.CollectOneRow(rows, func(row pgx.CollectableRow) (int32, error) {
|
|
var n int32
|
|
err := row.Scan(&n)
|
|
return n, err
|
|
})
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, int32(42), n)
|
|
})
|
|
}
|
|
|
|
func TestCollectOneRowNotFound(t *testing.T) {
|
|
defaultConnTestRunner.RunTest(context.Background(), t, func(ctx context.Context, t testing.TB, conn *pgx.Conn) {
|
|
rows, _ := conn.Query(ctx, `select 42 where false`)
|
|
n, err := pgx.CollectOneRow(rows, func(row pgx.CollectableRow) (int32, error) {
|
|
var n int32
|
|
err := row.Scan(&n)
|
|
return n, err
|
|
})
|
|
assert.ErrorIs(t, err, pgx.ErrNoRows)
|
|
assert.Equal(t, int32(0), n)
|
|
})
|
|
}
|
|
|
|
func TestCollectOneRowIgnoresExtraRows(t *testing.T) {
|
|
defaultConnTestRunner.RunTest(context.Background(), t, func(ctx context.Context, t testing.TB, conn *pgx.Conn) {
|
|
rows, _ := conn.Query(ctx, `select n from generate_series(42, 99) n`)
|
|
n, err := pgx.CollectOneRow(rows, func(row pgx.CollectableRow) (int32, error) {
|
|
var n int32
|
|
err := row.Scan(&n)
|
|
return n, err
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, int32(42), n)
|
|
})
|
|
}
|
|
|
|
func TestRowTo(t *testing.T) {
|
|
defaultConnTestRunner.RunTest(context.Background(), t, func(ctx context.Context, t testing.TB, conn *pgx.Conn) {
|
|
rows, _ := conn.Query(ctx, `select n from generate_series(0, 99) n`)
|
|
numbers, err := pgx.CollectRows(rows, pgx.RowTo[int32])
|
|
require.NoError(t, err)
|
|
|
|
assert.Len(t, numbers, 100)
|
|
for i := range numbers {
|
|
assert.Equal(t, int32(i), numbers[i])
|
|
}
|
|
})
|
|
}
|
|
|
|
func TestRowToAddrOf(t *testing.T) {
|
|
defaultConnTestRunner.RunTest(context.Background(), t, func(ctx context.Context, t testing.TB, conn *pgx.Conn) {
|
|
rows, _ := conn.Query(ctx, `select n from generate_series(0, 99) n`)
|
|
numbers, err := pgx.CollectRows(rows, pgx.RowToAddrOf[int32])
|
|
require.NoError(t, err)
|
|
|
|
assert.Len(t, numbers, 100)
|
|
for i := range numbers {
|
|
assert.Equal(t, int32(i), *numbers[i])
|
|
}
|
|
})
|
|
}
|
|
|
|
func TestRowToMap(t *testing.T) {
|
|
defaultConnTestRunner.RunTest(context.Background(), t, func(ctx context.Context, t testing.TB, conn *pgx.Conn) {
|
|
rows, _ := conn.Query(ctx, `select 'Joe' as name, n as age from generate_series(0, 9) n`)
|
|
slice, err := pgx.CollectRows(rows, pgx.RowToMap)
|
|
require.NoError(t, err)
|
|
|
|
assert.Len(t, slice, 10)
|
|
for i := range slice {
|
|
assert.Equal(t, "Joe", slice[i]["name"])
|
|
assert.EqualValues(t, i, slice[i]["age"])
|
|
}
|
|
})
|
|
}
|
|
|
|
func TestRowToStructPos(t *testing.T) {
|
|
type person struct {
|
|
Name string
|
|
Age int32
|
|
}
|
|
|
|
defaultConnTestRunner.RunTest(context.Background(), t, func(ctx context.Context, t testing.TB, conn *pgx.Conn) {
|
|
rows, _ := conn.Query(ctx, `select 'Joe' as name, n as age from generate_series(0, 9) n`)
|
|
slice, err := pgx.CollectRows(rows, pgx.RowToStructByPos[person])
|
|
require.NoError(t, err)
|
|
|
|
assert.Len(t, slice, 10)
|
|
for i := range slice {
|
|
assert.Equal(t, "Joe", slice[i].Name)
|
|
assert.EqualValues(t, i, slice[i].Age)
|
|
}
|
|
})
|
|
}
|
|
|
|
func TestRowToAddrOfStructPos(t *testing.T) {
|
|
type person struct {
|
|
Name string
|
|
Age int32
|
|
}
|
|
|
|
defaultConnTestRunner.RunTest(context.Background(), t, func(ctx context.Context, t testing.TB, conn *pgx.Conn) {
|
|
rows, _ := conn.Query(ctx, `select 'Joe' as name, n as age from generate_series(0, 9) n`)
|
|
slice, err := pgx.CollectRows(rows, pgx.RowToAddrOfStructByPos[person])
|
|
require.NoError(t, err)
|
|
|
|
assert.Len(t, slice, 10)
|
|
for i := range slice {
|
|
assert.Equal(t, "Joe", slice[i].Name)
|
|
assert.EqualValues(t, i, slice[i].Age)
|
|
}
|
|
})
|
|
}
|