reuse connections to DB (#76)

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
-	Dialector    gorm.Dialector
-	dbMigrations map[string]bool
+	sync.RWMutex
+	Time      func() time.Time
+	dbConn    *sql.DB
+	gormConns map[string]*gorm.DB
 }
 
-func NewHandler(t func() time.Time, d gorm.Dialector) *Users {
-	return &Users{Time: t, Dialector: d, dbMigrations: map[string]bool{}}
+func NewHandler(t func() time.Time, dbConn *sql.DB) *Users {
+	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{
-		NamingStrategy: schema.NamingStrategy{
-			TablePrefix: fmt.Sprintf("%s_", strings.ReplaceAll(acc.Issuer, "-", "")),
-		},
-	})
+	u.RLock()
+	if conn, ok := u.gormConns[acc.Issuer]; ok {
+		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)
-	db, err := gorm.Open(dial, &gorm.Config{
-		NamingStrategy: schema.NamingStrategy{TablePrefix: "micro_"},
-	})
+	sqlDB, err := sql.Open("pgx", addr)
 	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
-	if err := db.AutoMigrate(&handler.User{}, &handler.Token{}); err != nil {
-		t.Fatalf("Error migrating database: %v", err)
-	}
-	return handler.NewHandler(time.Now, dial)
+	return handler.NewHandler(time.Now, sqlDB)
 }
 
 func assertUsersMatch(t *testing.T, exp, act *pb.User) {