From 6823c967ce6b8ff1d387ce8f50a5d6c61c5c1f18 Mon Sep 17 00:00:00 2001 From: Georgi Dimitrov Date: Tue, 2 Dec 2025 18:31:53 +0200 Subject: [PATCH 01/12] use pgxpool and singelton to initialize it just once --- go.mod | 1 + go.sum | 1 - retriever/postgresqlretriever/database.go | 64 ++++++++++++++++++++++ retriever/postgresqlretriever/retriever.go | 36 ++++++------ 4 files changed, 81 insertions(+), 21 deletions(-) create mode 100644 retriever/postgresqlretriever/database.go diff --git a/go.mod b/go.mod index b44ceec0397..cdce964d207 100644 --- a/go.mod +++ b/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 diff --git a/go.sum b/go.sum index 60940a0f3c0..a12019d9cbf 100644 --- a/go.sum +++ b/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= diff --git a/retriever/postgresqlretriever/database.go b/retriever/postgresqlretriever/database.go new file mode 100644 index 00000000000..1d275c7f0b6 --- /dev/null +++ b/retriever/postgresqlretriever/database.go @@ -0,0 +1,64 @@ +package postgresqlretriever + +import ( + "context" + "fmt" + "sync" + + "github.com/jackc/pgx/v5/pgxpool" +) + +var poolInstance *pgxpool.Pool +var once sync.Once +var poolErr error + +// refCount tracks how many retrievers are currently using the pool. +var refCount int +var refCountMutex sync.Mutex + +func GetPool(ctx context.Context, uri string) (*pgxpool.Pool, error) { + // The sync.Once ensures that the inner function is executed only once, + // even if called by multiple goroutines concurrently. + once.Do(func() { + poolInstance, poolErr = pgxpool.New(ctx, uri) + if poolErr != nil { + poolErr = fmt.Errorf("failed to create connection pool: %w", poolErr) + return + } + + // Check connection immediately + if err := poolInstance.Ping(ctx); err != nil { + poolErr = fmt.Errorf("failed to ping database with new pool: %w", err) + // Don't close here, the pool remains valid for a retry connection + } + }) + + if poolErr != nil { + return nil, poolErr + } + + refCountMutex.Lock() + refCount++ + refCountMutex.Unlock() + + return poolInstance, nil +} + +func ReleasePool() { + refCountMutex.Lock() + defer refCountMutex.Unlock() + + if refCount > 0 { + refCount-- + } + + // Only close the physical connection when the last reference is released. + if refCount == 0 && poolInstance != nil { + poolInstance.Close() + poolInstance = nil + // Reset sync.Once to allow re-initialization if needed + once = sync.Once{} + poolErr = nil + } +} + diff --git a/retriever/postgresqlretriever/retriever.go b/retriever/postgresqlretriever/retriever.go index 7a8fa5c484a..b97a6fc34e8 100644 --- a/retriever/postgresqlretriever/retriever.go +++ b/retriever/postgresqlretriever/retriever.go @@ -5,8 +5,8 @@ import ( "encoding/json" "fmt" "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 +26,7 @@ type Retriever struct { logger *fflog.FFLogger status retriever.Status columns map[string]string - conn *pgx.Conn + pool *pgxpool.Pool flagset *string } @@ -40,23 +40,19 @@ func (r *Retriever) Init(ctx context.Context, logger *fflog.FFLogger, flagset *s 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 } @@ -71,37 +67,36 @@ func (r *Retriever) Status() retriever.Status { // 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 +116,7 @@ func (r *Retriever) Retrieve(ctx context.Context) ([]byte, error) { 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) } From 7a3eb537d049b3e7fdd77fc90e92d51792e6810b Mon Sep 17 00:00:00 2001 From: Georgi Dimitrov Date: Tue, 2 Dec 2025 18:47:49 +0200 Subject: [PATCH 02/12] format with gofmt --- retriever/postgresqlretriever/database.go | 5 ++--- retriever/postgresqlretriever/retriever.go | 4 ++-- retriever/postgresqlretriever/retriever_test.go | 2 +- 3 files changed, 5 insertions(+), 6 deletions(-) diff --git a/retriever/postgresqlretriever/database.go b/retriever/postgresqlretriever/database.go index 1d275c7f0b6..89e2336d223 100644 --- a/retriever/postgresqlretriever/database.go +++ b/retriever/postgresqlretriever/database.go @@ -25,7 +25,7 @@ func GetPool(ctx context.Context, uri string) (*pgxpool.Pool, error) { poolErr = fmt.Errorf("failed to create connection pool: %w", poolErr) return } - + // Check connection immediately if err := poolInstance.Ping(ctx); err != nil { poolErr = fmt.Errorf("failed to ping database with new pool: %w", err) @@ -57,8 +57,7 @@ func ReleasePool() { poolInstance.Close() poolInstance = nil // Reset sync.Once to allow re-initialization if needed - once = sync.Once{} + once = sync.Once{} poolErr = nil } } - diff --git a/retriever/postgresqlretriever/retriever.go b/retriever/postgresqlretriever/retriever.go index b97a6fc34e8..8ab93f8c00b 100644 --- a/retriever/postgresqlretriever/retriever.go +++ b/retriever/postgresqlretriever/retriever.go @@ -4,12 +4,12 @@ import ( "context" "encoding/json" "fmt" - "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" + "log/slog" ) var defaultColumns = map[string]string{ @@ -68,7 +68,7 @@ func (r *Retriever) Status() retriever.Status { // Shutdown closes the database connection. func (r *Retriever) Shutdown(ctx context.Context) error { ReleasePool() - return nil + return nil } // Retrieve fetches flag configuration from PostgreSQL. diff --git a/retriever/postgresqlretriever/retriever_test.go b/retriever/postgresqlretriever/retriever_test.go index dbe40a157e5..baabede471c 100644 --- a/retriever/postgresqlretriever/retriever_test.go +++ b/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) { From a0082b1f648eda8b61b9733bd4fa5b1d4918a48e Mon Sep 17 00:00:00 2001 From: Georgi Dimitrov Date: Tue, 2 Dec 2025 18:56:39 +0200 Subject: [PATCH 03/12] fix imports --- retriever/postgresqlretriever/database.go | 24 +++++++++++----------- retriever/postgresqlretriever/retriever.go | 20 ++++++++++-------- 2 files changed, 23 insertions(+), 21 deletions(-) diff --git a/retriever/postgresqlretriever/database.go b/retriever/postgresqlretriever/database.go index 89e2336d223..e6f2eb20696 100644 --- a/retriever/postgresqlretriever/database.go +++ b/retriever/postgresqlretriever/database.go @@ -1,16 +1,16 @@ package postgresqlretriever import ( - "context" - "fmt" - "sync" + "context" + "fmt" + "sync" - "github.com/jackc/pgx/v5/pgxpool" + "github.com/jackc/pgx/v5/pgxpool" ) var poolInstance *pgxpool.Pool var once sync.Once -var poolErr error +var errPool error // refCount tracks how many retrievers are currently using the pool. var refCount int @@ -20,21 +20,21 @@ func GetPool(ctx context.Context, uri string) (*pgxpool.Pool, error) { // The sync.Once ensures that the inner function is executed only once, // even if called by multiple goroutines concurrently. once.Do(func() { - poolInstance, poolErr = pgxpool.New(ctx, uri) - if poolErr != nil { - poolErr = fmt.Errorf("failed to create connection pool: %w", poolErr) + poolInstance, errPool = pgxpool.New(ctx, uri) + if errPool != nil { + errPool = fmt.Errorf("failed to create connection pool: %w", errPool) return } // Check connection immediately if err := poolInstance.Ping(ctx); err != nil { - poolErr = fmt.Errorf("failed to ping database with new pool: %w", err) + errPool = fmt.Errorf("failed to ping database with new pool: %w", err) // Don't close here, the pool remains valid for a retry connection } }) - if poolErr != nil { - return nil, poolErr + if errPool != nil { + return nil, errPool } refCountMutex.Lock() @@ -58,6 +58,6 @@ func ReleasePool() { poolInstance = nil // Reset sync.Once to allow re-initialization if needed once = sync.Once{} - poolErr = nil + errPool = nil } } diff --git a/retriever/postgresqlretriever/retriever.go b/retriever/postgresqlretriever/retriever.go index 8ab93f8c00b..4a15ebe3580 100644 --- a/retriever/postgresqlretriever/retriever.go +++ b/retriever/postgresqlretriever/retriever.go @@ -1,15 +1,17 @@ package postgresqlretriever import ( - "context" - "encoding/json" - "fmt" - "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" - "log/slog" + "context" + "encoding/json" + "fmt" + "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" ) var defaultColumns = map[string]string{ From e9ff21c7f1a259c9a4c4e3e61cecbc89baf12c91 Mon Sep 17 00:00:00 2001 From: Georgi Dimitrov Date: Tue, 2 Dec 2025 19:20:40 +0200 Subject: [PATCH 04/12] fix formatting issues --- retriever/postgresqlretriever/database.go | 8 ++++---- retriever/postgresqlretriever/retriever.go | 21 ++++++++++----------- 2 files changed, 14 insertions(+), 15 deletions(-) diff --git a/retriever/postgresqlretriever/database.go b/retriever/postgresqlretriever/database.go index e6f2eb20696..a16c56b49fd 100644 --- a/retriever/postgresqlretriever/database.go +++ b/retriever/postgresqlretriever/database.go @@ -1,11 +1,11 @@ package postgresqlretriever import ( - "context" - "fmt" - "sync" + "context" + "fmt" + "sync" - "github.com/jackc/pgx/v5/pgxpool" + "github.com/jackc/pgx/v5/pgxpool" ) var poolInstance *pgxpool.Pool diff --git a/retriever/postgresqlretriever/retriever.go b/retriever/postgresqlretriever/retriever.go index 4a15ebe3580..b0bb93397c6 100644 --- a/retriever/postgresqlretriever/retriever.go +++ b/retriever/postgresqlretriever/retriever.go @@ -1,17 +1,16 @@ package postgresqlretriever import ( - "context" - "encoding/json" - "fmt" - "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" + "context" + "encoding/json" + "fmt" + "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" ) var defaultColumns = map[string]string{ From 117ae1a132ec7f2e3536dcb8d7c6f112c9d86c1f Mon Sep 17 00:00:00 2001 From: Georgi Dimitrov Date: Tue, 2 Dec 2025 22:58:04 +0200 Subject: [PATCH 05/12] simplify postgres db initalization --- retriever/postgresqlretriever/database.go | 83 ++++++++++------------- 1 file changed, 35 insertions(+), 48 deletions(-) diff --git a/retriever/postgresqlretriever/database.go b/retriever/postgresqlretriever/database.go index a16c56b49fd..9e79c75907e 100644 --- a/retriever/postgresqlretriever/database.go +++ b/retriever/postgresqlretriever/database.go @@ -2,62 +2,49 @@ package postgresqlretriever import ( "context" - "fmt" "sync" "github.com/jackc/pgx/v5/pgxpool" ) -var poolInstance *pgxpool.Pool -var once sync.Once -var errPool error - -// refCount tracks how many retrievers are currently using the pool. -var refCount int -var refCountMutex sync.Mutex +var ( + pool *pgxpool.Pool + mu sync.Mutex + refCount int +) func GetPool(ctx context.Context, uri string) (*pgxpool.Pool, error) { - // The sync.Once ensures that the inner function is executed only once, - // even if called by multiple goroutines concurrently. - once.Do(func() { - poolInstance, errPool = pgxpool.New(ctx, uri) - if errPool != nil { - errPool = fmt.Errorf("failed to create connection pool: %w", errPool) - return - } - - // Check connection immediately - if err := poolInstance.Ping(ctx); err != nil { - errPool = fmt.Errorf("failed to ping database with new pool: %w", err) - // Don't close here, the pool remains valid for a retry connection - } - }) - - if errPool != nil { - return nil, errPool - } - - refCountMutex.Lock() - refCount++ - refCountMutex.Unlock() - - return poolInstance, nil + 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() { - refCountMutex.Lock() - defer refCountMutex.Unlock() - - if refCount > 0 { - refCount-- - } - - // Only close the physical connection when the last reference is released. - if refCount == 0 && poolInstance != nil { - poolInstance.Close() - poolInstance = nil - // Reset sync.Once to allow re-initialization if needed - once = sync.Once{} - errPool = nil - } + mu.Lock() + defer mu.Unlock() + + refCount-- + if refCount <= 0 { + if pool != nil { + pool.Close() + pool = nil + } + refCount = 0 + } } + From 4617a05e1336b72d5961e194d4cc1173d1beb246 Mon Sep 17 00:00:00 2001 From: Georgi Dimitrov Date: Tue, 2 Dec 2025 23:05:55 +0200 Subject: [PATCH 06/12] format --- retriever/postgresqlretriever/database.go | 65 +++++++++++------------ 1 file changed, 32 insertions(+), 33 deletions(-) diff --git a/retriever/postgresqlretriever/database.go b/retriever/postgresqlretriever/database.go index 9e79c75907e..14b3addded2 100644 --- a/retriever/postgresqlretriever/database.go +++ b/retriever/postgresqlretriever/database.go @@ -8,43 +8,42 @@ import ( ) var ( - pool *pgxpool.Pool - mu sync.Mutex - refCount int + 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 + 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 - } + mu.Lock() + defer mu.Unlock() + + refCount-- + if refCount <= 0 { + if pool != nil { + pool.Close() + pool = nil + } + refCount = 0 + } } - From 5d97bb002a96a8278ca19557e6d33dd38b6f2676 Mon Sep 17 00:00:00 2001 From: Georgi Dimitrov Date: Mon, 8 Dec 2025 15:06:11 +0200 Subject: [PATCH 07/12] multi postgres pools support --- retriever/postgresqlretriever/database.go | 49 ---------------- retriever/postgresqlretriever/postgres.go | 66 ++++++++++++++++++++++ retriever/postgresqlretriever/retriever.go | 6 +- 3 files changed, 71 insertions(+), 50 deletions(-) delete mode 100644 retriever/postgresqlretriever/database.go create mode 100644 retriever/postgresqlretriever/postgres.go diff --git a/retriever/postgresqlretriever/database.go b/retriever/postgresqlretriever/database.go deleted file mode 100644 index 14b3addded2..00000000000 --- a/retriever/postgresqlretriever/database.go +++ /dev/null @@ -1,49 +0,0 @@ -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 - } -} diff --git a/retriever/postgresqlretriever/postgres.go b/retriever/postgresqlretriever/postgres.go new file mode 100644 index 00000000000..7b9ef0c62fa --- /dev/null +++ b/retriever/postgresqlretriever/postgres.go @@ -0,0 +1,66 @@ +package postgresqlretriever + +import ( + "context" + "sync" + + "github.com/jackc/pgx/v5/pgxpool" +) + +type poolEntry struct { + pool *pgxpool.Pool + refCount int +} + +var ( + mu sync.Mutex + poolMap = make(map[string]*poolEntry) +) + +// GetPool returns a pool for a given URI, creating it if needed. +func GetPool(ctx context.Context, uri string) (*pgxpool.Pool, error) { + mu.Lock() + defer mu.Unlock() + + // If already exists, bump refCount + if entry, ok := poolMap[uri]; ok { + entry.refCount++ + return entry.pool, nil + } + + // Create a new pool + pool, err := pgxpool.New(ctx, uri) + if err != nil { + return nil, err + } + if err := pool.Ping(ctx); err != nil { + pool.Close() + return nil, err + } + + poolMap[uri] = &poolEntry{ + pool: pool, + refCount: 1, + } + + return pool, nil +} + + +// ReleasePool decreases refCount and closes/removes when it hits zero. +func ReleasePool(uri string) { + mu.Lock() + defer mu.Unlock() + + entry, ok := poolMap[uri] + if !ok { + return // nothing to do + } + + entry.refCount-- + if entry.refCount <= 0 { + entry.pool.Close() + delete(poolMap, uri) + } +} + diff --git a/retriever/postgresqlretriever/retriever.go b/retriever/postgresqlretriever/retriever.go index b0bb93397c6..febc52db524 100644 --- a/retriever/postgresqlretriever/retriever.go +++ b/retriever/postgresqlretriever/retriever.go @@ -68,10 +68,14 @@ func (r *Retriever) Status() retriever.Status { // Shutdown closes the database connection. func (r *Retriever) Shutdown(ctx context.Context) error { - ReleasePool() + if r.pool != nil { + ReleasePool(r.URI) + r.pool = nil + } return nil } + // Retrieve fetches flag configuration from PostgreSQL. func (r *Retriever) Retrieve(ctx context.Context) ([]byte, error) { if r.pool == nil { From cde16643b3fe984598cd6beec965e14e4389fd1b Mon Sep 17 00:00:00 2001 From: Georgi Dimitrov Date: Mon, 8 Dec 2025 15:11:37 +0200 Subject: [PATCH 08/12] format issues --- retriever/postgresqlretriever/postgres.go | 6 ++---- retriever/postgresqlretriever/retriever.go | 1 - 2 files changed, 2 insertions(+), 5 deletions(-) diff --git a/retriever/postgresqlretriever/postgres.go b/retriever/postgresqlretriever/postgres.go index 7b9ef0c62fa..f6299a5653f 100644 --- a/retriever/postgresqlretriever/postgres.go +++ b/retriever/postgresqlretriever/postgres.go @@ -13,8 +13,8 @@ type poolEntry struct { } var ( - mu sync.Mutex - poolMap = make(map[string]*poolEntry) + mu sync.Mutex + poolMap = make(map[string]*poolEntry) ) // GetPool returns a pool for a given URI, creating it if needed. @@ -46,7 +46,6 @@ func GetPool(ctx context.Context, uri string) (*pgxpool.Pool, error) { return pool, nil } - // ReleasePool decreases refCount and closes/removes when it hits zero. func ReleasePool(uri string) { mu.Lock() @@ -63,4 +62,3 @@ func ReleasePool(uri string) { delete(poolMap, uri) } } - diff --git a/retriever/postgresqlretriever/retriever.go b/retriever/postgresqlretriever/retriever.go index febc52db524..6a1bbe7ba52 100644 --- a/retriever/postgresqlretriever/retriever.go +++ b/retriever/postgresqlretriever/retriever.go @@ -75,7 +75,6 @@ func (r *Retriever) Shutdown(ctx context.Context) error { return nil } - // Retrieve fetches flag configuration from PostgreSQL. func (r *Retriever) Retrieve(ctx context.Context) ([]byte, error) { if r.pool == nil { From 28525772df21a328fc4e61f9691cd6ae406ab713 Mon Sep 17 00:00:00 2001 From: Georgi Dimitrov Date: Mon, 8 Dec 2025 15:53:03 +0200 Subject: [PATCH 09/12] Pass ctx to DB shutdown --- retriever/postgresqlretriever/postgres.go | 2 +- retriever/postgresqlretriever/retriever.go | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/retriever/postgresqlretriever/postgres.go b/retriever/postgresqlretriever/postgres.go index f6299a5653f..b2c2daa4ecf 100644 --- a/retriever/postgresqlretriever/postgres.go +++ b/retriever/postgresqlretriever/postgres.go @@ -47,7 +47,7 @@ func GetPool(ctx context.Context, uri string) (*pgxpool.Pool, error) { } // ReleasePool decreases refCount and closes/removes when it hits zero. -func ReleasePool(uri string) { +func ReleasePool(ctx context.Context, uri string) { mu.Lock() defer mu.Unlock() diff --git a/retriever/postgresqlretriever/retriever.go b/retriever/postgresqlretriever/retriever.go index 6a1bbe7ba52..e31b3747bb4 100644 --- a/retriever/postgresqlretriever/retriever.go +++ b/retriever/postgresqlretriever/retriever.go @@ -69,7 +69,7 @@ func (r *Retriever) Status() retriever.Status { // Shutdown closes the database connection. func (r *Retriever) Shutdown(ctx context.Context) error { if r.pool != nil { - ReleasePool(r.URI) + ReleasePool(ctx, r.URI) r.pool = nil } return nil From 228208a8b89c5327dc86d9273e817d782127295d Mon Sep 17 00:00:00 2001 From: Georgi Dimitrov Date: Thu, 11 Dec 2025 13:38:41 +0200 Subject: [PATCH 10/12] multiple pools test --- .../postgresqlretriever/postgres_test.go | 85 +++++++++++++++++++ 1 file changed, 85 insertions(+) create mode 100644 retriever/postgresqlretriever/postgres_test.go diff --git a/retriever/postgresqlretriever/postgres_test.go b/retriever/postgresqlretriever/postgres_test.go new file mode 100644 index 00000000000..2287b5e2ad6 --- /dev/null +++ b/retriever/postgresqlretriever/postgres_test.go @@ -0,0 +1,85 @@ +package postgresqlretriever_test + +import ( + "context" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/testcontainers/testcontainers-go" + "github.com/testcontainers/testcontainers-go/wait" + rr "github.com/thomaspoignant/go-feature-flag/retriever/postgresqlretriever" +) + +func TestGetPool_MultipleURIsAndReuse(t *testing.T) { + ctx := context.Background() + + // Setup Container + req := testcontainers.ContainerRequest{ + Image: "postgres:15-alpine", + ExposedPorts: []string{"5432/tcp"}, + Env: map[string]string{"POSTGRES_PASSWORD": "password"}, + // This waits until the log says the system is ready, preventing connection errors + WaitingFor: wait.ForAll( + wait.ForLog("database system is ready to accept connections").WithStartupTimeout(10*time.Second), + wait.ForListeningPort("5432/tcp").WithStartupTimeout(10*time.Second), + ), + } + + pg, err := testcontainers.GenericContainer(ctx, testcontainers.GenericContainerRequest{ + ContainerRequest: req, + Started: true, + }) + assert.NoError(t, err) + defer pg.Terminate(ctx) + + endpoint, err := pg.Endpoint(ctx, "") + assert.NoError(t, err) + baseURI := "postgres://postgres:password@" + endpoint + "/postgres?sslmode=disable" + + // Different URIs + uri1 := baseURI + "&application_name=A" + uri2 := baseURI + "&application_name=B" + + // First Calls (RefCount = 1) + pool1a, err := rr.GetPool(ctx, uri1) + assert.NoError(t, err) + assert.NotNil(t, pool1a) + + pool2a, err := rr.GetPool(ctx, uri2) + assert.NoError(t, err) + assert.NotNil(t, pool2a) + + // Verify distinct pools + assert.NotEqual(t, pool1a, pool2a) + + // Reuse Logic (RefCount = 2 for URI1) + pool1b, err := rr.GetPool(ctx, uri1) + assert.NoError(t, err) + assert.Equal(t, pool1a, pool1b, "Should return exact same pool instance") + + // Release Logic + // URI1 RefCount: 2 -> 1 + rr.ReleasePool(ctx, uri1) + + // URI1 RefCount: 1 -> 0 (Closed & Removed) + rr.ReleasePool(ctx, uri1) + + // Recreation Logic + // URI1 should now create a NEW pool + pool1c, err := rr.GetPool(ctx, uri1) + assert.NoError(t, err) + assert.NotEqual(t, pool1a, pool1c, "Should be a new pool instance after full release") + + // Cleanup new pool + rr.ReleasePool(ctx, uri1) + + // URI2 Cleanup verification + rr.ReleasePool(ctx, uri2) // RefCount -> 0 + + pool2b, err := rr.GetPool(ctx, uri2) + assert.NoError(t, err) + assert.NotEqual(t, pool2a, pool2b, "URI2 should be recreated") + + rr.ReleasePool(ctx, uri2) +} From e3f6a1184161bd5affc8d628bb35beecfbf859ce Mon Sep 17 00:00:00 2001 From: Georgi Dimitrov Date: Thu, 11 Dec 2025 13:51:53 +0200 Subject: [PATCH 11/12] merge main --- go.mod | 1 + 1 file changed, 1 insertion(+) diff --git a/go.mod b/go.mod index 6f3a9b61a66..86f19e5265e 100644 --- a/go.mod +++ b/go.mod @@ -179,6 +179,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.6.0 // indirect github.com/jcmturner/aescts/v2 v2.0.0 // indirect github.com/jcmturner/dnsutils/v2 v2.0.0 // indirect From b990bbedf1feb9d86485605c18cc7230b8570084 Mon Sep 17 00:00:00 2001 From: Georgi Dimitrov Date: Mon, 15 Dec 2025 17:27:03 +0200 Subject: [PATCH 12/12] change error message / fix test --- retriever/postgresqlretriever/retriever_test.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/retriever/postgresqlretriever/retriever_test.go b/retriever/postgresqlretriever/retriever_test.go index baabede471c..f788a11b522 100644 --- a/retriever/postgresqlretriever/retriever_test.go +++ b/retriever/postgresqlretriever/retriever_test.go @@ -215,7 +215,7 @@ func TestRetrieverErrorHandling(t *testing.T) { _, err := r.Retrieve(context.Background()) assert.Error(t, err) - assert.Contains(t, err.Error(), "database connection is not initialized") + assert.Contains(t, err.Error(), "database connection pool is not initialized") }) t.Run("Status - nil receiver", func(t *testing.T) {