Fix broken tests

contrib/database/sql/conn.go

@@ -77,7 +77,7 @@ func (tc *tracedConn) PrepareContext(ctx context.Context, query string) (stmt dr
 				span.Finish(tracer.WithError(err))
 			}()
 		}
-		stmt, err := connPrepareCtx.PrepareContext(ctx, sqlCommentCarrier.CommentedQuery(query))
+		stmt, err := connPrepareCtx.PrepareContext(tracer.ContextWithSpan(ctx, span), sqlCommentCarrier.CommentedQuery(query))
 		if err != nil {
 			return nil, err
 		}
@@ -139,14 +139,15 @@ func (tc *tracedConn) ExecContext(ctx context.Context, query string, args []driv
 func (tc *tracedConn) Ping(ctx context.Context) (err error) {
 	start := time.Now()
 	if pinger, ok := tc.Conn.(driver.Pinger); ok {
+		span := tc.tryStartTrace(ctx, queryTypePing, "", start, &tracer.SQLCommentCarrier{}, err)
+		if span != nil {
+			go func() {
+				span.Finish(tracer.WithError(err))
+			}()
+		}
 		err = pinger.Ping(ctx)
 	}
-	span := tc.tryStartTrace(ctx, queryTypePing, "", start, &tracer.SQLCommentCarrier{}, err)
-	if span != nil {
-		go func() {
-			span.Finish(tracer.WithError(err))
-		}()
-	}
+
 	return err
 }
 
@@ -224,6 +225,7 @@ func WithSpanTags(ctx context.Context, tags map[string]string) context.Context {
 
 // tryStartTrace will create a span using the given arguments, but will act as a no-op when err is driver.ErrSkip.
 func (tp *traceParams) tryStartTrace(ctx context.Context, qtype queryType, query string, startTime time.Time, sqlCommentCarrier *tracer.SQLCommentCarrier, err error) (span tracer.Span) {
+	fmt.Printf("Executing query type %s\n", qtype)
 	if err == driver.ErrSkip {
 		// Not a user error: driver is telling sql package that an
 		// optional interface method is not implemented. There is

contrib/database/sql/sql.go

@@ -131,7 +131,7 @@ type tracedConnector struct {
 	cfg        *config
 }
 
-func (t *tracedConnector) Connect(ctx context.Context) (driver.Conn, error) {
+func (t *tracedConnector) Connect(ctx context.Context) (c driver.Conn, err error) {
 	tp := &traceParams{
 		driverName: t.driverName,
 		cfg:        t.cfg,
@@ -142,13 +142,13 @@ func (t *tracedConnector) Connect(ctx context.Context) (driver.Conn, error) {
 		tp.meta, _ = internal.ParseDSN(t.driverName, t.cfg.dsn)
 	}
 	start := time.Now()
-	conn, err := t.connector.Connect(ctx)
 	span := tp.tryStartTrace(ctx, queryTypeConnect, "", start, &tracer.SQLCommentCarrier{}, err)
 	if span != nil {
 		go func() {
 			span.Finish(tracer.WithError(err))
 		}()
 	}
+	conn, err := t.connector.Connect(tracer.ContextWithSpan(ctx, span))
 	if err != nil {
 		return nil, err
 	}

contrib/database/sql/stmt.go

@@ -27,13 +27,14 @@ type tracedStmt struct {
 // Close sends a span before closing a statement
 func (s *tracedStmt) Close() (err error) {
 	start := time.Now()
-	err = s.Stmt.Close()
 	span := s.tryStartTrace(s.ctx, queryTypeClose, "", start, &tracer.SQLCommentCarrier{}, err)
 	if span != nil {
 		go func() {
 			span.Finish(tracer.WithError(err))
 		}()
 	}
+	err = s.Stmt.Close()
+
 	return err
 }
 
@@ -75,13 +76,13 @@ func (s *tracedStmt) ExecContext(ctx context.Context, args []driver.NamedValue)
 func (s *tracedStmt) QueryContext(ctx context.Context, args []driver.NamedValue) (rows driver.Rows, err error) {
 	start := time.Now()
 	if stmtQueryContext, ok := s.Stmt.(driver.StmtQueryContext); ok {
+		rows, err := stmtQueryContext.QueryContext(ctx, args)
 		span := s.tryStartTrace(ctx, queryTypeQuery, s.query, start, &tracer.SQLCommentCarrier{}, err)
 		if span != nil {
 			go func() {
 				span.Finish(tracer.WithError(err))
 			}()
 		}
-		rows, err := stmtQueryContext.QueryContext(ctx, args)
 
 		return rows, err
 	}

contrib/database/sql/tx.go

@@ -39,11 +39,12 @@ func (t *tracedTx) Commit() (err error) {
 func (t *tracedTx) Rollback() (err error) {
 	start := time.Now()
 	span := t.tryStartTrace(t.ctx, queryTypeRollback, "", start, &tracer.SQLCommentCarrier{}, err)
-	err = t.Tx.Rollback()
 	if span != nil {
 		go func() {
 			span.Finish(tracer.WithError(err))
 		}()
 	}
+	err = t.Tx.Rollback()
+
 	return err
 }