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	"github.com/jinzhu/gorm"
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