Skip to content

Commit

Permalink
contrib/jackc/pgx: implement Connect/Batch/CopyFrom/Prepare tracers
Browse files Browse the repository at this point in the history
  • Loading branch information
mrkagelui committed Dec 27, 2023
1 parent c9d2cb5 commit 6bb79bb
Show file tree
Hide file tree
Showing 5 changed files with 215 additions and 2 deletions.
66 changes: 66 additions & 0 deletions contrib/jackc/pgx/batch.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,66 @@
package pgx

import (
"context"
"math"
"time"

"github.com/jackc/pgx/v5"

"gopkg.in/DataDog/dd-trace-go.v1/ddtrace"
"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/ext"
ddtracer "gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
)

// TraceBatchStart marks the start of a batch, implementing pgx.BatchTracer
func (t *tracer) TraceBatchStart(ctx context.Context, _ *pgx.Conn, _ pgx.TraceBatchStartData) context.Context {
opts := []ddtrace.StartSpanOption{
ddtracer.ServiceName(t.serviceName),
ddtracer.SpanType(ext.SpanTypeSQL),
ddtracer.StartTime(time.Now()),
ddtracer.Tag("sql.query_type", "Batch"),
ddtracer.Tag(ext.ResourceName, "pgx.batch"),
}
for key, tag := range t.tags {
opts = append(opts, ddtracer.Tag(key, tag))
}
if !math.IsNaN(t.analyticsRate) {
opts = append(opts, ddtracer.Tag(ext.EventSampleRate, t.analyticsRate))
}
_, ctx = ddtracer.StartSpanFromContext(ctx, "pgx.batch", opts...)

return ctx
}

// TraceBatchQuery traces the query of a batch, implementing pgx.BatchTracer
func (t *tracer) TraceBatchQuery(ctx context.Context, c *pgx.Conn, data pgx.TraceBatchQueryData) {
opts := []ddtrace.StartSpanOption{
ddtracer.ServiceName(t.serviceName),
ddtracer.SpanType(ext.SpanTypeSQL),
ddtracer.StartTime(time.Now()),
ddtracer.Tag(ext.ResourceName, data.SQL),
}
if t.traceArgs {
opts = append(opts, ddtracer.Tag("sql.args", data.Args))
}
for key, tag := range t.tags {
opts = append(opts, ddtracer.Tag(key, tag))
}
if !math.IsNaN(t.analyticsRate) {
opts = append(opts, ddtracer.Tag(ext.EventSampleRate, t.analyticsRate))
}
ddtracer.StartSpanFromContext(ctx, "pgx.batch_query", opts...)
}

// TraceBatchEnd marks the end of a batch, implementing pgx.BatchTracer
func (t *tracer) TraceBatchEnd(ctx context.Context, _ *pgx.Conn, data pgx.TraceBatchEndData) {
span, exists := ddtracer.SpanFromContext(ctx)
if !exists {
return
}

if data.Err != nil {
span.SetTag(ext.Error, data.Err)
}
span.Finish()
}
45 changes: 45 additions & 0 deletions contrib/jackc/pgx/connect.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,45 @@
package pgx

import (
"context"
"math"
"time"

"github.com/jackc/pgx/v5"

"gopkg.in/DataDog/dd-trace-go.v1/ddtrace"
"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/ext"
ddtracer "gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
)

// TraceConnectStart marks the start of a pgx connect operation, implementing pgx.ConnectTracer
func (t *tracer) TraceConnectStart(ctx context.Context, _ pgx.TraceConnectStartData) context.Context {
opts := []ddtrace.StartSpanOption{
ddtracer.ServiceName(t.serviceName),
ddtracer.SpanType(ext.SpanTypeSQL),
ddtracer.StartTime(time.Now()),
ddtracer.Tag("sql.query_type", "Connect"),
}
for key, tag := range t.tags {
opts = append(opts, ddtracer.Tag(key, tag))
}
if !math.IsNaN(t.analyticsRate) {
opts = append(opts, ddtracer.Tag(ext.EventSampleRate, t.analyticsRate))
}
_, ctx = ddtracer.StartSpanFromContext(ctx, "pgx.connect", opts...)

return ctx
}

// TraceConnectEnd marks the end of a pgx connect operation, implementing pgx.ConnectTracer
func (t *tracer) TraceConnectEnd(ctx context.Context, data pgx.TraceConnectEndData) {
span, exists := ddtracer.SpanFromContext(ctx)
if !exists {
return
}

if data.Err != nil {
span.SetTag(ext.Error, data.Err)
}
span.Finish()
}
53 changes: 53 additions & 0 deletions contrib/jackc/pgx/copyfrom.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,53 @@
package pgx

