Skip to content

feat: shared pgx-pool #4393

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
gogbog wants to merge 14 commits into
base: main
Choose a base branch
from
Open

feat: shared pgx-pool #4393

Show file tree
Hide file tree
Changes from 6 commits
Commits
Show all changes
14 commits
Select commit Hold shift + click to select a range
6823c96
use pgxpool and singelton to initialize it just once
Dec 2, 2025
7a3eb53
format with gofmt
Dec 2, 2025
a0082b1
fix imports
Dec 2, 2025
e9ff21c
fix formatting issues
Dec 2, 2025
117ae1a
simplify postgres db initalization
Dec 2, 2025
4617a05
format
Dec 2, 2025
5d97bb0
multi postgres pools support
Dec 8, 2025
cde1664
format issues
Dec 8, 2025
2852577
Pass ctx to DB shutdown
Dec 8, 2025
fa090e2
Merge branch 'main' into feature/shared-pgxpool
gogbog Dec 8, 2025
228208a
multiple pools test
Dec 11, 2025
de4ed7f
Merge branch 'main' into feature/shared-pgxpool
Dec 11, 2025
e3f6a11
merge main
Dec 11, 2025
a0bd5ce
Merge branch 'main' into feature/shared-pgxpool
thomaspoignant Dec 13, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions go.mod
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
1 change: 0 additions & 1 deletion go.sum
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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=
Expand Down
49 changes: 49 additions & 0 deletions retriever/postgresqlretriever/database.go
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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) {

Check warning on line 16 in retriever/postgresqlretriever/database.go

View check run for this annotation

SonarQubeCloud / SonarCloud Code Analysis

Remove the 'Get' prefix from this function name. 


See more on https://sonarcloud.io/project/issues?id=thomaspoignant_go-feature-flag&issues=AZrf8MN0sn8GrSqJf98R&open=AZrf8MN0sn8GrSqJf98R&pullRequest=4393
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
}
}
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot Dec 3, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

critical

The current singleton implementation for the connection pool doesn't support multiple database URIs. The global pool variable will be initialized once with the URI from the first retriever, and all subsequent retrievers, regardless of their configured URI, will receive this same pool. This can lead to retrievers connecting to the wrong database.

To fix this, the pool management should be URI-aware. I suggest using maps to store a connection pool and a reference count for each unique URI. This ensures that each database has its own dedicated pool, which is shared only by retrievers connecting to that same database.

package postgresqlretriever

import (
	"context"
	"sync"

	"github.com/jackc/pgx/v5/pgxpool"
)

var (
	pools     = make(map[string]*pgxpool.Pool)
	refCounts = make(map[string]int)
	mu        sync.Mutex
)

// GetPool retrieves a shared connection pool for the given URI.
// It creates a new pool if one doesn't exist for the URI, otherwise, it returns the existing one.
// It uses reference counting to manage the lifecycle of the pool.
func GetPool(ctx context.Context, uri string) (*pgxpool.Pool, error) {
	mu.Lock()
	defer mu.Unlock()

	// If a pool for this URI already exists, increment its ref count and return it.
	if pool, ok := pools[uri]; ok {
		refCounts[uri]++
		return pool, nil
	}

	// Otherwise, create a new pool for this URI.
	p, err := pgxpool.New(ctx, uri)
	if err != nil {
		return nil, err
	}
	if err := p.Ping(ctx); err != nil {
		p.Close()
		return nil, err
	}

	pools[uri] = p
	refCounts[uri] = 1
	return p, nil
}

// ReleasePool decrements the reference count for a pool.
// If the reference count drops to zero, the pool is closed and removed.
func ReleasePool(uri string) {
	mu.Lock()
	defer mu.Unlock()

	// Only proceed if the URI is being tracked.
	if _, ok := refCounts[uri]; !ok {
		return
	}

	refCounts[uri]--
	if refCounts[uri] <= 0 {
		if pool, ok := pools[uri]; ok {
			pool.Close()
			delete(pools, uri)
		}
		delete(refCounts, uri)
	}
}

35 changes: 16 additions & 19 deletions retriever/postgresqlretriever/retriever.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -26,7 +27,7 @@
logger *fflog.FFLogger
status retriever.Status
columns map[string]string
conn *pgx.Conn
pool *pgxpool.Pool
flagset *string
}

Expand All @@ -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
}
Expand All @@ -70,38 +67,37 @@
}

// Shutdown closes the database connection.
func (r *Retriever) Shutdown(ctx context.Context) error {

Check warning on line 70 in retriever/postgresqlretriever/retriever.go

View check run for this annotation

SonarQubeCloud / SonarCloud Code Analysis

Unused parameter 'ctx' should be removed. 


See more on https://sonarcloud.io/project/issues?id=thomaspoignant_go-feature-flag&issues=AZrf8MI6sn8GrSqJf98Q&open=AZrf8MI6sn8GrSqJf98Q&pullRequest=4393
if r.conn == nil {
return nil
}
return r.conn.Close(ctx)
ReleasePool()
return nil
}
Comment on lines 70 to 77
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot Dec 3, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

To support the URI-aware pool manager, the Shutdown function should pass the retriever's URI to ReleasePool. This ensures that the reference count for the correct connection pool is decremented.

I've also added a nil check for r.pool to make the shutdown process more robust, preventing a call to ReleasePool if the retriever's initialization failed before the pool was assigned.

Suggested change
func (r *Retriever) Shutdown(ctx context.Context) error {
if r.conn == nil {
return nil
}
return r.conn.Close(ctx)
ReleasePool()
return nil
}
func (r *Retriever) Shutdown(ctx context.Context) error {
if r.pool != nil {
ReleasePool(r.URI)
}
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
Expand All @@ -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)
}
Expand Down
2 changes: 1 addition & 1 deletion retriever/postgresqlretriever/retriever_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down
Loading