From ff8c3ee47c58dffaa44ca24dce1e668e5a72d8a0 Mon Sep 17 00:00:00 2001
From: Andrew Kane <andrew@ankane.org>
Date: Wed, 23 Oct 2024 02:58:05 -0700
Subject: [PATCH] Added example and test for pgx pool - #18

---
 README.md   | 10 +++++++++-
 pgx_test.go | 23 +++++++++++++++++++++++
 2 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 949e5d3..25930fd 100644
--- a/README.md
+++ b/README.md
@@ -48,12 +48,20 @@ Enable the extension
 _, err := conn.Exec(ctx, "CREATE EXTENSION IF NOT EXISTS vector")
 ```
 
-Register the types
+Register the types with the connection
 
 ```go
 err := pgxvector.RegisterTypes(ctx, conn)
 ```
 
+or the pool
+
+```go
+config.AfterConnect = func(ctx context.Context, conn *pgx.Conn) error {
+    return pgxvector.RegisterTypes(ctx, conn)
+}
+```
+
 Create a table
 
 ```go
diff --git a/pgx_test.go b/pgx_test.go
index 815aa62..6b89bb5 100644
--- a/pgx_test.go
+++ b/pgx_test.go
@@ -8,6 +8,7 @@ import (
 
 	"github.com/jackc/pgx/v5"
 	"github.com/jackc/pgx/v5/pgtype"
+	"github.com/jackc/pgx/v5/pgxpool"
 	"github.com/pgvector/pgvector-go"
 	pgxvector "github.com/pgvector/pgvector-go/pgx"
 )
@@ -146,4 +147,26 @@ func TestPgx(t *testing.T) {
 	if err != nil {
 		panic(err)
 	}
+
+	config, err := pgxpool.ParseConfig("postgres://localhost/pgvector_go_test")
+	if err != nil {
+		panic(err)
+	}
+	config.AfterConnect = func(ctx context.Context, conn *pgx.Conn) error {
+		return pgxvector.RegisterTypes(ctx, conn)
+	}
+	pool, err := pgxpool.NewWithConfig(ctx, config)
+	defer pool.Close()
+
+	_, err = pool.CopyFrom(
+		ctx,
+		pgx.Identifier{"pgx_items"},
+		[]string{"embedding", "binary_embedding", "sparse_embedding"},
+		pgx.CopyFromSlice(1, func(i int) ([]any, error) {
+			return []interface{}{"[1,2,3]", "101", "{1:1,2:2,3:3}/3"}, nil
+		}),
+	)
+	if err != nil {
+		panic(err)
+	}
 }