@@ -3,8 +3,10 @@ package adapter
33import (
44 "context"
55 "fmt"
6+ "strings"
67 "time"
78
9+ "github.com/exaring/otelpgx"
810 "github.com/jackc/pgx/v5"
911 "github.com/jackc/pgx/v5/pgconn"
1012 "github.com/jackc/pgx/v5/pgxpool"
@@ -25,6 +27,7 @@ type PostgresConfig struct {
2527 MaxConnIdleTime time.Duration
2628 SearchPath string
2729 Timezone string
30+ EnableOTel bool
2831}
2932
3033// NewPostgres creates a new database connection pool.
@@ -69,6 +72,17 @@ func NewPostgres(cfg PostgresConfig) (*PostgresDB, error) {
6972 poolConfig .ConnConfig .RuntimeParams ["TimeZone" ] = cfg .Timezone
7073 poolConfig .ConnConfig .RuntimeParams ["timezone" ] = cfg .Timezone
7174
75+ // Initialize OpenTelemetry tracing for pgx if enabled
76+ // Uses otelpgx package to create a tracer that trims SQL in span names
77+ // and uses a custom function to generate span names from statements
78+ // See pgxSpanNameFunc below for details
79+ if cfg .EnableOTel {
80+ poolConfig .ConnConfig .Tracer = otelpgx .NewTracer (
81+ otelpgx .WithTrimSQLInSpanName (),
82+ otelpgx .WithSpanNameFunc (pgxSpanNameFunc ),
83+ )
84+ }
85+
7286 pool , err := pgxpool .NewWithConfig (context .Background (), poolConfig )
7387 if err != nil {
7488 return nil , fmt .Errorf ("failed to create pgxpool: %w" , err )
@@ -170,8 +184,8 @@ func (db *PostgresDB) GetPoolConnection(ctx context.Context) (*pgxpool.Conn, err
170184 return db .Pool .Acquire (ctx )
171185}
172186
173- // TestConnection tests database connectivity with timeout
174- func TestConnection (cfg PostgresConfig ) error {
187+ // CheckPostgresConnection tests database connectivity with timeout
188+ func CheckPostgresConnection (cfg PostgresConfig ) error {
175189 ctx , cancel := context .WithTimeout (context .Background (), 10 * time .Second )
176190 defer cancel ()
177191
@@ -193,7 +207,7 @@ func TestConnection(cfg PostgresConfig) error {
193207}
194208
195209// NewSingleConnection creates and returns a single pgx.Conn using the provided PostgresConfig.
196- // This is used by TestConnection and other helpers that need a standalone connection.
210+ // This is used by CheckPostgresConnection and other helpers that need a standalone connection.
197211func NewSingleConnection (cfg PostgresConfig ) (* pgx.Conn , error ) {
198212 // apply sensible defaults used elsewhere
199213 if cfg .SearchPath == "" {
@@ -221,9 +235,31 @@ func NewSingleConnection(cfg PostgresConfig) (*pgx.Conn, error) {
221235 connCfg .ConnectTimeout = 10 * time .Second
222236 }
223237
238+ // Initialize OpenTelemetry tracing for pgx if enabled
239+ // Uses otelpgx package to create a tracer that trims SQL in span names
240+ // and uses a custom function to generate span names from statements
241+ // See pgxSpanNameFunc below for details
242+ if cfg .EnableOTel {
243+ connCfg .Tracer = otelpgx .NewTracer (
244+ otelpgx .WithTrimSQLInSpanName (),
245+ otelpgx .WithSpanNameFunc (pgxSpanNameFunc ),
246+ )
247+ }
248+
224249 conn , err := pgx .ConnectConfig (context .Background (), connCfg )
225250 if err != nil {
226251 return nil , fmt .Errorf ("failed to create single connection: %w" , err )
227252 }
228253 return conn , nil
229254}
255+
256+ // pgxSpanNameFunc trims "-- name: " prefix and returns the statement name for tracing
257+ func pgxSpanNameFunc (stmt string ) string {
258+ // If stmt is of the sqlc form "-- name: Example :one\n...",
259+ // extract "Example". Otherwise, leave as-is.
260+ stmt = strings .TrimPrefix (stmt , "-- name: " )
261+ if i := strings .IndexByte (stmt , ' ' ); i != - 1 {
262+ return stmt [:i ]
263+ }
264+ return stmt
265+ }
0 commit comments