storage/postgres: pgx client tracing (#5438)

* fix testcontainers docker client using the global tracer provider

* storage/postgres: pgx client tracing

* skip postgres test on macos

pkg/storage/postgres/backend.go

@@ -9,6 +9,7 @@ import (
 	"time"
 
 	"github.com/cenkalti/backoff/v4"
+	"github.com/exaring/otelpgx"
 	"github.com/jackc/pgx/v5/pgxpool"
 	"google.golang.org/protobuf/types/known/fieldmaskpb"
 	"google.golang.org/protobuf/types/known/timestamppb"
@@ -353,6 +354,11 @@ func (backend *Backend) init(ctx context.Context) (serverVersion uint64, pool *p
 		return serverVersion, nil, err
 	}
 
+	if backend.cfg.tracerProvider != nil {
+		config.ConnConfig.Tracer = otelpgx.NewTracer(
+			otelpgx.WithTracerProvider(backend.cfg.tracerProvider))
+	}
+
 	pool, err = pgxpool.NewWithConfig(context.Background(), config)
 	if err != nil {
 		return serverVersion, nil, err

pkg/storage/postgres/option.go

@@ -2,6 +2,8 @@ package postgres
 
 import (
 	"time"
+
+	oteltrace "go.opentelemetry.io/otel/trace"
 )
 
 const (
@@ -10,8 +12,9 @@ const (
 )
 
 type config struct {
-	expiry      time.Duration
-	registryTTL time.Duration
+	expiry         time.Duration
+	registryTTL    time.Duration
+	tracerProvider oteltrace.TracerProvider
 }
 
 // Option customizes a Backend.
@@ -31,6 +34,12 @@ func WithRegistryTTL(ttl time.Duration) Option {
 	}
 }
 
+func WithTracerProvider(tracerProvider oteltrace.TracerProvider) Option {
+	return func(cfg *config) {
+		cfg.tracerProvider = tracerProvider
+	}
+}
+
 func getConfig(options ...Option) *config {
 	cfg := new(config)
 	WithExpiry(defaultExpiry)(cfg)

pkg/storage/postgres/tracing_test.go

@@ -0,0 +1,70 @@
+package postgres_test
+
+import (
+	"context"
+	"io"
+	"net/http"
+	"os"
+	"runtime"
+	"testing"
+
+	"github.com/pomerium/pomerium/config"
+	"github.com/pomerium/pomerium/internal/testenv"
+	"github.com/pomerium/pomerium/internal/testenv/scenarios"
+	"github.com/pomerium/pomerium/internal/testenv/snippets"
+	"github.com/pomerium/pomerium/internal/testenv/upstreams"
+	"github.com/pomerium/pomerium/internal/testutil"
+	"github.com/pomerium/pomerium/internal/testutil/tracetest"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestQueryTracing(t *testing.T) {
+	if os.Getenv("GITHUB_ACTION") != "" && runtime.GOOS == "darwin" {
+		t.Skip("Github action can not run docker on MacOS")
+	}
+
+	testutil.WithTestPostgres(t, func(dsn string) {
+		receiver := scenarios.NewOTLPTraceReceiver()
+		env := testenv.New(t, testenv.WithTraceDebugFlags(testenv.StandardTraceDebugFlags), testenv.WithTraceClient(receiver.NewGRPCClient()))
+		env.Add(receiver)
+
+		env.Add(testenv.ModifierFunc(func(_ context.Context, cfg *config.Config) {
+			cfg.Options.DataBrokerStorageType = config.StoragePostgresName
+			cfg.Options.DataBrokerStorageConnectionString = dsn
+		}))
+		up := upstreams.HTTP(nil, upstreams.WithDisplayName("Upstream"))
+		up.Handle("/foo", func(w http.ResponseWriter, _ *http.Request) {
+			w.Write([]byte("OK"))
+		})
+		env.Add(scenarios.NewIDP([]*scenarios.User{{Email: "user@example.com"}}))
+
+		route := up.Route().
+			From(env.SubdomainURL("postgres-test")).
+			PPL(`{"allow":{"and":["email":{"is":"user@example.com"}]}}`)
+		env.AddUpstream(up)
+
+		env.Start()
+		snippets.WaitStartupComplete(env)
+
+		resp, err := up.Get(route, upstreams.AuthenticateAs("user@example.com"), upstreams.Path("/foo"))
+		assert.NoError(t, err)
+		io.ReadAll(resp.Body)
+		resp.Body.Close()
+
+		env.Stop()
+
+		results := tracetest.NewTraceResults(receiver.FlushResourceSpans())
+		traces, exists := results.GetTraces().ByParticipant["Data Broker"]
+		require.True(t, exists)
+		require.Len(t, traces, 1)
+		var found bool
+		for _, span := range traces[0].Spans {
+			if span.Scope.GetName() == "github.com/exaring/otelpgx" {
+				found = true
+				break
+			}
+		}
+		assert.True(t, found, "no spans with otelpgx scope found")
+	})
+}