import (
"context"
"math"
"time"

"github.com/jackc/pgx/v5"

"gopkg.in/DataDog/dd-trace-go.v1/ddtrace"
"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/ext"
ddtracer "gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
)

// TraceCopyFromStart marks the start of a CopyFrom query, implementing pgx.CopyFromTracer
func (t *tracer) TraceCopyFromStart(ctx context.Context, _ *pgx.Conn, data pgx.TraceCopyFromStartData) context.Context {
opts := []ddtrace.StartSpanOption{
ddtracer.ServiceName(t.serviceName),
ddtracer.SpanType(ext.SpanTypeSQL),
ddtracer.StartTime(time.Now()),
ddtracer.Tag("sql.query_type", "Query"),
ddtracer.Tag(ext.ResourceName, "pgx.copyfrom"),
ddtracer.Tag("pgx.table_name", data.TableName),
ddtracer.Tag("pgx.columns", data.ColumnNames),
}
for key, tag := range t.tags {
opts = append(opts, ddtracer.Tag(key, tag))
}
if !math.IsNaN(t.analyticsRate) {
opts = append(opts, ddtracer.Tag(ext.EventSampleRate, t.analyticsRate))
}

_, ctx = ddtracer.StartSpanFromContext(ctx, "pgx.copyfrom", opts...)

return ctx
}

// TraceCopyFromEnd marks the end of a CopyFrom query, implementing pgx.CopyFromTracer
func (t *tracer) TraceCopyFromEnd(ctx context.Context, _ *pgx.Conn, data pgx.TraceCopyFromEndData) {
span, exists := ddtracer.SpanFromContext(ctx)
if !exists {
return
}

if t.traceStatus {
span.SetTag("pgx.status", data.CommandTag.String())
}

if data.Err != nil {
span.SetTag(ext.Error, data.Err)
}
span.Finish()
}
49 changes: 49 additions & 0 deletions contrib/jackc/pgx/prepare.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,49 @@
package pgx

import (
"context"
"math"
"time"

"github.com/jackc/pgx/v5"

"gopkg.in/DataDog/dd-trace-go.v1/ddtrace"
"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/ext"
ddtracer "gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
)

// TracePrepareStart marks the start of a pgx prepare operation, implementing pgx.PrepareTracer
func (t *tracer) TracePrepareStart(ctx context.Context, _ *pgx.Conn, data pgx.TracePrepareStartData) context.Context {
opts := []ddtrace.StartSpanOption{
ddtracer.ServiceName(t.serviceName),
ddtracer.SpanType(ext.SpanTypeSQL),
ddtracer.StartTime(time.Now()),
ddtracer.Tag("sql.query_type", "Prepare"),
ddtracer.Tag("pgx.prepared_statement_name", data.Name),
ddtracer.Tag(ext.ResourceName, data.SQL),
}
for key, tag := range t.tags {
opts = append(opts, ddtracer.Tag(key, tag))
}
if !math.IsNaN(t.analyticsRate) {
opts = append(opts, ddtracer.Tag(ext.EventSampleRate, t.analyticsRate))
}
_, ctx = ddtracer.StartSpanFromContext(ctx, "pgx.prepare", opts...)

return ctx
}

// TracePrepareEnd marks the end of a pgx prepare operation, implementing pgx.PrepareTracer
func (t *tracer) TracePrepareEnd(ctx context.Context, _ *pgx.Conn, data pgx.TracePrepareEndData) {
span, exists := ddtracer.SpanFromContext(ctx)
if !exists {
return
}

span.SetTag("pgx.already_prepared", data.AlreadyPrepared)

if data.Err != nil {
span.SetTag(ext.Error, data.Err)
}
span.Finish()
}
4 changes: 2 additions & 2 deletions contrib/jackc/pgx/query.go
Original file line number Diff line number Diff line change
Expand Up @@ -5,11 +5,11 @@ import (
"math"
"time"

"github.com/jackc/pgx/v5"

"gopkg.in/DataDog/dd-trace-go.v1/ddtrace"
"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/ext"
ddtracer "gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"

"github.com/jackc/pgx/v5"
)

// TraceQueryStart marks the start of a query, implementing pgx.QueryTracer
Expand Down

0 comments on commit 6bb79bb

Please sign in to comment.