Skip to content

Commit deb3be3

Browse files
abhishek9686abhishek9686
authored
Merge pull request #3488 from gravitl/patch/db-conn-pool
Patch: Use single db handle and use connection pool
2 parents b86aea9 + d9f910c commit deb3be3

File tree

10 files changed

+99
-57
lines changed

10 files changed

+99
-57
lines changed

‎controllers/network_test.go‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,8 @@ package controller
22

33
import (
44
"context"
5+
"github.com/gravitl/netmaker/db"
6+
"github.com/gravitl/netmaker/schema"
57
"os"
68
"testing"
79

@@ -23,6 +25,9 @@ type NetworkValidationTestCase struct {
2325
var netHost models.Host
2426

2527
func TestMain(m *testing.M) {
28+
db.InitializeDB(schema.ListModels()...)
29+
defer db.CloseDB()
30+
2631
database.InitializeDatabase()
2732
defer database.CloseDB()
2833
logic.CreateSuperAdmin(&models.User{

‎controllers/server.go‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -233,6 +233,9 @@ func getConfig(w http.ResponseWriter, r *http.Request) {
233233
if servercfg.IsPro {
234234
scfg.IsPro = "yes"
235235
}
236+
237+
scfg.ClientID = logic.Mask()
238+
scfg.ClientSecret = logic.Mask()
236239
json.NewEncoder(w).Encode(scfg)
237240
// w.WriteHeader(http.StatusOK)
238241
}

‎database/postgres.go‎

Lines changed: 11 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,12 @@
11
package database
22

33
import (
4+
"context"
45
"database/sql"
56
"errors"
6-
"fmt"
7+
"github.com/gravitl/netmaker/db"
8+
"time"
79

8-
"github.com/gravitl/netmaker/servercfg"
910
_ "github.com/lib/pq"
1011
)
1112

@@ -25,24 +26,19 @@ var PG_FUNCTIONS = map[string]interface{}{
2526
isConnected: pgIsConnected,
2627
}
2728

28-
func getPGConnString() string {
29-
pgconf := servercfg.GetSQLConf()
30-
pgConn := fmt.Sprintf("host=%s port=%d user=%s "+
31-
"password=%s dbname=%s sslmode=%s connect_timeout=5",
32-
pgconf.Host, pgconf.Port, pgconf.Username, pgconf.Password, pgconf.DB, pgconf.SSLMode)
33-
return pgConn
34-
}
35-
3629
func initPGDB() error {
37-
connString := getPGConnString()
30+
gormDB := db.FromContext(db.WithContext(context.TODO()))
31+
3832
var dbOpenErr error
39-
PGDB, dbOpenErr = sql.Open("postgres", connString)
33+
PGDB, dbOpenErr = gormDB.DB()
4034
if dbOpenErr != nil {
4135
return dbOpenErr
4236
}
43-
dbOpenErr = PGDB.Ping()
4437

45-
return dbOpenErr
38+
PGDB.SetMaxOpenConns(5)
39+
PGDB.SetConnMaxLifetime(time.Hour)
40+
41+
return PGDB.Ping()
4642
}
4743

4844
func pgCreateTable(tableName string) error {
@@ -134,7 +130,7 @@ func pgFetchRecords(tableName string) (map[string]string, error) {
134130
}
135131

136132
func pgCloseDB() {
137-
PGDB.Close()
133+
//PGDB.Close()
138134
}
139135

140136
func pgIsConnected() bool {

‎database/sqlite.go‎

Lines changed: 11 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,15 @@
11
package database
22

33
import (
4+
"context"
45
"database/sql"
56
"errors"
6-
"os"
7-
"path/filepath"
7+
"github.com/gravitl/netmaker/db"
8+
"time"
89

910
_ "github.com/mattn/go-sqlite3" // need to blank import this package
1011
)
1112

12-
// == sqlite ==
13-
const dbFilename = "netmaker.db"
14-
1513
// SqliteDB is the db object for sqlite database connections
1614
var SqliteDB *sql.DB
1715

@@ -29,21 +27,17 @@ var SQLITE_FUNCTIONS = map[string]interface{}{
2927
}
3028

3129
func initSqliteDB() error {
32-
// == create db file if not present ==
33-
if _, err := os.Stat("data"); os.IsNotExist(err) {
34-
os.Mkdir("data", 0700)
35-
}
36-
dbFilePath := filepath.Join("data", dbFilename)
37-
if _, err := os.Stat(dbFilePath); os.IsNotExist(err) {
38-
os.Create(dbFilePath)
39-
}
40-
// == "connect" the database ==
30+
gormDB := db.FromContext(db.WithContext(context.TODO()))
31+
4132
var dbOpenErr error
42-
SqliteDB, dbOpenErr = sql.Open("sqlite3", dbFilePath)
33+
SqliteDB, dbOpenErr = gormDB.DB()
4334
if dbOpenErr != nil {
4435
return dbOpenErr
4536
}
46-
SqliteDB.SetMaxOpenConns(1)
37+
38+
SqliteDB.SetMaxOpenConns(5)
39+
SqliteDB.SetConnMaxLifetime(time.Hour)
40+
4741
return nil
4842
}
4943

@@ -134,7 +128,7 @@ func sqliteFetchRecords(tableName string) (map[string]string, error) {
134128
}
135129

136130
func sqliteCloseDB() {
137-
SqliteDB.Close()
131+
//SqliteDB.Close()
138132
}
139133

140134
func sqliteConnected() bool {

‎db/db.go‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -110,3 +110,24 @@ func BeginTx(ctx context.Context) context.Context {
110110

111111
return context.WithValue(ctx, dbCtxKey, dbInCtx.Begin())
112112
}
113+
114+
// CloseDB close a connection to the database
115+
// (if one exists). It panics if any error
116+
// occurs.
117+
func CloseDB() {
118+
if db == nil {
119+
return
120+
}
121+
122+
sqlDB, err := db.DB()
123+
if err != nil {
124+
panic(err)
125+
}
126+
127+
err = sqlDB.Close()
128+
if err != nil {
129+
panic(err)
130+
}
131+
132+
db = nil
133+
}

‎db/postgres.go‎

Lines changed: 4 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ package db
22

33
import (
44
"fmt"
5+
"github.com/gravitl/netmaker/servercfg"
56
"os"
67
"strconv"
78

@@ -18,7 +19,7 @@ type postgresConnector struct{}
1819
// postgresConnector.connect connects and
1920
// initializes a connection to postgres.
2021
func (pg *postgresConnector) connect() (*gorm.DB, error) {
21-
pgConf := GetSQLConf()
22+
pgConf := servercfg.GetSQLConf()
2223
dsn := fmt.Sprintf(
2324
"host=%s port=%d user=%s password=%s dbname=%s sslmode=%s connect_timeout=5",
2425
pgConf.Host,
@@ -29,27 +30,11 @@ func (pg *postgresConnector) connect() (*gorm.DB, error) {
2930
pgConf.SSLMode,
3031
)
3132

32-
db, err := gorm.Open(postgres.Open(dsn), &gorm.Config{
33+
return gorm.Open(postgres.Open(dsn), &gorm.Config{
3334
Logger: logger.Default.LogMode(logger.Silent),
3435
})
35-
if err != nil {
36-
return nil, err
37-
}
38-
39-
// ensure netmaker_v1 schema exists.
40-
err = db.Exec("CREATE SCHEMA IF NOT EXISTS netmaker_v1").Error
41-
if err != nil {
42-
return nil, err
43-
}
44-
45-
// set the netmaker_v1 schema as the default schema.
46-
err = db.Exec("SET search_path TO netmaker_v1").Error
47-
if err != nil {
48-
return nil, err
49-
}
50-
51-
return db, nil
5236
}
37+
5338
func GetSQLConf() config.SQLConfig {
5439
var cfg config.SQLConfig
5540
cfg.Host = GetSQLHost()

‎functions/helpers_test.go‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,8 @@ package functions
33
import (
44
"context"
55
"encoding/json"
6+
"github.com/gravitl/netmaker/db"
7+
"github.com/gravitl/netmaker/schema"
68
"os"
79
"testing"
810

@@ -23,6 +25,9 @@ var (
2325
)
2426

2527
func TestMain(m *testing.M) {
28+
db.InitializeDB(schema.ListModels()...)
29+
defer db.CloseDB()
30+
2631
database.InitializeDatabase()
2732
defer database.CloseDB()
2833
logic.CreateSuperAdmin(&models.User{

‎logic/enrollmentkey_test.go‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,8 @@
11
package logic
22

33
import (
4+
"github.com/gravitl/netmaker/db"
5+
"github.com/gravitl/netmaker/schema"
46
"testing"
57
"time"
68

@@ -11,6 +13,9 @@ import (
1113
)
1214

1315
func TestCreateEnrollmentKey(t *testing.T) {
16+
db.InitializeDB(schema.ListModels()...)
17+
defer db.CloseDB()
18+
1419
database.InitializeDatabase()
1520
defer database.CloseDB()
1621
t.Run("Can_Not_Create_Key", func(t *testing.T) {
@@ -60,6 +65,9 @@ func TestCreateEnrollmentKey(t *testing.T) {
6065
}
6166

6267
func TestDelete_EnrollmentKey(t *testing.T) {
68+
db.InitializeDB(schema.ListModels()...)
69+
defer db.CloseDB()
70+
6371
database.InitializeDatabase()
6472
defer database.CloseDB()
6573
newKey, _ := CreateEnrollmentKey(0, time.Time{}, []string{"mynet", "skynet"}, nil, nil, true, uuid.Nil, false, false)
@@ -81,6 +89,9 @@ func TestDelete_EnrollmentKey(t *testing.T) {
8189
}
8290

8391
func TestDecrement_EnrollmentKey(t *testing.T) {
92+
db.InitializeDB(schema.ListModels()...)
93+
defer db.CloseDB()
94+
8495
database.InitializeDatabase()
8596
defer database.CloseDB()
8697
newKey, _ := CreateEnrollmentKey(1, time.Time{}, nil, nil, nil, false, uuid.Nil, false, false)
@@ -105,6 +116,9 @@ func TestDecrement_EnrollmentKey(t *testing.T) {
105116
}
106117

107118
func TestUsability_EnrollmentKey(t *testing.T) {
119+
db.InitializeDB(schema.ListModels()...)
120+
defer db.CloseDB()
121+
108122
database.InitializeDatabase()
109123
defer database.CloseDB()
110124
key1, _ := CreateEnrollmentKey(1, time.Time{}, nil, nil, nil, false, uuid.Nil, false, false)
@@ -143,6 +157,9 @@ func removeAllEnrollments() {
143157
//Test that cheks if it can't tokenize
144158

145159
func TestTokenize_EnrollmentKeys(t *testing.T) {
160+
db.InitializeDB(schema.ListModels()...)
161+
defer db.CloseDB()
162+
146163
database.InitializeDatabase()
147164
defer database.CloseDB()
148165
newKey, _ := CreateEnrollmentKey(0, time.Time{}, []string{"mynet", "skynet"}, nil, nil, true, uuid.Nil, false, false)
@@ -176,6 +193,9 @@ func TestTokenize_EnrollmentKeys(t *testing.T) {
176193
}
177194

178195
func TestDeTokenize_EnrollmentKeys(t *testing.T) {
196+
db.InitializeDB(schema.ListModels()...)
197+
defer db.CloseDB()
198+
179199
database.InitializeDatabase()
180200
defer database.CloseDB()
181201
newKey, _ := CreateEnrollmentKey(0, time.Time{}, []string{"mynet", "skynet"}, nil, nil, true, uuid.Nil, false, false)

‎logic/host_test.go‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,8 @@ package logic
33
import (
44
"context"
55
"fmt"
6+
"github.com/gravitl/netmaker/db"
7+
"github.com/gravitl/netmaker/schema"
68
"net"
79
"os"
810
"testing"
@@ -14,6 +16,9 @@ import (
1416
)
1517

1618
func TestMain(m *testing.M) {
19+
db.InitializeDB(schema.ListModels()...)
20+
defer db.CloseDB()
21+
1722
database.InitializeDatabase()
1823
defer database.CloseDB()
1924
peerUpdate := make(chan *models.Node)
@@ -41,6 +46,9 @@ func TestCheckPorts(t *testing.T) {
4146
}
4247
//not sure why this initialization is required but without it
4348
// RemoveHost returns database is closed
49+
db.InitializeDB(schema.ListModels()...)
50+
defer db.CloseDB()
51+
4452
database.InitializeDatabase()
4553
RemoveHost(&h, true)
4654
CreateHost(&h)

‎main.go‎

Lines changed: 11 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,8 @@ import (
77
"encoding/json"
88
"flag"
99
"fmt"
10+
"github.com/gravitl/netmaker/db"
11+
"github.com/gravitl/netmaker/schema"
1012
"os"
1113
"os/signal"
1214
"path/filepath"
@@ -18,15 +20,13 @@ import (
1820
"github.com/gravitl/netmaker/config"
1921
controller "github.com/gravitl/netmaker/controllers"
2022
"github.com/gravitl/netmaker/database"
21-
"github.com/gravitl/netmaker/db"
2223
"github.com/gravitl/netmaker/functions"
2324
"github.com/gravitl/netmaker/logger"
2425
"github.com/gravitl/netmaker/logic"
2526
"github.com/gravitl/netmaker/migrate"
2627
"github.com/gravitl/netmaker/models"
2728
"github.com/gravitl/netmaker/mq"
2829
"github.com/gravitl/netmaker/netclient/ncutils"
29-
"github.com/gravitl/netmaker/schema"
3030
"github.com/gravitl/netmaker/servercfg"
3131
"github.com/gravitl/netmaker/serverctl"
3232
_ "go.uber.org/automaxprocs"
@@ -62,6 +62,7 @@ func main() {
6262
if servercfg.DeployedByOperator() && !servercfg.IsPro {
6363
logic.SetFreeTierLimits()
6464
}
65+
defer db.CloseDB()
6566
defer database.CloseDB()
6667
ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGTERM, os.Interrupt)
6768
defer stop()
@@ -102,15 +103,19 @@ func initialize() { // Client Mode Prereq Check
102103
logger.FatalLog("error: must set NODE_ID, currently blank")
103104
}
104105

105-
if err = database.InitializeDatabase(); err != nil {
106-
logger.FatalLog("Error connecting to database: ", err.Error())
107-
}
108106
// initialize sql schema db.
109107
err = db.InitializeDB(schema.ListModels()...)
110108
if err != nil {
111-
logger.FatalLog("Error connecting to v1 database: ", err.Error())
109+
logger.FatalLog("error connecting to database: ", err.Error())
112110
}
111+
113112
logger.Log(0, "database successfully connected")
113+
114+
// initialize kv schema db.
115+
if err = database.InitializeDatabase(); err != nil {
116+
logger.FatalLog("error initializing database: ", err.Error())
117+
}
118+
114119
initializeUUID()
115120
//initialize cache
116121
_, _ = logic.GetNetworks()

0 commit comments

Comments
 (0)