From 11a3cc4dd501ecbfe1d39af674fafe54cf64ed1e Mon Sep 17 00:00:00 2001
From: Patrik Helia <patrik.helia@kiwi.com>
Date: Wed, 10 Jun 2020 01:09:51 +0200
Subject: [PATCH] Add support for pg_go

Signed-off-by: Patrik Helia <patashelia@gmail.com>
---
 contrib/go-pg/go-pg/pg_go.go      | 52 +++++++++++++++++++++++++++++++
 contrib/go-pg/go-pg/pg_go_test.go | 43 +++++++++++++++++++++++++
 2 files changed, 95 insertions(+)
 create mode 100644 contrib/go-pg/go-pg/pg_go.go
 create mode 100644 contrib/go-pg/go-pg/pg_go_test.go

diff --git a/contrib/go-pg/go-pg/pg_go.go b/contrib/go-pg/go-pg/pg_go.go
new file mode 100644
index 0000000000..e9010a76f0
--- /dev/null
+++ b/contrib/go-pg/go-pg/pg_go.go
@@ -0,0 +1,52 @@
+package pggo
+
+import (
+	"context"
+	"github.com/go-pg/pg/v10"
+	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace"
+	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/ext"
+	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
+	"time"
+)
+
+const (
+	gopgStartSpan  = "dd-trace-go:span"
+)
+
+func Wrap(db *pg.DB) *pg.DB {
+	db.AddQueryHook(&Hook{})
+	return db
+}
+
+type Hook struct {}
+
+func (h *Hook) BeforeQuery(ctx context.Context,qe *pg.QueryEvent) (context.Context, error){
+	if qe.Stash == nil{
+		qe.Stash = make(map[interface{}]interface{})
+	}
+	qe.Stash[gopgStartSpan] = time.Now()
+	return ctx, qe.Err
+}
+
+
+func (h *Hook) AfterQuery(ctx context.Context,qe *pg.QueryEvent) error{
+	startSpan, ok := qe.Stash[gopgStartSpan]
+	if !ok{
+		return nil // TODO Return some error
+	}
+	
+	spanStart := startSpan.(time.Time)
+	unformatedSql, _ := qe.UnformattedQuery()
+
+	opts := []ddtrace.StartSpanOption{
+		tracer.StartTime(spanStart),
+		//tracer.ServiceName(tracer), // Should be took from context, or consider if can be set through WithMethod
+		tracer.SpanType(ext.SpanTypeSQL),
+		tracer.ResourceName(string(unformatedSql)),
+	}
+
+	span, _ := tracer.StartSpanFromContext(ctx, "SELECT", opts...) // TODO Write sql scanner for operation name
+	span.Finish(tracer.WithError(qe.Err))
+
+	return qe.Err
+}
diff --git a/contrib/go-pg/go-pg/pg_go_test.go b/contrib/go-pg/go-pg/pg_go_test.go
new file mode 100644
index 0000000000..4d840d9ec2
--- /dev/null
+++ b/contrib/go-pg/go-pg/pg_go_test.go
@@ -0,0 +1,43 @@
+package pggo
+
+import (
+	"context"
+	"github.com/go-pg/pg/v10"
+	"github.com/stretchr/testify/assert"
+	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/ext"
+	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/mocktracer"
+	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
+	"testing"
+)
+
+
+func GetPostgresConnection() *pg.DB{
+	db := pg.Connect(&pg.Options{
+		User:     "postgres",
+		Password: "",
+		Database: "postgres",
+	})
+	return db
+}
+
+func TestSelect(t *testing.T){
+	assert := assert.New(t)
+	mt := mocktracer.Start()
+	defer mt.Stop()
+
+	conn := GetPostgresConnection()
+	Wrap(conn)
+
+	parentSpan, ctx := tracer.StartSpanFromContext(context.Background(), "http.request",
+		tracer.ServiceName("fake-http-server"),
+		tracer.SpanType(ext.SpanTypeWeb),
+	)
+
+	var n int
+	_, _ = conn.WithContext(ctx).QueryOne(pg.Scan(&n), "SELECT 1")
+	parentSpan.Finish()
+	spans := mt.FinishedSpans()
+
+	assert.True(len(spans) >= 2)
+	// TODO Complete tests
+}