feat: use postgres

pkg/postgres/open.go

@@ -0,0 +1,136 @@
+package postgres
+
+import (
+	"context"
+	"database/sql"
+	"errors"
+	"fmt"
+
+	"github.com/XSAM/otelsql"
+	"github.com/jackc/pgx/v5"
+	"github.com/jackc/pgx/v5/stdlib"
+	"go.opentelemetry.io/otel/attribute"
+	"go.opentelemetry.io/otel/metric"
+	"go.opentelemetry.io/otel/trace"
+)
+
+// dbSystemAttribute identifies the wrapped backend in OpenTelemetry spans
+// without locking the package to a specific semconv release.
+var dbSystemAttribute = attribute.String("db.system", "postgresql")
+
+// Option configures the OpenTelemetry providers attached to a connection by
+// OpenPrimary, OpenReplicas, and InstrumentDBStats. Unset providers fall back
+// to the OpenTelemetry global providers.
+type Option func(*options)
+
+type options struct {
+	tracerProvider trace.TracerProvider
+	meterProvider  metric.MeterProvider
+}
+
+// WithTracerProvider sets the tracer provider used for SQL statement spans.
+func WithTracerProvider(tp trace.TracerProvider) Option {
+	return func(o *options) {
+		o.tracerProvider = tp
+	}
+}
+
+// WithMeterProvider sets the meter provider used for connection-pool stats.
+func WithMeterProvider(mp metric.MeterProvider) Option {
+	return func(o *options) {
+		o.meterProvider = mp
+	}
+}
+
+func evalOptions(opts []Option) options {
+	var resolved options
+	for _, opt := range opts {
+		if opt == nil {
+			continue
+		}
+		opt(&resolved)
+	}
+	return resolved
+}
+
+func (o options) otelsqlOpenOptions() []otelsql.Option {
+	out := []otelsql.Option{otelsql.WithAttributes(dbSystemAttribute)}
+	if o.tracerProvider != nil {
+		out = append(out, otelsql.WithTracerProvider(o.tracerProvider))
+	}
+	if o.meterProvider != nil {
+		out = append(out, otelsql.WithMeterProvider(o.meterProvider))
+	}
+	return out
+}
+
+// OpenPrimary opens the primary `*sql.DB` from cfg. ctx bounds individual
+// pgx connect attempts via the parsed pgx config's ConnectTimeout (set to
+// cfg.OperationTimeout). The returned pool has SetMaxOpenConns,
+// SetMaxIdleConns and SetConnMaxLifetime applied.
+func OpenPrimary(ctx context.Context, cfg Config, opts ...Option) (*sql.DB, error) {
+	if err := cfg.Validate(); err != nil {
+		return nil, fmt.Errorf("open postgres primary: %w", err)
+	}
+	db, err := openDB(ctx, cfg, cfg.PrimaryDSN, evalOptions(opts))
+	if err != nil {
+		return nil, fmt.Errorf("open postgres primary: %w", err)
+	}
+	return db, nil
+}
+
+// OpenReplicas opens one `*sql.DB` per replica DSN. It returns nil when no
+// replicas are configured. When opening a replica fails mid-way, every
+// already-opened replica is closed before returning the error.
+func OpenReplicas(ctx context.Context, cfg Config, opts ...Option) ([]*sql.DB, error) {
+	if err := cfg.Validate(); err != nil {
+		return nil, fmt.Errorf("open postgres replicas: %w", err)
+	}
+	if len(cfg.ReplicaDSNs) == 0 {
+		return nil, nil
+	}
+
+	resolved := evalOptions(opts)
+	pools := make([]*sql.DB, 0, len(cfg.ReplicaDSNs))
+	for index, dsn := range cfg.ReplicaDSNs {
+		db, err := openDB(ctx, cfg, dsn, resolved)
+		if err != nil {
+			for _, opened := range pools {
+				_ = opened.Close()
+			}
+			return nil, fmt.Errorf("open postgres replica at index %d: %w", index, err)
+		}
+		pools = append(pools, db)
+	}
+	return pools, nil
+}
+
+func openDB(ctx context.Context, cfg Config, dsn string, opts options) (*sql.DB, error) {
+	if ctx.Err() != nil {
+		return nil, ctx.Err()
+	}
+
+	pgxCfg, err := pgx.ParseConfig(dsn)
+	if err != nil {
+		return nil, fmt.Errorf("parse dsn: %w", err)
+	}
+	pgxCfg.ConnectTimeout = cfg.OperationTimeout
+
+	registeredName := stdlib.RegisterConnConfig(pgxCfg)
+
+	db, err := otelsql.Open("pgx", registeredName, opts.otelsqlOpenOptions()...)
+	if err != nil {
+		stdlib.UnregisterConnConfig(registeredName)
+		return nil, fmt.Errorf("otelsql open: %w", err)
+	}
+	if db == nil {
+		stdlib.UnregisterConnConfig(registeredName)
+		return nil, errors.New("otelsql open returned nil db")
+	}
+
+	db.SetMaxOpenConns(cfg.MaxOpenConns)
+	db.SetMaxIdleConns(cfg.MaxIdleConns)
+	db.SetConnMaxLifetime(cfg.ConnMaxLifetime)
+
+	return db, nil
+}