[hw-7] add metrics, tracing

loms/internal/infra/db/postgres/postgres.go

@@ -0,0 +1,60 @@
+package postgres
+
+import (
+	"context"
+
+	"github.com/jackc/pgx/v5/pgxpool"
+	"github.com/pkg/errors"
+)
+
+// From https://gitlab.ozon.dev/go/classroom-18/students/week-4-workshop/-/blob/master/internal/infra/postgres/postgres.go
+
+func NewPool(ctx context.Context, dsn string) (*pgxpool.Pool, error) {
+	config, err := pgxpool.ParseConfig(dsn)
+	if err != nil {
+		return nil, errors.Wrap(err, "pgxpool.ParseConfig")
+	}
+
+	config.ConnConfig.Tracer = promTracer{}
+
+	pool, err := pgxpool.NewWithConfig(ctx, config)
+	if err != nil {
+		return nil, errors.Wrap(err, "pgxpool.NewWithConfig")
+	}
+
+	return pool, nil
+}
+
+func NewPools(ctx context.Context, DSNs ...string) ([]*pgxpool.Pool, error) {
+	pools := make([]*pgxpool.Pool, len(DSNs))
+
+	for i, dsn := range DSNs {
+		cfg, err := pgxpool.ParseConfig(dsn)
+		if err != nil {
+			closeOpened(pools[:i])
+
+			return nil, errors.Wrap(err, "pgxpool.ParseConfig")
+		}
+
+		cfg.ConnConfig.Tracer = promTracer{}
+
+		pool, err := pgxpool.NewWithConfig(ctx, cfg)
+		if err != nil {
+			closeOpened(pools[:i])
+
+			return nil, errors.Wrap(err, "pgxpool.NewWithConfig")
+		}
+
+		pools[i] = pool
+	}
+
+	return pools, nil
+}
+
+func closeOpened(pools []*pgxpool.Pool) {
+	for _, p := range pools {
+		if p != nil {
+			p.Close()
+		}
+	}
+}

loms/internal/infra/db/postgres/tracer.go

@@ -0,0 +1,38 @@
+package postgres
+
+import (
+	"context"
+	"strings"
+	"time"
+
+	"route256/loms/internal/infra/grpc/metrics"
+
+	"github.com/jackc/pgx/v5"
+)
+
+type startKey struct{}
+
+type promTracer struct{}
+
+func (promTracer) TraceQueryStart(ctx context.Context, _ *pgx.Conn, _ pgx.TraceQueryStartData) context.Context {
+	return context.WithValue(ctx, startKey{}, time.Now())
+}
+
+func (promTracer) TraceQueryEnd(ctx context.Context, _ *pgx.Conn, data pgx.TraceQueryEndData) {
+	start, _ := ctx.Value(startKey{}).(time.Time)
+
+	sql := strings.TrimSpace(strings.ToLower(data.CommandTag.String()))
+	parts := strings.SplitN(sql, " ", 2)
+	category := "unknown"
+	if len(parts) > 0 {
+		category = parts[0]
+	}
+
+	errLabel := "none"
+	if data.Err != nil {
+		errLabel = "error"
+	}
+
+	metrics.IncDBQueryCount(category, errLabel)
+	metrics.StoreDBQueryDuration(category, errLabel, time.Since(start))
+}

loms/internal/infra/db/postgres/tx.go

@@ -0,0 +1,87 @@
+package postgres
+
+// From https://gitlab.ozon.dev/go/classroom-18/students/week-4-workshop/-/blob/master/internal/infra/postgres/tx.go
+
+import (
+	"context"
+
+	"github.com/jackc/pgx/v5"
+	"github.com/jackc/pgx/v5/pgxpool"
+	"github.com/opentracing/opentracing-go"
+)
+
+// Tx транзакция.
+type Tx pgx.Tx
+
+type txKey struct{}
+
+func ctxWithTx(ctx context.Context, tx pgx.Tx) context.Context {
+	return context.WithValue(ctx, txKey{}, tx)
+}
+
+func TxFromCtx(ctx context.Context) (pgx.Tx, bool) {
+	tx, ok := ctx.Value(txKey{}).(pgx.Tx)
+
+	return tx, ok
+}
+
+type TxManager struct {
+	write *pgxpool.Pool
+	read  *pgxpool.Pool
+}
+
+func NewTxManager(write, read *pgxpool.Pool) *TxManager {
+	return &TxManager{
+		write: write,
+		read:  read,
+	}
+}
+
+// WithTransaction выполняет fn в транзакции с дефолтным уровнем изоляции.
+func (m *TxManager) WriteWithTransaction(ctx context.Context, fn func(ctx context.Context) error) (err error) {
+	return m.withTx(ctx, m.write, pgx.TxOptions{}, fn)
+}
+
+func (m *TxManager) ReadWithTransaction(ctx context.Context, fn func(ctx context.Context) error) (err error) {
+	return m.withTx(ctx, m.read, pgx.TxOptions{}, fn)
+}
+
+// WithTransaction выполняет fn в транзакции с уровнем изоляции RepeatableRead.
+func (m *TxManager) WriteWithRepeatableRead(ctx context.Context, fn func(ctx context.Context) error) (err error) {
+	return m.withTx(ctx, m.write, pgx.TxOptions{IsoLevel: pgx.RepeatableRead}, fn)
+}
+
+func (m *TxManager) ReadWithRepeatableRead(ctx context.Context, fn func(ctx context.Context) error) (err error) {
+	return m.withTx(ctx, m.read, pgx.TxOptions{IsoLevel: pgx.RepeatableRead}, fn)
+}
+
+// WithTx выполняет fn в транзакции.
+func (m *TxManager) withTx(ctx context.Context, pool *pgxpool.Pool, options pgx.TxOptions, fn func(ctx context.Context) error) (err error) {
+	var span opentracing.Span
+	span, ctx = opentracing.StartSpanFromContext(ctx, "Transaction")
+	defer span.Finish()
+
+	tx, err := pool.BeginTx(ctx, options)
+	if err != nil {
+		return
+	}
+	ctx = ctxWithTx(ctx, tx)
+
+	defer func() {
+		if p := recover(); p != nil {
+			// a panic occurred, rollback and repanic
+			_ = tx.Rollback(ctx)
+			panic(p)
+		} else if err != nil {
+			// something went wrong, rollback
+			_ = tx.Rollback(ctx)
+		} else {
+			// all good, commit
+			err = tx.Commit(ctx)
+		}
+	}()
+
+	err = fn(ctx)
+
+	return
+}