feat: implement module for pgx/v5

module/apmpgxv5/batch_tracer.go

@@ -0,0 +1,94 @@
+// Licensed to Elasticsearch B.V. under one or more contributor
+// license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright
+// ownership. Elasticsearch B.V. licenses this file to you under
+// the Apache License, Version 2.0 (the "License"); you may
+// not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package apmpgxv5
+
+import (
+	"context"
+	"github.com/jackc/pgx/v5"
+	"go.elastic.co/apm/v2"
+	"time"
+)
+
+// BatchTracer traces SendBatch
+type BatchTracer struct{}
+
+var _ pgx.BatchTracer = (*BatchTracer)(nil)
+
+const (
+	batchSpanType = "db.postgresql.batch"
+)
+
+func (b BatchTracer) TraceBatchStart(ctx context.Context, conn *pgx.Conn, _ pgx.TraceBatchStartData) context.Context {
+	span, apmCtx := apm.StartSpanOptions(ctx, "BATCH", batchSpanType, apm.SpanOptions{
+		Start:    time.Now(),
+		ExitSpan: false,
+	})
+
+	if span.Dropped() {
+		span.End()
+		return nil // todo: should be discussed
+	}
+	span.Action = "batch"
+
+	return apmCtx
+}
+
+func (b BatchTracer) TraceBatchQuery(ctx context.Context, conn *pgx.Conn, data pgx.TraceBatchQueryData) {
+	span, apmCtx := apm.StartSpanOptions(ctx, data.SQL, querySpanType, apm.SpanOptions{
+		Start:    time.Now(),
+		ExitSpan: false,
+	})
+	defer span.End()
+
+	span.Action = "batch"
+	span.Context.SetDatabase(apm.DatabaseSpanContext{
+		Instance:  conn.Config().Database,
+		Statement: data.SQL,
+		Type:      "sql",
+		User:      conn.Config().User,
+	})
+	span.Context.SetDestinationAddress(conn.Config().Host, int(conn.Config().Port))
+
+	if apmErr := apm.CaptureError(apmCtx, data.Err); apmErr != nil {
+		apmErr.SetSpan(span)
+		apmErr.Send()
+		return
+	}
+}
+
+func (b BatchTracer) TraceBatchEnd(ctx context.Context, conn *pgx.Conn, data pgx.TraceBatchEndData) {
+	span := apm.SpanFromContext(ctx)
+	defer span.End()
+
+	if span.Dropped() {
+		span.End()
+		return
+	}
+
+	// check if span is ended or not
+	if span.SpanData != nil {
+		span.Context.SetDestinationAddress(conn.Config().Host, int(conn.Config().Port))
+		span.Action = "batch"
+	}
+
+	if apmErr := apm.CaptureError(ctx, data.Err); apmErr != nil {
+		apmErr.SetSpan(span)
+		apmErr.Send()
+		return
+	}
+}

module/apmpgxv5/batch_tracer_test.go

@@ -0,0 +1,101 @@
+package apmpgxv5_test
+
+import (
+	"context"
+	"fmt"
+	"github.com/gvencadze/apm-agent-go/module/apmpgxv5"
+	"github.com/jackc/pgx/v5"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"go.elastic.co/apm/v2/apmtest"
+	"go.elastic.co/apm/v2/model"
+	_ "go.elastic.co/apm/v2/model"
+	"os"
+	"testing"
+)
+
+func TestBatchTrace(t *testing.T) {
+	host := os.Getenv("PGHOST")
+	if host == "" {
+		t.Skipf("PGHOST not specified")
+	}
+
+	cfg, err := pgx.ParseConfig(fmt.Sprintf("postgres://postgres:hunter2@%s:5432/test_db", host))
+	require.NoError(t, err)
+
+	ctx := context.TODO()
+
+	apmpgxv5.Instrument(cfg)
+
+	conn, err := pgx.ConnectConfig(ctx, cfg)
+	require.NoError(t, err)
+
+	_, err = conn.Exec(ctx, "CREATE TABLE IF NOT EXISTS foo (bar INT)")
+	require.NoError(t, err)
+
+	testcases := []struct {
+		name       string
+		expectErr  bool
+		queryQueue []string
+	}{
+		{
+			name:      "BATCH spans, success",
+			expectErr: false,
+			queryQueue: []string{
+				"SELECT * FROM foo WHERE bar = 1",
+				"SELECT bar FROM foo WHERE bar = 1",
+			},
+		},
+		{
+			name:      "BATCH spans, error",
+			expectErr: true,
+			queryQueue: []string{
+				"SELECT * FROM foo WHERE bar = 1",
+				"SELECT bar FROM foo2",
+			},
+		},
+	}
+
+	for _, tt := range testcases {
+		t.Run(tt.name, func(t *testing.T) {
+			_, spans, errs := apmtest.WithTransaction(func(ctx context.Context) {
+				batch := &pgx.Batch{}
+
+				for _, query := range tt.queryQueue {
+					batch.Queue(query)
+				}
+
+				br := conn.SendBatch(ctx, batch)
+				defer func() {
+					_ = br.Close()
+				}()
+			})
+
+			if tt.expectErr {
+				require.Len(t, errs, 2)
+				assert.Equal(t, "failure", spans[0].Outcome)
+			} else {
+				for i, expectedStmt := range tt.queryQueue {
+					assert.Equal(t, "success", spans[i].Outcome)
+					assert.Equal(t, "db", spans[i].Type)
+					assert.Equal(t, "postgresql", spans[i].Subtype)
+					assert.Equal(t, "batch", spans[i].Action)
+					assert.Equal(t, expectedStmt, spans[i].Name)
+
+					assert.Equal(t, &model.SpanContext{
+						Destination: &model.DestinationSpanContext{
+							Address: cfg.Host,
+							Port:    int(cfg.Port),
+						},
+						Database: &model.DatabaseSpanContext{
+							Instance:  cfg.Database,
+							Statement: expectedStmt,
+							Type:      "sql",
+							User:      cfg.User,
+						},
+					}, spans[i].Context)
+				}
+			}
+		})
+	}
+}

