feat(module): add module for tracing SQL queries created by jackc/pgx (…

…#1301)

* feat: implementation of pgx query tracer

* feat: add logger support

* feat: add copy protocol support

* test: add tests for query and copy trace

* test: add table tests to test errors in different cases + add Log function test

* feat: implement batch trace

* docs: add godoc

* feat: add stacktrace to span

* test: add stacktrace length check

* refac: make tracer private

* refac: rename package from apmpgx to apmpgx_test to test usage like an external API

* test: add t.Parallel() to table tests

* chore(tracer): remove set stacktrace because it's unnecessary and bump pgx library version up to v4.17 and write comment for error with min required version

* docs(instrumenting): add description of library that serves at "https://www.elastic.co/guide/en/apm/agent/go/master/builtin-modules.html"

* docs(instrumenting): add example for apmgpx usage without pool

* refac: replace NewTracer constructor with Instrument function that accepts pgx.ConnConfig. Move duplicate code to separate func called startSpan to make ...Trace functions small and simple. Add type casting checks to avoid instrument runtime panics.

* test: replace NewTracer call with creating config and then override it with Instrument func for further tests

* docs(ci): add license header, run update-licenses,update-modules, fmt and add apmpgx entry to Dockerfile-testing

* fix(ci): add vanity import path

* refac(tracer): add bool return argument check from startSpan() in CopyTrace func

* refac(query): add statement parameter to startSpan func for full statement in span.Context.Database.Statement field

* docs(apmpgx): update docs due changes in package API

* test(e2e): add end-to-end tests for query and copy statements spans

* feat(copy): add statement to copy from trace

* test(copy): add statement to DatabaseSpanContext into E2E_CopyTrace test

* fix(test): revert postgres connection in E2E_CopyTrace

* fix(test): replace interface type cast from []string to pgx.Identifier according to stored value in data map

* fix(tracer): add type switch for column names that stored in data map

* docs(changelog): add apmpgx module description entry into unreleased section

Co-authored-by: Stepan Rabotkin <[email protected]>

CHANGELOG.asciidoc

@@ -26,6 +26,7 @@ https://github.com/elastic/apm-agent-go/compare/v2.1.0...main[View commits]
 - Global labels are now parsed when the tracer is constructed, instead of parsing only once on package initialization {pull}1290[#(1290)]
 - Rename span_frames_min_duration to span_stack_trace_min_duration {pull}1285[#(1285)]
 - Ignore *principal* headers by default {pull}1332[#(1332)]
+- Add *apmpgx* module for postgres tracing with jackc/pgx driver enhanced support e.g. Copy and Batch statements {pull}1301[#(1301)]
 
 [[release-notes-2.x]]
 === Go Agent version 2.x

docs/instrumenting.asciidoc

@@ -32,6 +32,7 @@ that framework's API. The request context can be used for reporting <<custom-ins
 * <<builtin-modules-apmmongo>>
 * <<builtin-modules-apmawssdkgo>>
 * <<builtin-modules-apmazure>>
+* <<builtin-modules-apmpgx>>
 
 [[builtin-modules-apmhttp]]
 ==== module/apmhttp
@@ -969,3 +970,55 @@ func main() {
   ...
 }
 ----
+
+[[builtin-modules-apmpgx]]
+== module/apmpgx
+Package apmpgx provides a means of instrumenting the
+https://github.com/jackc/pgx[Pgx] for v4.17+,
+so that SQL queries are reported as spans within the current transaction.
+Also this lib have extended support of pgx, such as COPY FROM queries and BATCH which have their own span types:
+`db.postgresql.copy` and `db.postgresql.batch` accordingly.
+
+To report `pgx` queries, create `pgx` connection, and then provide config to `apmpgx.Instrument()`. If logger is presented in config,
+then traces will be written to log. (It's safe to use without logger)
+
+Spans will be created for queries as long as they have context associated, and the
+context includes a transaction.
+
+Example with pool:
+[source,go]
+----
+import (
+    "github.com/jackc/pgx/v4/pgxpool"
+
+    "go.elastic.com/apm/module/apmpgx/v2"
+)
+
+func main() {
+    c, err := pgxpool.ParseConfig("dsn_string")
+    ...
+    pool, err := pgxpool.ParseConfig("dsn")
+    ...
+    // set custom logger before instrumenting
+    apmpgx.Instrument(pool.ConnConfig)
+    ...
+}
+----
+
+Example without pool:
+[source,go]
+----
+import (
+    "github.com/jackc/pgx/v4"
+
+    "go.elastic.com/apm/module/apmpgx/v2"
+)
+
+func main() {
+    c, err := pgx.Connect(context.Background, "dsn_string")
+    ...
+    // set custom logger before instrumenting
+    apmpgx.Instrument(c.Config())
+    ...
+}
+----

module/apmpgx/doc.go

@@ -0,0 +1,19 @@
+// 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 apmpgx provides helpers for tracing github.com/jackc/pgx/v4. Minimal required version is v4.17
+package apmpgx // import "go.elastic.co/apm/module/apmpgx/v2"

module/apmpgx/e2e_test.go

@@ -0,0 +1,212 @@
+// 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 apmpgx_test
+
+import (
+	"context"
+	"fmt"
+	"os"
+	_ "os"
+	"testing"
+
+	"github.com/jackc/pgx/v4"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"go.elastic.co/apm/module/apmpgx/v2"
+	"go.elastic.co/apm/v2/apmtest"
+	"go.elastic.co/apm/v2/model"
+)
+
+func Test_E2E_QueryTrace(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()
+
+	apmpgx.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
+		query     string
+	}{
+		{
+			name:      "QUERY span, success",
+			expectErr: false,
+			query:     "SELECT * FROM foo",
+		},
+		{
+			name:      "QUERY span, error",
+			expectErr: true,
+			query:     "SELECT * FROM foo2",
+		},
+	}
+
+	for _, tt := range testcases {
+		t.Run(tt.name, func(t *testing.T) {
+			_, spans, errs := apmtest.WithTransaction(func(ctx context.Context) {
+				rows, _ := conn.Query(ctx, tt.query)
+				defer rows.Close()
+			})
+
+			assert.NotNil(t, spans[0].ID)
+
+			if tt.expectErr {
+				require.Len(t, errs, 1)
+				assert.Equal(t, "failure", spans[0].Outcome)
+			} else {
+				assert.Equal(t, "success", spans[0].Outcome)
+
+				assert.Equal(t, "SELECT FROM foo", spans[0].Name)
+				assert.Equal(t, "postgresql", spans[0].Subtype)
+				assert.Equal(t, "success", spans[0].Outcome)
+
+				assert.Equal(t, &model.SpanContext{
+					Destination: &model.DestinationSpanContext{
+						Address: cfg.Host,
+						Port:    int(cfg.Port),
+						Service: &model.DestinationServiceSpanContext{
+							Type:     "db",
+							Name:     "postgresql",
+							Resource: "postgresql",
+						},
+					},
+					Service: &model.ServiceSpanContext{
+						Target: &model.ServiceTargetSpanContext{
+							Type: "postgresql",
+							Name: cfg.Database,
+						},
+					},
+					Database: &model.DatabaseSpanContext{
+						Instance:  cfg.Database,
+						Statement: "SELECT * FROM foo",
+						Type:      "sql",
+						User:      "postgres",
+					},
+				}, spans[0].Context)
+			}
+		})
+	}
+}
+
+func Test_E2E_CopyTrace(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()
+
+	apmpgx.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
+		tableName   pgx.Identifier
+		columnNames pgx.Identifier
+		rows        [][]interface{}
+	}{
+		{
+			name:        "COPY span, success",
+			expectErr:   false,
+			tableName:   pgx.Identifier{"foo"},
+			columnNames: pgx.Identifier{"bar"},
+			rows: [][]interface{}{
+				{int32(36)},
+				{int32(29)},
+			},
+		},
+		{
+			name:        "COPY span, fail",
+			expectErr:   true,
+			tableName:   pgx.Identifier{"foo"},
+			columnNames: pgx.Identifier{"bar"},
+			rows: [][]interface{}{
+				{"error"},
+			},
+		},
+	}
+
+	for _, tt := range testcases {
+		t.Run(tt.name, func(t *testing.T) {
+			_, spans, errs := apmtest.WithTransaction(func(ctx context.Context) {
+				_, _ = conn.CopyFrom(ctx,
+					tt.tableName,
+					tt.columnNames,
+					pgx.CopyFromRows(tt.rows))
+			})
+
+			assert.NotNil(t, spans[0].ID)
+
+			if tt.expectErr {
+				require.Len(t, errs, 1)
+				assert.Equal(t, "failure", spans[0].Outcome)
+			} else {
+				assert.Equal(t, "success", spans[0].Outcome)
+
+				assert.Equal(t, "COPY TO foo", spans[0].Name)
+				assert.Equal(t, "postgresql", spans[0].Subtype)
+
+				assert.Equal(t, &model.SpanContext{
+					Destination: &model.DestinationSpanContext{
+						Address: cfg.Host,
+						Port:    int(cfg.Port),
+						Service: &model.DestinationServiceSpanContext{
+							Type:     "db",
+							Name:     "postgresql",
+							Resource: "postgresql",
+						},
+					},
+					Service: &model.ServiceSpanContext{
+						Target: &model.ServiceTargetSpanContext{
+							Type: "postgresql",
+							Name: cfg.Database,
+						},
+					},
+					Database: &model.DatabaseSpanContext{
+						Instance:  cfg.Database,
+						Statement: "COPY TO foo(bar)",
+						Type:      "sql",
+						User:      "postgres",
+					},
+				}, spans[0].Context)
+			}
+		})
+	}
+}

module/apmpgx/go.mod

@@ -0,0 +1,14 @@
+module go.elastic.co/apm/module/apmpgx/v2
+
+go 1.15
+
+require (
+	github.com/jackc/pgx/v4 v4.17.0
+	github.com/stretchr/testify v1.8.0
+	go.elastic.co/apm/module/apmsql/v2 v2.1.0
+	go.elastic.co/apm/v2 v2.1.0
+)
+
+replace go.elastic.co/apm/v2 => ../..
+
+replace go.elastic.co/apm/module/apmsql/v2 => ../apmsql