reuse connections to DB (#76)

go.mod

@@ -11,6 +11,7 @@ require (
 	github.com/google/uuid v1.1.2
 	github.com/gosimple/slug v1.9.0
 	github.com/hailocab/go-geoindex v0.0.0-20160127134810-64631bfe9711
+	github.com/jackc/pgx/v4 v4.10.1
 	github.com/micro/dev v0.0.0-20201117163752-d3cfc9788dfa
 	github.com/micro/micro/v3 v3.1.2-0.20210311170414-40583563ada6
 	github.com/miekg/dns v1.1.31 // indirect

users/handler/handler.go

@@ -2,15 +2,18 @@ package handler
 
 import (
 	"context"
+	"database/sql"
 	"fmt"
 	"regexp"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/micro/micro/v3/service/auth"
 	"github.com/micro/micro/v3/service/errors"
 	pb "github.com/micro/services/users/proto"
 	"golang.org/x/crypto/bcrypt"
+	"gorm.io/driver/postgres"
 	"gorm.io/gorm"
 	"gorm.io/gorm/schema"
 )
@@ -63,13 +66,14 @@ type Token struct {
 }
 
 type Users struct {
-	Time         func() time.Time
+	sync.RWMutex
-	Dialector    gorm.Dialector
+	Time      func() time.Time
-	dbMigrations map[string]bool
+	dbConn    *sql.DB
+	gormConns map[string]*gorm.DB
 }
 
-func NewHandler(t func() time.Time, d gorm.Dialector) *Users {
+func NewHandler(t func() time.Time, dbConn *sql.DB) *Users {
-	return &Users{Time: t, Dialector: d, dbMigrations: map[string]bool{}}
+	return &Users{Time: t, dbConn: dbConn, gormConns: map[string]*gorm.DB{}}
 }
 
 func (u *Users) getDBConn(ctx context.Context) (*gorm.DB, error) {
@@ -77,23 +81,36 @@ func (u *Users) getDBConn(ctx context.Context) (*gorm.DB, error) {
 	if !ok {
 		return nil, fmt.Errorf("missing account from context")
 	}
-	db, err := gorm.Open(u.Dialector, &gorm.Config{
+	u.RLock()
-		NamingStrategy: schema.NamingStrategy{
+	if conn, ok := u.gormConns[acc.Issuer]; ok {
-			TablePrefix: fmt.Sprintf("%s_", strings.ReplaceAll(acc.Issuer, "-", "")),
+		u.RUnlock()
-		},
+		return conn, nil
-	})
+	}
+	u.RUnlock()
+	u.Lock()
+	// double check
+	if conn, ok := u.gormConns[acc.Issuer]; ok {
+		u.Unlock()
+		return conn, nil
+	}
+	defer u.Unlock()
+	db, err := gorm.Open(
+		postgres.New(postgres.Config{
+			Conn: u.dbConn,
+		}),
+		&gorm.Config{
+			NamingStrategy: schema.NamingStrategy{
+				TablePrefix: fmt.Sprintf("%s_", strings.ReplaceAll(acc.Issuer, "-", "")),
+			},
+		})
 	if err != nil {
 		return nil, err
 	}
-	// skip migration if we've already done it
-	if u.dbMigrations[acc.Issuer] {
-		return db, nil
-	}
 	if err := db.AutoMigrate(&User{}, &Token{}); err != nil {
 		return nil, err
 	}
 	// record success
-	u.dbMigrations[acc.Issuer] = true
+	u.gormConns[acc.Issuer] = db
 	return db, nil
 }
 

users/handler/handler_test.go

@@ -2,18 +2,18 @@ package handler_test
 
 import (
 	"context"
+	"database/sql"
 	"os"
 	"testing"
 	"time"
 
 	"github.com/micro/micro/v3/service/auth"
 	"github.com/stretchr/testify/assert"
-	"gorm.io/gorm/schema"
 
 	"github.com/micro/services/users/handler"
 	pb "github.com/micro/services/users/proto"
-	"gorm.io/driver/postgres"
+
-	"gorm.io/gorm"
+	_ "github.com/jackc/pgx/v4/stdlib"
 )
 
 func testHandler(t *testing.T) *handler.Users {
@@ -22,24 +22,16 @@ func testHandler(t *testing.T) *handler.Users {
 	if len(addr) == 0 {
 		addr = "postgresql://postgres@localhost:5432/postgres?sslmode=disable"
 	}
-	dial := postgres.Open(addr)
+	sqlDB, err := sql.Open("pgx", addr)
-	db, err := gorm.Open(dial, &gorm.Config{
-		NamingStrategy: schema.NamingStrategy{TablePrefix: "micro_"},
-	})
 	if err != nil {
-		t.Fatalf("Error connecting to database: %v", err)
+		t.Fatalf("Failed to open connection to DB %s", err)
 	}
-
 	// clean any data from a previous run
-	if err := db.Exec("DROP TABLE IF EXISTS micro_users, micro_tokens CASCADE").Error; err != nil {
+	if _, err := sqlDB.Exec("DROP TABLE IF EXISTS micro_users, micro_tokens CASCADE"); err != nil {
 		t.Fatalf("Error cleaning database: %v", err)
 	}
 
-	// migrate the database
+	return handler.NewHandler(time.Now, sqlDB)
-	if err := db.AutoMigrate(&handler.User{}, &handler.Token{}); err != nil {
-		t.Fatalf("Error migrating database: %v", err)
-	}
-	return handler.NewHandler(time.Now, dial)
 }
 
 func assertUsersMatch(t *testing.T, exp, act *pb.User) {

users/main.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"database/sql"
 	"time"
 
 	"github.com/micro/micro/v3/service"
@@ -8,7 +9,8 @@ import (
 	"github.com/micro/micro/v3/service/logger"
 	"github.com/micro/services/users/handler"
 	pb "github.com/micro/services/users/proto"
-	"gorm.io/driver/postgres"
+
+	_ "github.com/jackc/pgx/v4/stdlib"
 )
 
 var dbAddress = "postgresql://postgres:postgres@localhost:5432/users?sslmode=disable"
@@ -26,8 +28,13 @@ func main() {
 		logger.Fatalf("Error loading config: %v", err)
 	}
 	addr := cfg.String(dbAddress)
+
 	// Register handler
-	pb.RegisterUsersHandler(srv.Server(), handler.NewHandler(time.Now, postgres.Open(addr)))
+	sqlDB, err := sql.Open("pgx", addr)
+	if err != nil {
+		logger.Fatalf("Failed to open connection to DB %s", err)
+	}
+	pb.RegisterUsersHandler(srv.Server(), handler.NewHandler(time.Now, sqlDB))
 
 	// Run service
 	if err := srv.Run(); err != nil {