module/apmpgxv5/connect_tracer.go

@@ -0,0 +1,103 @@
+// Licensed to Elasticsearch B.V. under one or more contributor
+// license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright
+// ownership. Elasticsearch B.V. licenses this file to you under
+// the Apache License, Version 2.0 (the "License"); you may
+// not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package apmpgxv5
+
+import (
+	"context"
+	"github.com/jackc/pgx/v5"
+	"go.elastic.co/apm/v2"
+	"time"
+)
+
+// ConnectTracer traces Connect and ConnectConfig
+type ConnectTracer struct{}
+
+var _ pgx.ConnectTracer = (*ConnectTracer)(nil)
+
+const (
+	connectSpanType = "db.postgresql.connect"
+)
+
+type contextKey string
+
+const (
+	dataContextKey contextKey = "data"
+)
+
+func (c ConnectTracer) TraceConnectStart(ctx context.Context, _ pgx.TraceConnectStartData) context.Context {
+	statement := "connect"
+
+	span, apmCtx := apm.StartSpanOptions(ctx, statement, connectSpanType, apm.SpanOptions{
+		ExitSpan: false,
+	})
+
+	newCtx := context.WithValue(apmCtx, dataContextKey, values{
+		start:     time.Now(),
+		statement: statement,
+	})
+
+	return apm.ContextWithSpan(newCtx, span)
+}
+
+func (c ConnectTracer) TraceConnectEnd(ctx context.Context, data pgx.TraceConnectEndData) {
+	span := apm.SpanFromContext(ctx)
+	defer span.End()
+
+	var (
+		statement string
+		db        string
+		user      string
+		host      string
+		port      int
+	)
+
+	if d, ok := ctx.Value(dataContextKey).(values); ok {
+		statement = d.statement
+	}
+
+	// TODO: refactor
+	if data.Conn != nil {
+		db = data.Conn.Config().Database
+		user = data.Conn.Config().User
+		host = data.Conn.Config().Host
+		port = int(data.Conn.Config().Port)
+	}
+
+	if span.Dropped() {
+		return
+	}
+
+	span.Context.SetDatabase(apm.DatabaseSpanContext{
+		Instance:  db,
+		Statement: statement,
+		Type:      "sql",
+		User:      user,
+	})
+	span.Context.SetDestinationAddress(host, port)
+	span.Context.SetServiceTarget(apm.ServiceTargetSpanContext{
+		Name: "postgresql",
+		Type: "db",
+	})
+
+	if apmErr := apm.CaptureError(ctx, data.Err); apmErr != nil {
+		apmErr.SetSpan(span)
+		apmErr.Send()
+	}
+
+	return
+}

module/apmpgxv5/connect_tracer_test.go

@@ -0,0 +1,62 @@
+package apmpgxv5_test
+
+import (
+	"context"
+	"fmt"
+	"github.com/gvencadze/apm-agent-go/module/apmpgxv5"
+	"github.com/jackc/pgx/v5"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"go.elastic.co/apm/v2/apmtest"
+	"os"
+	"testing"
+)
+
+// TODO: add tests for connect
+func Test_Connect(t *testing.T) {
+	host := os.Getenv("PGHOST")
+	if host == "" {
+		t.Skipf("PGHOST not specified")
+	}
+
+	testcases := []struct {
+		name      string
+		dsn       string
+		expectErr bool
+	}{
+		{
+			name:      "CONNECT span, success",
+			expectErr: false,
+			dsn:       "postgres://postgres:hunter2@%s:5432/test_db",
+		},
+		{
+			name:      "CONNECT span, failure",
+			expectErr: true,
+			dsn:       "postgres://postgres:hunter2@%s:5432/non_existing_db",
+		},
+	}
+
+	for _, tt := range testcases {
+		t.Run(tt.name, func(t *testing.T) {
+			cfg, err := pgx.ParseConfig(fmt.Sprintf(tt.dsn, host))
+			require.NoError(t, err)
+
+			apmpgxv5.Instrument(cfg)
+
+			_, spans, errs := apmtest.WithTransaction(func(ctx context.Context) {
+				_, _ = pgx.ConnectConfig(ctx, cfg)
+			})
+
+			assert.NotNil(t, spans[0].ID)
+
+			if tt.expectErr {
+				require.Len(t, errs, 1)
+				assert.Equal(t, "failure", spans[0].Outcome)
+				assert.Equal(t, "connect", spans[0].Name)
+			} else {
+				assert.Equal(t, "success", spans[0].Outcome)
+				assert.Equal(t, "connect", spans[0].Name)
+			}
+		})
+	}
+}