feat: shared pgx-pool

go.mod

@@ -180,6 +180,7 @@ require (
 	github.com/inconshreveable/mousetrap v1.1.0 // indirect
 	github.com/jackc/pgpassfile v1.0.0 // indirect
 	github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect
+	github.com/jackc/puddle/v2 v2.2.2 // indirect
 	github.com/jaegertracing/jaeger-idl v0.5.0 // indirect
 	github.com/jcmturner/aescts/v2 v2.0.0 // indirect
 	github.com/jcmturner/dnsutils/v2 v2.0.0 // indirect

go.sum

@@ -660,7 +660,6 @@ github.com/jackc/pgx/v5 v5.7.6/go.mod h1:aruU7o91Tc2q2cFp5h4uP3f6ztExVpyVv88Xl/8
 github.com/jackc/puddle v0.0.0-20190413234325-e4ced69a3a2b/go.mod h1:m4B5Dj62Y0fbyuIc15OsIqK0+JU8nkqQjsgx7dvjSWk=
 github.com/jackc/puddle v0.0.0-20190608224051-11cab39313c9/go.mod h1:m4B5Dj62Y0fbyuIc15OsIqK0+JU8nkqQjsgx7dvjSWk=
 github.com/jackc/puddle v1.1.3/go.mod h1:m4B5Dj62Y0fbyuIc15OsIqK0+JU8nkqQjsgx7dvjSWk=
-github.com/jackc/puddle v1.2.1 h1:gI8os0wpRXFd4FiAY2dWiqRK037tjj3t7rKFeO4X5iw=
 github.com/jackc/puddle v1.2.1/go.mod h1:m4B5Dj62Y0fbyuIc15OsIqK0+JU8nkqQjsgx7dvjSWk=
 github.com/jackc/puddle/v2 v2.2.2 h1:PR8nw+E/1w0GLuRFSmiioY6UooMp6KJv0/61nB7icHo=
 github.com/jackc/puddle/v2 v2.2.2/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4=

retriever/postgresqlretriever/database.go

@@ -0,0 +1,49 @@
+package postgresqlretriever
+
+import (
+	"context"
+	"sync"
+
+	"github.com/jackc/pgx/v5/pgxpool"
+)
+
+var (
+	pool     *pgxpool.Pool
+	mu       sync.Mutex
+	refCount int
+)
+
+func GetPool(ctx context.Context, uri string) (*pgxpool.Pool, error) {
+	mu.Lock()
+	defer mu.Unlock()
+
+	if pool == nil {
+		p, err := pgxpool.New(ctx, uri)
+		if err != nil {
+			return nil, err
+		}
+		if err := p.Ping(ctx); err != nil {
+			p.Close()
+			return nil, err
+		}
+
+		pool = p
+	}
+
+	refCount++
+	return pool, nil
+}
+
+func ReleasePool() {
+	mu.Lock()
+	defer mu.Unlock()
+
+	refCount--
+	if refCount <= 0 {
+		if pool != nil {
+			pool.Close()
+			pool = nil
+		}
+		refCount = 0
+	}
+}

retriever/postgresqlretriever/retriever.go

@@ -7,6 +7,7 @@
 	"log/slog"
 
 	"github.com/jackc/pgx/v5"
+	"github.com/jackc/pgx/v5/pgxpool"
 	"github.com/thomaspoignant/go-feature-flag/retriever"
 	"github.com/thomaspoignant/go-feature-flag/utils"
 	"github.com/thomaspoignant/go-feature-flag/utils/fflog"
@@ -26,7 +27,7 @@
 	logger  *fflog.FFLogger
 	status  retriever.Status
 	columns map[string]string
-	conn    *pgx.Conn
+	pool    *pgxpool.Pool
 	flagset *string
 }
 
@@ -40,23 +41,19 @@
 	r.columns = r.getColumnNames()
 	r.flagset = flagset
 
-	if r.conn == nil {
+	if r.pool == nil {
 		r.logger.Info("Initializing PostgreSQL retriever")
 		r.logger.Debug("Using columns", "columns", r.columns)
 
-		conn, err := pgx.Connect(ctx, r.URI)
+		pool, err := GetPool(ctx, r.URI)
 		if err != nil {
 			r.status = retriever.RetrieverError
 			return err
 		}
 
-		if err := conn.Ping(ctx); err != nil {
-			r.status = retriever.RetrieverError
-			return err
-		}
-
-		r.conn = conn
+		r.pool = pool
 	}
+
 	r.status = retriever.RetrieverReady
 	return nil
 }
@@ -70,38 +67,37 @@
 }
 
 // Shutdown closes the database connection.
 func (r *Retriever) Shutdown(ctx context.Context) error {
-	if r.conn == nil {
-		return nil
-	}
-	return r.conn.Close(ctx)
+	ReleasePool()
+	return nil
 }
 
 // Retrieve fetches flag configuration from PostgreSQL.
 func (r *Retriever) Retrieve(ctx context.Context) ([]byte, error) {
-	if r.conn == nil {
-		return nil, fmt.Errorf("database connection is not initialized")
+	if r.pool == nil {
+		return nil, fmt.Errorf("database connection pool is not initialized")
 	}
 
 	// Build the query using the configured table and column names
 	query := r.buildQuery()
 
 	// Build the arguments for the query
-	args := []interface{}{}
+	args := []any{}
 	if r.getFlagset() != "" {
-		args = []interface{}{r.getFlagset()}
+		// If a flagset is defined, it becomes the first ($1) argument.
+		args = []any{r.getFlagset()}
 	}
 
 	r.logger.Debug("Executing PostgreSQL query", slog.String("query", query), slog.Any("args", args))
 
-	rows, err := r.conn.Query(ctx, query, args...)
+	rows, err := r.pool.Query(ctx, query, args...)
 	if err != nil {
 		return nil, fmt.Errorf("failed to execute query: %w", err)
 	}
 	defer rows.Close()
 
 	// Map to store flag configurations with flag_name as key
-	flagConfigs := make(map[string]interface{})
+	flagConfigs := make(map[string]any)
 
 	for rows.Next() {
 		var flagName string
@@ -121,6 +117,7 @@
 		flagConfigs[flagName] = config
 	}
 
+	// Check for any errors that occurred during row iteration
 	if err := rows.Err(); err != nil {
 		return nil, fmt.Errorf("rows iteration error: %w", err)
 	}

retriever/postgresqlretriever/retriever_test.go

@@ -378,7 +378,7 @@ func TestRetrieverDataHandlingErrors(t *testing.T) {
 			assert.Contains(t, err.Error(), "failed to execute query")
 		}
 	})
-	}
+}
 
 // TestRetrieverEdgeCases tests additional edge cases and boundary conditions
 func TestRetrieverEdgeCases(t *testing.T) {