Repository: jackc/pgx
Pull request #1776 (2023-10-20 09:07:08)

Proposal: CopyFromCh
Summary
This pull request proposes a new helper function to extend the functionality of CopyFromSlice() and CopyFromRow(). The new function aims to facilitate streaming data to the database via channels, which is particularly useful in concurrent programming scenarios.


jackc: I don't this is quite the right abstraction. Usually, its best for libraries to expose non-channel interfaces and let the application add channels on top if desired.

In this case there are a couple concerning points.

1. The `CopyFromSource` interface isn't expected to block. But the channel could block indefinitely. The context passed to `CopyFrom` couldn't cancel this and then `CopyFrom` could block.
2. Someone may want to pass `[]any` instead of `T` (similar to `CopyFromRows`).

If anything should be added I think a better interface might be something like:

```go
func CopyFromFunc(next func() ([]any, error)) CopyFromSource
```

The `next` function could read from a channel. `next` could use a closure to gain access to a cancellation channel (like `ctx.Done()`) to avoid blocking. `next` could also handle `[]any`, a struct with reflect, or a struct manually converted to `[]any`.

`CopyFromFunc` could also be used to read from any other source.

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index fb2304a2..63ee30d5 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,3 +1,17 @@
+# 5.5.0 (November 4, 2023)
+
+* Add CollectExactlyOneRow. (Julien GOTTELAND)
+* Add OpenDBFromPool to create *database/sql.DB from *pgxpool.Pool. (Lev Zakharov)
+* Prepare can automatically choose statement name based on sql. This makes it easier to explicitly manage prepared statements.
+* Statement cache now uses deterministic, stable statement names.
+* database/sql prepared statement names are deterministically generated.
+* Fix: SendBatch wasn't respecting context cancellation.
+* Fix: Timeout error from pipeline is now normalized.
+* Fix: database/sql encoding json.RawMessage to []byte.
+* CancelRequest: Wait for the cancel request to be acknowledged by the server. This should improve PgBouncer compatibility. (Anton Levakin)
+* stdlib: Use Ping instead of CheckConn in ResetSession
+* Add json.Marshaler and json.Unmarshaler for Float4, Float8 (Kirill Mironov)
+
 # 5.4.3 (August 5, 2023)
 
 * Fix: QCharArrayOID was defined with the wrong OID (Christoph Engelbert)
diff --git a/README.md b/README.md
index 621830a3..2a9efc23 100644
--- a/README.md
+++ b/README.md
@@ -86,6 +86,10 @@ It is also possible to use the `database/sql` interface and convert a connection
 
 See CONTRIBUTING.md for setup instructions.
 
