File size: 2,587 Bytes
fea99b3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 | package pgdriver
import (
"context"
"database/sql"
"fmt"
"time"
"github.com/XSAM/otelsql"
"github.com/jackc/pgx/v5/pgxpool"
pgxstdlib "github.com/jackc/pgx/v5/stdlib"
"go.opentelemetry.io/otel/metric"
semconv "go.opentelemetry.io/otel/semconv/v1.27.0"
"go.opentelemetry.io/otel/trace"
"github.com/openmeterio/openmeter/pkg/pgxpoolobserver"
)
type Option interface {
apply(*options)
}
type optionFunc func(c *options)
func (fn optionFunc) apply(c *options) {
fn(c)
}
func WithTracerProvider(p trace.TracerProvider) Option {
return optionFunc(func(o *options) {
o.otelOptions = append(o.otelOptions, otelsql.WithTracerProvider(p))
})
}
func WithMeterProvider(p metric.MeterProvider) Option {
return optionFunc(func(o *options) {
o.otelOptions = append(o.otelOptions, otelsql.WithMeterProvider(p))
})
}
func WithMetricMeter(m metric.Meter) Option {
return optionFunc(func(o *options) {
o.metricMeter = m
})
}
func WithSpanOptions(opt otelsql.SpanOptions) Option {
return optionFunc(func(o *options) {
o.otelOptions = append(o.otelOptions, otelsql.WithSpanOptions(opt))
})
}
func WithLockTimeout(timeout time.Duration) Option {
return optionFunc(func(o *options) {
o.connConfig.ConnConfig.RuntimeParams["lock_timeout"] = fmt.Sprintf("%d", timeout.Milliseconds())
})
}
type options struct {
connConfig *pgxpool.Config
otelOptions []otelsql.Option
metricMeter metric.Meter
}
type Driver struct {
pool *pgxpool.Pool
db *sql.DB
}
func (d *Driver) DB() *sql.DB {
return d.db
}
func (d *Driver) Close() error {
d.pool.Close()
return nil
}
func NewPostgresDriver(ctx context.Context, url string, opts ...Option) (*Driver, error) {
config, err := pgxpool.ParseConfig(url)
if err != nil {
return nil, fmt.Errorf("failed to parse postgres url: %w", err)
}
o := &options{
connConfig: config,
otelOptions: []otelsql.Option{
otelsql.WithAttributes(
semconv.DBSystemPostgreSQL,
),
},
}
for _, opt := range opts {
opt.apply(o)
}
pool, err := pgxpool.NewWithConfig(ctx, o.connConfig)
if err != nil {
return nil, fmt.Errorf("failed to create postgres pool: %w", err)
}
if o.metricMeter != nil {
if err := pgxpoolobserver.ObservePoolMetrics(o.metricMeter, pool); err != nil {
return nil, err
}
}
db := otelsql.OpenDB(pgxstdlib.GetPoolConnector(pool), o.otelOptions...)
// Set maximum idle connections to 0 as connections are managed from pgx.Pool.
// See: https://github.com/jackc/pgx/blob/v5.6.0/stdlib/sql.go#L204-L208
db.SetMaxIdleConns(0)
return &Driver{
pool: pool,
db: db,
}, nil
}
|