go: add pgx observability

Signed-off-by: Victor Login <batazor@evrone.com>
shortlink-org · Apr 8, 2023 · f057b75 · f057b75
1 parent 743cd8d
commit f057b75
Show file tree

Hide file tree

Showing 6 changed files with 41 additions and 5 deletions.
diff --git a/internal/pkg/db/postgres/postgres.go b/internal/pkg/db/postgres/postgres.go
@@ -12,11 +12,10 @@ import (
 	"github.com/jackc/pgx/v5/pgxpool"
 	"github.com/johejo/golang-migrate-extra/source/iofs"
 	_ "github.com/lib/pq" // need for init PostgreSQL interface
+	"github.com/shortlink-org/shortlink/internal/pkg/db/options"
 	"github.com/spf13/viper"
 	"github.com/uptrace/opentelemetry-go-extra/otelsql"
 	semconv "go.opentelemetry.io/otel/semconv/v1.12.0"
-
-	"github.com/shortlink-org/shortlink/internal/pkg/db/options"
 )
 
 var (
@@ -43,6 +42,9 @@ func (p *Store) Init(ctx context.Context) error {
 		return err
 	}
 
+	// Instrument the pgxpool config with OpenTelemetry.
+	cnf.Tracer = &p.tracer
+
 	// Create pool config
 	cnfPool, err := pgxpool.ParseConfig("")
 	if err != nil {

diff --git a/internal/pkg/db/postgres/tracer.go b/internal/pkg/db/postgres/tracer.go
@@ -0,0 +1,32 @@
+package postgres
+
+import (
+	"context"
+
+	"github.com/jackc/pgx/v5"
+	semconv "go.opentelemetry.io/otel/semconv/v1.12.0"
+	"go.opentelemetry.io/otel/trace"
+)
+
+type Tracer struct{}
+
+func (t *Tracer) TraceQueryStart(ctx context.Context, conn *pgx.Conn, data pgx.TraceQueryStartData) context.Context {
+	span := trace.SpanFromContext(ctx)
+
+	span.SetAttributes(
+		semconv.DBSystemPostgreSQL,
+		semconv.DBStatementKey.String(data.SQL),
+	)
+
+	return ctx
+}
+
+func (t *Tracer) TraceQueryEnd(ctx context.Context, conn *pgx.Conn, data pgx.TraceQueryEndData) {
+	span := trace.SpanFromContext(ctx)
+
+	if data.Err != nil {
+		span.RecordError(data.Err)
+	}
+
+	span.End()
+}
diff --git a/internal/pkg/db/postgres/type.go b/internal/pkg/db/postgres/type.go
@@ -14,4 +14,6 @@ type Config struct {
 type Store struct {
 	client *pgxpool.Pool
 	config Config
+
+	tracer Tracer
 }
diff --git a/internal/services/api/application/grpc_web/v1/server.go b/internal/services/api/application/grpc_web/v1/server.go
@@ -61,7 +61,7 @@ func (api *API) Run(
 		return err
 	}
 
-	api.http = http_server.New(ctx, mux, config)
+	api.http = http_server.New(ctx, mux, config, tracer)
 
 	// Start HTTP server (and proxy calls to gRPC server endpoint)
 	log.Info(fmt.Sprintf("API run on port %d", config.Port))

diff --git a/internal/services/api/application/http-chi/server.go b/internal/services/api/application/http-chi/server.go
@@ -101,7 +101,7 @@ func (api *API) Run(
 	r.Mount("/api/cqrs", cqrs_api.Routes(link_command, link_query))
 	r.Mount("/api/sitemap", sitemap_api.Routes(sitemap_rpc))
 
-	srv := http_server.New(ctx, r, config)
+	srv := http_server.New(ctx, r, config, tracer)
 
 	// start HTTP-server
 	log.Info(i18n.Sprintf("API run on port %d", config.Port))

diff --git a/internal/services/billing/infrastructure/api/http/http-chi/server.go b/internal/services/billing/infrastructure/api/http/http-chi/server.go
@@ -117,7 +117,7 @@ func (api *API) Run(
 		tariffRoutes.Routes(router)
 	}))
 
-	srv := http_server.New(ctx, r, config)
+	srv := http_server.New(ctx, r, config, tracer)
 
 	// start HTTP-server
 	log.Info(fmt.Sprintf("API run on port %d", config.Port))