+## Architecture
+
+See the presentation at Golang Estonia, [PGX Top to Bottom](https://www.youtube.com/watch?v=sXMSWhcHCf8) for a description of pgx architecture.
+
 ## Supported Go and PostgreSQL Versions
 
 pgx supports the same versions of Go and PostgreSQL that are supported by their respective teams. For [Go](https://golang.org/doc/devel/release.html#policy) that is the two most recent major releases and for [PostgreSQL](https://www.postgresql.org/support/versioning/) the major releases in the last 5 years. This means pgx supports Go 1.20 and higher and PostgreSQL 11 and higher. pgx also is tested against the latest version of [CockroachDB](https://www.cockroachlabs.com/product/).
diff --git a/copy_from.go b/copy_from.go
index a2c227fd..abcd2239 100644
--- a/copy_from.go
+++ b/copy_from.go
@@ -64,6 +64,33 @@ 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 &copyFromFunc{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
diff --git a/pgtype/float4.go b/pgtype/float4.go
index 2540f9e5..91ca0147 100644
--- a/pgtype/float4.go
+++ b/pgtype/float4.go
@@ -3,6 +3,7 @@ package pgtype
 import (
 	"database/sql/driver"
 	"encoding/binary"
+	"encoding/json"
 	"fmt"
 	"math"
 	"strconv"
@@ -65,6 +66,29 @@ func (f Float4) Value() (driver.Value, error) {
 	return float64(f.Float32), nil
 }
 
+func (f Float4) MarshalJSON() ([]byte, error) {
+	if !f.Valid {
+		return []byte("null"), nil
+	}
+	return json.Marshal(f.Float32)
+}
+
+func (f *Float4) UnmarshalJSON(b []byte) error {
+	var n *float32
+	err := json.Unmarshal(b, &n)
+	if err != nil {
+		return err
+	}
+
+	if n == nil {
+		*f = Float4{}
+	} else {
+		*f = Float4{Float32: *n, Valid: true}
+	}
+
+	return nil
+}
+
 type Float4Codec struct{}
 
 func (Float4Codec) FormatSupported(format int16) bool {
diff --git a/pgtype/float8.go b/pgtype/float8.go
index 65e5d8b3..9c923c9a 100644
--- a/pgtype/float8.go
+++ b/pgtype/float8.go
@@ -74,6 +74,29 @@ func (f Float8) Value() (driver.Value, error) {
 	return f.Float64, nil
 }
 
+func (f Float8) MarshalJSON() ([]byte, error) {
+	if !f.Valid {
+		return []byte("null"), nil
+	}
+	return json.Marshal(f.Float64)
+}
+
+func (f *Float8) UnmarshalJSON(b []byte) error {
+	var n *float64
+	err := json.Unmarshal(b, &n)
+	if err != nil {
+		return err
+	}
+
+	if n == nil {
+		*f = Float8{}
+	} else {
+		*f = Float8{Float64: *n, Valid: true}
+	}
+
+	return nil
+}
+
 type Float8Codec struct{}
 
 func (Float8Codec) FormatSupported(format int16) bool {
@@ -109,13 +132,6 @@ func (Float8Codec) PlanEncode(m *Map, oid uint32, format int16, value any) Encod
 	return nil
 }
 
-func (f *Float8) MarshalJSON() ([]byte, error) {
-	if !f.Valid {
-		return []byte("null"), nil
-	}
-	return json.Marshal(f.Float64)
-}
-
 type encodePlanFloat8CodecBinaryFloat64 struct{}
 
 func (encodePlanFloat8CodecBinaryFloat64) Encode(value any, buf []byte) (newBuf []byte, err error) {
diff --git a/stdlib/sql.go b/stdlib/sql.go
index f688f70c..c5be1a3f 100644
--- a/stdlib/sql.go
+++ b/stdlib/sql.go
@@ -548,7 +548,7 @@ func (c *Conn) ResetSession(ctx context.Context) error {
 
 	now := time.Now()
 	if now.Sub(c.lastResetSessionTime) > time.Second {
-		if err := c.conn.PgConn().CheckConn(); err != nil {
+		if err := c.conn.PgConn().Ping(ctx); err != nil {
 			return driver.ErrBadConn
 		}
 	}


Test patch:
diff --git a/copy_from_test.go b/copy_from_test.go
index ac2ccaab..9da23c04 100644
--- a/copy_from_test.go
+++ b/copy_from_test.go
@@ -802,3 +802,59 @@ func TestConnCopyFromAutomaticStringConversion(t *testing.T) {
 
 	ensureConnValid(t, conn)
 }
+
+func TestCopyFromFunc(t *testing.T) {
+	t.Parallel()
+
+	conn := mustConnectString(t, os.Getenv("PGX_TEST_DATABASE"))
+	defer closeConn(t, conn)
+
+	mustExec(t, conn, `create temporary table foo(
+		a int
+	)`)
+
+	dataCh := make(chan int, 1)
+
+	const channelItems = 10
+	go func() {
+		for i := 0; i < channelItems; i++ {
+			dataCh <- i
+		}
+		close(dataCh)
+	}()
+
+	copyCount, err := conn.CopyFrom(context.Background(), pgx.Identifier{"foo"}, []string{"a"},
+		pgx.CopyFromFunc(func() ([]any, error) {
+			v, ok := <-dataCh
+			if !ok {
+				return nil, nil
+			}
+			return []any{v}, nil
+		}))
+
+	require.ErrorIs(t, err, nil)
+	require.EqualValues(t, channelItems, copyCount)
+
+	rows, err := conn.Query(context.Background(), "select * from foo order by a")
+	require.NoError(t, err)
+	nums, err := pgx.CollectRows(rows, pgx.RowTo[int64])
+	require.NoError(t, err)
+	require.Equal(t, []int64{0, 1, 2, 3, 4, 5, 6, 7, 8, 9}, nums)
+
+	// simulate a failure
+	copyCount, err = conn.CopyFrom(context.Background(), pgx.Identifier{"foo"}, []string{"a"},
+		pgx.CopyFromFunc(func() func() ([]any, error) {
+			x := 9
+			return func() ([]any, error) {
+				x++
+				if x > 100 {
+					return nil, fmt.Errorf("simulated error")
+				}
+				return []any{x}, nil
+			}
+		}()))
+	require.NotErrorIs(t, err, nil)
+	require.EqualValues(t, 0, copyCount) // no change, due to error
+
+	ensureConnValid(t, conn)
+}
diff --git a/pgtype/float4_test.go b/pgtype/float4_test.go
index f155ed97..bc74921c 100644
--- a/pgtype/float4_test.go
+++ b/pgtype/float4_test.go
@@ -21,3 +21,44 @@ func TestFloat4Codec(t *testing.T) {
 		{nil, new(*float32), isExpectedEq((*float32)(nil))},
 	})
 }
+
+func TestFloat4MarshalJSON(t *testing.T) {
+	successfulTests := []struct {
+		source pgtype.Float4
+		result string
+	}{
+		{source: pgtype.Float4{Float32: 0}, result: "null"},
+		{source: pgtype.Float4{Float32: 1.23, Valid: true}, result: "1.23"},
+	}
+	for i, tt := range successfulTests {
+		r, err := tt.source.MarshalJSON()
+		if err != nil {
+			t.Errorf("%d: %v", i, err)
+		}
+
+		if string(r) != tt.result {
+			t.Errorf("%d: expected %v to convert to %v, but it was %v", i, tt.source, tt.result, string(r))
+		}
+	}
+}
+
+func TestFloat4UnmarshalJSON(t *testing.T) {
+	successfulTests := []struct {
+		source string
+		result pgtype.Float4
+	}{
+		{source: "null", result: pgtype.Float4{Float32: 0}},
+		{source: "1.23", result: pgtype.Float4{Float32: 1.23, Valid: true}},
+	}
+	for i, tt := range successfulTests {
+		var r pgtype.Float4
+		err := r.UnmarshalJSON([]byte(tt.source))
+		if err != nil {
+			t.Errorf("%d: %v", i, err)
+		}
+
+		if r != tt.result {
+			t.Errorf("%d: expected %v to convert to %v, but it was %v", i, tt.source, tt.result, r)
+		}
+	}
+}
diff --git a/pgtype/float8_test.go b/pgtype/float8_test.go
index 496b718b..64593d97 100644
--- a/pgtype/float8_test.go
+++ b/pgtype/float8_test.go
@@ -21,3 +21,44 @@ func TestFloat8Codec(t *testing.T) {
 		{nil, new(*float64), isExpectedEq((*float64)(nil))},
 	})
 }
+
+func TestFloat8MarshalJSON(t *testing.T) {
+	successfulTests := []struct {
+		source pgtype.Float8
+		result string
+	}{
+		{source: pgtype.Float8{Float64: 0}, result: "null"},
+		{source: pgtype.Float8{Float64: 1.23, Valid: true}, result: "1.23"},
+	}
+	for i, tt := range successfulTests {
+		r, err := tt.source.MarshalJSON()
+		if err != nil {
+			t.Errorf("%d: %v", i, err)
+		}
+
+		if string(r) != tt.result {
+			t.Errorf("%d: expected %v to convert to %v, but it was %v", i, tt.source, tt.result, string(r))
+		}
+	}
+}
+
+func TestFloat8UnmarshalJSON(t *testing.T) {
+	successfulTests := []struct {
+		source string
+		result pgtype.Float8
+	}{
+		{source: "null", result: pgtype.Float8{Float64: 0}},
+		{source: "1.23", result: pgtype.Float8{Float64: 1.23, Valid: true}},
+	}
+	for i, tt := range successfulTests {
+		var r pgtype.Float8
+		err := r.UnmarshalJSON([]byte(tt.source))
+		if err != nil {
+			t.Errorf("%d: %v", i, err)
+		}
+
+		if r != tt.result {
+			t.Errorf("%d: expected %v to convert to %v, but it was %v", i, tt.source, tt.result, r)
+		}
+	}
+}