pgvector restore WithConnectionURL option

vectorstores/pgvector/options.go

@@ -55,6 +55,14 @@ func WithCollectionTableName(name string) Option {
 	}
 }
 
+// WithConnectionURL is an option for specifying the Postgres connection URL. Either this
+// or WithConn must be used.
+func WithConnectionURL(connectionURL string) Option {
+	return func(p *Store) {
+		p.connURL = connectionURL
+	}
+}
+
 // WithConn is an option for specifying the Postgres connection.
 // From pgx doc: it is not safe for concurrent usage.Use a connection pool to manage access
 // to multiple database connections from multiple goroutines.
@@ -106,7 +114,7 @@ func applyClientOptions(opts ...Option) (Store, error) {
 		opt(o)
 	}
 
-	if o.conn == nil {
+	if o.conn == nil && o.connURL == "" {
 		return Store{}, fmt.Errorf("%w: missing postgres connection", ErrInvalidOptions)
 	}
 

vectorstores/pgvector/pgvector.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"io"
 	"strings"
 
 	"github.com/google/uuid"
@@ -45,9 +46,14 @@ type PGXConn interface {
 	SendBatch(ctx context.Context, batch *pgx.Batch) pgx.BatchResults
 }
 
+type CloseNoErr interface {
+	Close()
+}
+
 // Store is a wrapper around the pgvector client.
 type Store struct {
 	embedder            embeddings.Embedder
+	connURL             string
 	conn                PGXConn
 	embeddingTableName  string
 	collectionTableName string
@@ -73,6 +79,12 @@ func New(ctx context.Context, opts ...Option) (Store, error) {
 	if err != nil {
 		return Store{}, err
 	}
+	if store.conn == nil {
+		store.conn, err = pgx.Connect(ctx, store.connURL)
+		if err != nil {
+			return Store{}, err
+		}
+	}
 	if err = store.conn.Ping(ctx); err != nil {
 		return Store{}, err
 	}
@@ -82,6 +94,17 @@ func New(ctx context.Context, opts ...Option) (Store, error) {
 	return store, nil
 }
 
+// Close closes the connection.
+func (s Store) Close() error {
+	if closer, ok := s.conn.(io.Closer); ok {
+		return closer.Close()
+	}
+	if closer, ok := s.conn.(CloseNoErr); ok {
+		closer.Close()
+	}
+	return nil
+}
+
 func (s *Store) init(ctx context.Context) error {
 	tx, err := s.conn.Begin(ctx)
 	if err != nil {