1// Unless explicitly stated otherwise all files in this repository are licensed 2// under the Apache License Version 2.0. 3// This product includes software developed at Datadog (https://www.datadoghq.com/). 4// Copyright 2016 Datadog, Inc. 5 6package gorm 7 8import ( 9 "math" 10 11 "gopkg.in/DataDog/dd-trace-go.v1/internal" 12 13 "gopkg.in/jinzhu/gorm.v1" 14) 15 16type config struct { 17 serviceName string 18 analyticsRate float64 19 dsn string 20 tagFns map[string]func(scope *gorm.Scope) interface{} 21 errCheck func(err error) bool 22} 23 24// Option represents an option that can be passed to Register, Open or OpenDB. 25type Option func(*config) 26 27func defaults(cfg *config) { 28 cfg.serviceName = "gorm.db" 29 // cfg.analyticsRate = globalconfig.AnalyticsRate() 30 if internal.BoolEnv("DD_TRACE_GORM_ANALYTICS_ENABLED", false) { 31 cfg.analyticsRate = 1.0 32 } else { 33 cfg.analyticsRate = math.NaN() 34 } 35 cfg.errCheck = func(error) bool { return true } 36} 37 38// WithServiceName sets the given service name when registering a driver, 39// or opening a database connection. 40func WithServiceName(name string) Option { 41 return func(cfg *config) { 42 cfg.serviceName = name 43 } 44} 45 46// WithAnalytics enables Trace Analytics for all started spans. 47func WithAnalytics(on bool) Option { 48 return func(cfg *config) { 49 if on { 50 cfg.analyticsRate = 1.0 51 } else { 52 cfg.analyticsRate = math.NaN() 53 } 54 } 55} 56 57// WithAnalyticsRate sets the sampling rate for Trace Analytics events 58// correlated to started spans. 59func WithAnalyticsRate(rate float64) Option { 60 return func(cfg *config) { 61 if rate >= 0.0 && rate <= 1.0 { 62 cfg.analyticsRate = rate 63 } else { 64 cfg.analyticsRate = math.NaN() 65 } 66 } 67} 68 69// WithCustomTag will cause the given tagFn to be evaluated after executing 70// a query and attach the result to the span tagged by the key. 71func WithCustomTag(tag string, tagFn func(scope *gorm.Scope) interface{}) Option { 72 return func(cfg *config) { 73 if cfg.tagFns == nil { 74 cfg.tagFns = make(map[string]func(scope *gorm.Scope) interface{}) 75 } 76 cfg.tagFns[tag] = tagFn 77 } 78} 79 80// WithErrorCheck specifies a function fn which determines whether the passed 81// error should be marked as an error. The fn is called whenever a gorm operation 82// finishes with an error 83func WithErrorCheck(fn func(err error) bool) Option { 84 return func(cfg *config) { 85 cfg.errCheck = fn 86 } 87} 88