Passed
Push — main ( 0795ea...57a23e )
by Yume
02:19 queued 01:06
created

infrastructures.ConnectDB   B

Complexity

Conditions 5

Size

Total Lines 48
Code Lines 32

Duplication

Lines 0
Ratio 0 %

Importance

Changes 0
Metric Value
cc 5
eloc 32
nop 0
dl 0
loc 48
rs 8.6453
c 0
b 0
f 0
1
package infrastructures
2
3
import (
4
	"log"
5
	"os"
6
	"time"
7
8
	"github.com/memnix/memnix-rest/config"
9
	"github.com/memnix/memnix-rest/pkg/env"
10
	"github.com/pkg/errors"
11
	"gorm.io/driver/mysql"
12
	"gorm.io/gorm"
13
	"gorm.io/gorm/logger"
14
	"gorm.io/plugin/opentelemetry/tracing"
15
	"gorm.io/plugin/prometheus"
16
)
17
18
// DBConn is the database connection object
19
var DBConn *gorm.DB
20
21
// GetDBConn returns the database connection object
22
func GetDBConn() *gorm.DB {
23
	return DBConn
24
}
25
26
// getDSN returns the database connection string
27
// see: utils/env.go
28
func getDSN(env *env.Env) string {
29
	var dsn string
30
31
	// Get database configuration from environment variables
32
	if config.IsDevelopment() {
33
		dsn = env.GetEnv("DEBUG_DB_DSN")
34
	} else {
35
		dsn = env.GetEnv("DB_DSN")
36
	}
37
	return dsn
38
}
39
40
// ConnectDB creates a connection to database
41
//
42
// see: utils/config.go and utils/env.go for more details
43
func ConnectDB() error {
44
	newLogger := logger.New(
45
		log.New(os.Stdout, "\r\n", log.LstdFlags), // io writer
46
		logger.Config{
47
			LogLevel:                  logger.Silent, // Log level
48
			IgnoreRecordNotFoundError: true,          // Ignore ErrRecordNotFound error for logger
49
			Colorful:                  true,          // Disable color
50
		},
51
	)
52
53
	dsn := getDSN(config.EnvHelper)
54
55
	// Open connection
56
	conn, err := gorm.Open(mysql.Open(dsn), &gorm.Config{
57
		Logger:                                   newLogger, // Logger
58
		SkipDefaultTransaction:                   true,      // Skip default transaction
59
		DisableForeignKeyConstraintWhenMigrating: true,      // Disable foreign key constraint when migrating (planetscale recommends this)
60
	})
61
	if err != nil {
62
		return errors.Wrap(err, "failed to connect to database")
63
	}
64
65
	sqlDB, err := conn.DB() // Get sql.DB object from gorm.DB
66
	if err != nil {
67
		return errors.Wrap(err, "failed to get sql.DB object")
68
	}
69
	sqlDB.SetMaxIdleConns(config.SQLMaxIdleConns) // Set max idle connections
70
	sqlDB.SetMaxOpenConns(config.SQLMaxOpenConns) // Set max open connections
71
	sqlDB.SetConnMaxLifetime(time.Second)         // Set max connection lifetime
72
73
	if err = conn.Use(prometheus.New(prometheus.Config{
74
		DBName:          "db1",                                // `DBName` as metrics label
75
		RefreshInterval: config.GormPrometheusRefreshInterval, // refresh metrics interval (default 15 seconds)
76
		StartServer:     false,                                // start http server to expose metrics
77
		MetricsCollector: []prometheus.MetricsCollector{
78
			&prometheus.MySQL{VariableNames: []string{"Threads_running"}},
79
		},
80
	})); err != nil {
81
		return err
82
	}
83
84
	if err = conn.Use(tracing.NewPlugin()); err != nil {
85
		return err
86
	}
87
88
	DBConn = conn
89
90
	return nil
91
}
92
93
// DisconnectDB closes the database connection
94
func DisconnectDB() error {
95
	sqlDB, err := GetDBConn().DB() // Get sql.DB object from gorm.DB
96
	if err != nil {
97
		return errors.Wrap(err, "failed to get sql.DB object")
98
	}
99
	if err = sqlDB.Close(); err != nil {
100
		return errors.Wrap(err, "failed to close database connection")
101
	}
102
103
	return nil
104
}
105