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

     1  /*
     2  
     3  Copyright (c) 2024 - 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  	opentracingExt "github.com/opentracing/opentracing-go/ext"
    18  
    19  	"github.com/blend/go-sdk/db"
    20  	"github.com/blend/go-sdk/tracing"
    21  )
    22  
    23  var (
    24  	_ db.Tracer = (*dbTracer)(nil)
    25  )
    26  
    27  // Tracer returns a db tracer.
    28  func Tracer(tracer opentracing.Tracer) db.Tracer {
    29  	return &dbTracer{tracer: tracer}
    30  }
    31  
    32  type dbTracer struct {
    33  	tracer opentracing.Tracer
    34  }
    35  
    36  func (dbt dbTracer) Prepare(ctx context.Context, cfg db.Config, statement string) db.TraceFinisher {
    37  	startOptions := []opentracing.StartSpanOption{
    38  		opentracingExt.SpanKindRPCClient,
    39  		opentracing.Tag{Key: tracing.TagKeySpanType, Value: tracing.SpanTypeSQL},
    40  		// Ensure lib is using expected DB span tags for this DB span
    41  		// https://docs.datadoghq.com/tracing/trace_collection/tracing_naming_convention/#database
    42  		opentracing.Tag{Key: string(opentracingExt.DBInstance), Value: cfg.DatabaseOrDefault()},
    43  		opentracing.Tag{Key: string(opentracingExt.DBUser), Value: cfg.Username},
    44  		opentracing.Tag{Key: string(opentracingExt.DBStatement), Value: statement},
    45  		tracing.TagMeasured(),
    46  		opentracing.StartTime(time.Now().UTC()),
    47  	}
    48  	span, _ := tracing.StartSpanFromContext(ctx, dbt.tracer, tracing.OperationSQLPrepare, startOptions...)
    49  	return dbTraceFinisher{span: span}
    50  }
    51  
    52  func (dbt dbTracer) Query(ctx context.Context, cfg db.Config, label, statement string) db.TraceFinisher {
    53  	startOptions := []opentracing.StartSpanOption{
    54  		opentracingExt.SpanKindRPCClient,
    55  		opentracing.Tag{Key: tracing.TagKeyResourceName, Value: label},
    56  		opentracing.Tag{Key: tracing.TagKeySpanType, Value: tracing.SpanTypeSQL},
    57  		// Ensure lib is using expected DB span tags for this DB span
    58  		// https://docs.datadoghq.com/tracing/trace_collection/tracing_naming_convention/#database
    59  		opentracing.Tag{Key: string(opentracingExt.DBInstance), Value: cfg.DatabaseOrDefault()},
    60  		opentracing.Tag{Key: string(opentracingExt.DBUser), Value: cfg.Username},
    61  		opentracing.Tag{Key: string(opentracingExt.DBStatement), Value: statement},
    62  		tracing.TagMeasured(),
    63  		opentracing.StartTime(time.Now().UTC()),
    64  	}
    65  	span, _ := tracing.StartSpanFromContext(ctx, dbt.tracer, tracing.OperationSQLQuery, startOptions...)
    66  	return dbTraceFinisher{span: span}
    67  }
    68  
    69  type dbTraceFinisher struct {
    70  	span opentracing.Span
    71  }
    72  
    73  func (dbtf dbTraceFinisher) FinishPrepare(ctx context.Context, err error) {
    74  	if dbtf.span == nil {
    75  		return
    76  	}
    77  	if err == driver.ErrSkip {
    78  		return
    79  	}
    80  	tracing.SpanError(dbtf.span, err)
    81  	dbtf.span.Finish()
    82  }
    83  
    84  func (dbtf dbTraceFinisher) FinishQuery(ctx context.Context, res sql.Result, err error) {
    85  	if dbtf.span == nil {
    86  		return
    87  	}
    88  	if err == driver.ErrSkip {
    89  		return
    90  	}
    91  	if res != nil {
    92  		affected, _ := res.RowsAffected()
    93  		dbtf.span.SetTag(tracing.TagKeyDBRowsAffected, affected)
    94  	}
    95  	tracing.SpanError(dbtf.span, err)
    96  	dbtf.span.Finish()
    97  }