github.com/blend/go-sdk@v1.20220411.3/tracing/dbtrace/tracer.go (about)

     1  /*
     2  
     3  Copyright (c) 2022 - Present. Blend Labs, Inc. All rights reserved
     4  Use of this source code is governed by a MIT license that can be found in the LICENSE file.
     5  
     6  */
     7  
     8  package dbtrace
     9  
    10  import (
    11  	"context"
    12  	"database/sql"
    13  	"database/sql/driver"
    14  	"time"
    15  
    16  	opentracing "github.com/opentracing/opentracing-go"
    17  
    18  	"github.com/blend/go-sdk/db"
    19  	"github.com/blend/go-sdk/tracing"
    20  )
    21  
    22  var (
    23  	_ db.Tracer = (*dbTracer)(nil)
    24  )
    25  
    26  // Tracer returns a db tracer.
    27  func Tracer(tracer opentracing.Tracer) db.Tracer {
    28  	return &dbTracer{tracer: tracer}
    29  }
    30  
    31  type dbTracer struct {
    32  	tracer opentracing.Tracer
    33  }
    34  
    35  func (dbt dbTracer) Prepare(ctx context.Context, cfg db.Config, statement string) db.TraceFinisher {
    36  	startOptions := []opentracing.StartSpanOption{
    37  		opentracing.Tag{Key: tracing.TagKeySpanType, Value: tracing.SpanTypeSQL},
    38  		opentracing.Tag{Key: tracing.TagKeyDBName, Value: cfg.DatabaseOrDefault()},
    39  		opentracing.Tag{Key: tracing.TagKeyDBUser, Value: cfg.Username},
    40  		opentracing.Tag{Key: TagKeyQuery, Value: statement},
    41  		opentracing.Tag{Key: TagKeySQLCommand, Value: statement},
    42  		tracing.TagMeasured(),
    43  		opentracing.StartTime(time.Now().UTC()),
    44  	}
    45  	span, _ := tracing.StartSpanFromContext(ctx, dbt.tracer, tracing.OperationSQLPrepare, startOptions...)
    46  	return dbTraceFinisher{span: span}
    47  }
    48  
    49  func (dbt dbTracer) Query(ctx context.Context, cfg db.Config, label, statement string) db.TraceFinisher {
    50  	startOptions := []opentracing.StartSpanOption{
    51  		opentracing.Tag{Key: tracing.TagKeyResourceName, Value: label},
    52  		opentracing.Tag{Key: tracing.TagKeySpanType, Value: tracing.SpanTypeSQL},
    53  		opentracing.Tag{Key: tracing.TagKeyDBName, Value: cfg.DatabaseOrDefault()},
    54  		opentracing.Tag{Key: tracing.TagKeyDBUser, Value: cfg.Username},
    55  		opentracing.Tag{Key: TagKeyQuery, Value: statement},
    56  		opentracing.Tag{Key: TagKeySQLCommand, Value: statement},
    57  		tracing.TagMeasured(),
    58  		opentracing.StartTime(time.Now().UTC()),
    59  	}
    60  	span, _ := tracing.StartSpanFromContext(ctx, dbt.tracer, tracing.OperationSQLQuery, startOptions...)
    61  	return dbTraceFinisher{span: span}
    62  }
    63  
    64  type dbTraceFinisher struct {
    65  	span opentracing.Span
    66  }
    67  
    68  func (dbtf dbTraceFinisher) FinishPrepare(ctx context.Context, err error) {
    69  	if dbtf.span == nil {
    70  		return
    71  	}
    72  	if err == driver.ErrSkip {
    73  		return
    74  	}
    75  	tracing.SpanError(dbtf.span, err)
    76  	dbtf.span.Finish()
    77  }
    78  
    79  func (dbtf dbTraceFinisher) FinishQuery(ctx context.Context, res sql.Result, err error) {
    80  	if dbtf.span == nil {
    81  		return
    82  	}
    83  	if err == driver.ErrSkip {
    84  		return
    85  	}
    86  	if res != nil {
    87  		affected, _ := res.RowsAffected()
    88  		dbtf.span.SetTag(tracing.TagKeyDBRowsAffected, affected)
    89  	}
    90  	tracing.SpanError(dbtf.span, err)
    91  	dbtf.span.Finish()
    92  }