github.com/blend/go-sdk@v1.20240719.1/tracing/dbtrace/tracer_test.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/driver"
    13  	"fmt"
    14  	"testing"
    15  
    16  	opentracing "github.com/opentracing/opentracing-go"
    17  	opentracingExt "github.com/opentracing/opentracing-go/ext"
    18  	"github.com/opentracing/opentracing-go/mocktracer"
    19  
    20  	"github.com/blend/go-sdk/assert"
    21  	"github.com/blend/go-sdk/tracing"
    22  )
    23  
    24  func TestPrepare(t *testing.T) {
    25  	assert := assert.New(t)
    26  	mockTracer := mocktracer.New()
    27  	dbTracer := Tracer(mockTracer)
    28  
    29  	dbCfg, err := defaultDB().Config.Reparse()
    30  	assert.Nil(err)
    31  
    32  	dbtf := dbTracer.Prepare(context.Background(), dbCfg, "select * from test_table limit 1")
    33  	span := dbtf.(dbTraceFinisher).span
    34  	mockSpan := span.(*mocktracer.MockSpan)
    35  	assert.Equal(tracing.OperationSQLPrepare, mockSpan.OperationName)
    36  
    37  	assert.Len(mockSpan.Tags(), 6)
    38  	assert.Equal(tracing.SpanTypeSQL, mockSpan.Tags()[tracing.TagKeySpanType])
    39  	assert.Equal(dbCfg.Database, mockSpan.Tags()[string(opentracingExt.DBInstance)])
    40  	assert.Equal(dbCfg.Username, mockSpan.Tags()[string(opentracingExt.DBUser)])
    41  	assert.Equal("select * from test_table limit 1", mockSpan.Tags()[string(opentracingExt.DBStatement)])
    42  	assert.True(mockSpan.FinishTime.IsZero())
    43  }
    44  
    45  func TestPrepareWithParentSpan(t *testing.T) {
    46  	assert := assert.New(t)
    47  	mockTracer := mocktracer.New()
    48  	dbTracer := Tracer(mockTracer)
    49  
    50  	dbCfg, err := defaultDB().Config.Reparse()
    51  	assert.Nil(err)
    52  
    53  	parentSpan := mockTracer.StartSpan("test_op")
    54  	ctx := opentracing.ContextWithSpan(context.Background(), parentSpan)
    55  
    56  	dbtf := dbTracer.Prepare(ctx, dbCfg, "select * from test_table limit 1")
    57  	span := dbtf.(dbTraceFinisher).span
    58  	mockSpan := span.(*mocktracer.MockSpan)
    59  	assert.Equal(tracing.OperationSQLPrepare, mockSpan.OperationName)
    60  
    61  	mockParentSpan := parentSpan.(*mocktracer.MockSpan)
    62  	assert.Equal(mockSpan.ParentID, mockParentSpan.SpanContext.SpanID)
    63  }
    64  
    65  func TestQuery(t *testing.T) {
    66  	assert := assert.New(t)
    67  	mockTracer := mocktracer.New()
    68  	dbTracer := Tracer(mockTracer)
    69  
    70  	dbCfg, err := defaultDB().Config.Reparse()
    71  	assert.Nil(err)
    72  
    73  	statement := "SELECT 1 FROM test_table WHERE id = $1"
    74  	invocation := defaultDB().Invoke()
    75  	invocation.Label = "test_table_exists"
    76  
    77  	dbtf := dbTracer.Query(context.Background(), dbCfg, invocation.Label, statement)
    78  	span := dbtf.(dbTraceFinisher).span
    79  	mockSpan := span.(*mocktracer.MockSpan)
    80  	assert.Equal(tracing.OperationSQLQuery, mockSpan.OperationName)
    81  
    82  	assert.Len(mockSpan.Tags(), 7)
    83  	assert.Equal("test_table_exists", mockSpan.Tags()[tracing.TagKeyResourceName])
    84  	assert.Equal(tracing.SpanTypeSQL, mockSpan.Tags()[tracing.TagKeySpanType])
    85  	assert.Equal(dbCfg.Database, mockSpan.Tags()[string(opentracingExt.DBInstance)])
    86  	assert.Equal(dbCfg.Username, mockSpan.Tags()[string(opentracingExt.DBUser)])
    87  	assert.Equal(statement, mockSpan.Tags()[string(opentracingExt.DBStatement)])
    88  	assert.True(mockSpan.FinishTime.IsZero())
    89  }
    90  
    91  func TestQueryWithParentSpan(t *testing.T) {
    92  	assert := assert.New(t)
    93  	mockTracer := mocktracer.New()
    94  	dbTracer := Tracer(mockTracer)
    95  
    96  	dbCfg, err := defaultDB().Config.Reparse()
    97  	assert.Nil(err)
    98  
    99  	parentSpan := mockTracer.StartSpan("test_op")
   100  	ctx := opentracing.ContextWithSpan(context.Background(), parentSpan)
   101  
   102  	statement := "SELECT 1 FROM test_table WHERE id = $1"
   103  	invocation := defaultDB().Invoke()
   104  	invocation.Label = "test_table_exists"
   105  
   106  	dbtf := dbTracer.Query(ctx, dbCfg, invocation.Label, statement)
   107  	span := dbtf.(dbTraceFinisher).span
   108  	mockSpan := span.(*mocktracer.MockSpan)
   109  	assert.Equal(tracing.OperationSQLQuery, mockSpan.OperationName)
   110  
   111  	mockParentSpan := parentSpan.(*mocktracer.MockSpan)
   112  	assert.Equal(mockSpan.ParentID, mockParentSpan.SpanContext.SpanID)
   113  }
   114  
   115  func TestFinishQuery(t *testing.T) {
   116  	assert := assert.New(t)
   117  	mockTracer := mocktracer.New()
   118  	dbTracer := Tracer(mockTracer)
   119  
   120  	dbCfg, err := defaultDB().Config.Reparse()
   121  	assert.Nil(err)
   122  
   123  	dbtf := dbTracer.Query(context.Background(), dbCfg, "ok", "select 'ok1'")
   124  	dbtf.FinishQuery(context.TODO(), nil, nil)
   125  
   126  	span := dbtf.(dbTraceFinisher).span
   127  	mockSpan := span.(*mocktracer.MockSpan)
   128  	assert.Nil(mockSpan.Tags()[tracing.TagKeyError])
   129  	assert.False(mockSpan.FinishTime.IsZero())
   130  }
   131  
   132  func TestFinishPrepare(t *testing.T) {
   133  	assert := assert.New(t)
   134  	mockTracer := mocktracer.New()
   135  	dbTracer := Tracer(mockTracer)
   136  
   137  	dbCfg, err := defaultDB().Config.Reparse()
   138  	assert.Nil(err)
   139  
   140  	dbtf := dbTracer.Prepare(context.Background(), dbCfg, "select 'ok1'")
   141  	dbtf.FinishPrepare(context.TODO(), nil)
   142  
   143  	span := dbtf.(dbTraceFinisher).span
   144  	mockSpan := span.(*mocktracer.MockSpan)
   145  	assert.Nil(mockSpan.Tags()[tracing.TagKeyError])
   146  	assert.False(mockSpan.FinishTime.IsZero())
   147  }
   148  
   149  func TestFinishQueryError(t *testing.T) {
   150  	assert := assert.New(t)
   151  	mockTracer := mocktracer.New()
   152  	dbTracer := Tracer(mockTracer)
   153  
   154  	dbCfg, err := defaultDB().Config.Reparse()
   155  	assert.Nil(err)
   156  
   157  	ctx := context.Background()
   158  	dbtf := dbTracer.Query(ctx, dbCfg, "ok", "select 'ok1'")
   159  	dbtf.FinishQuery(ctx, nil, fmt.Errorf("error"))
   160  
   161  	span := dbtf.(dbTraceFinisher).span
   162  	mockSpan := span.(*mocktracer.MockSpan)
   163  	assert.Equal("error", mockSpan.Tags()[tracing.TagKeyError])
   164  	assert.False(mockSpan.FinishTime.IsZero())
   165  }
   166  
   167  func TestFinishPrepareError(t *testing.T) {
   168  	assert := assert.New(t)
   169  	mockTracer := mocktracer.New()
   170  	dbTracer := Tracer(mockTracer)
   171  
   172  	dbCfg, err := defaultDB().Config.Reparse()
   173  	assert.Nil(err)
   174  
   175  	ctx := context.Background()
   176  	dbtf := dbTracer.Prepare(ctx, dbCfg, "select 'ok1'")
   177  	dbtf.FinishPrepare(ctx, fmt.Errorf("error"))
   178  
   179  	span := dbtf.(dbTraceFinisher).span
   180  	mockSpan := span.(*mocktracer.MockSpan)
   181  	assert.Equal("error", mockSpan.Tags()[tracing.TagKeyError])
   182  	assert.False(mockSpan.FinishTime.IsZero())
   183  }
   184  
   185  func TestFinishQueryErrorSkip(t *testing.T) {
   186  	assert := assert.New(t)
   187  	mockTracer := mocktracer.New()
   188  	dbTracer := Tracer(mockTracer)
   189  
   190  	dbCfg, err := defaultDB().Config.Reparse()
   191  	assert.Nil(err)
   192  
   193  	ctx := context.Background()
   194  	dbtf := dbTracer.Query(ctx, dbCfg, "ok", "select 'ok1'")
   195  	dbtf.FinishQuery(ctx, nil, driver.ErrSkip)
   196  
   197  	span := dbtf.(dbTraceFinisher).span
   198  	mockSpan := span.(*mocktracer.MockSpan)
   199  	assert.Nil(mockSpan.Tags()[tracing.TagKeyError])
   200  }
   201  
   202  func TestFinishQueryNil(t *testing.T) {
   203  	assert := assert.New(t)
   204  
   205  	dbtf := dbTraceFinisher{}
   206  	dbtf.FinishQuery(context.TODO(), nil, nil)
   207  	assert.Nil(dbtf.span)
   208  }