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