From 40738bedc8b344b19100c13b13271e0d24756a62 Mon Sep 17 00:00:00 2001 From: Dominic Wong Date: Wed, 24 Mar 2021 10:22:48 +0000 Subject: [PATCH] reuse connections to DB (#76) --- go.mod | 1 + users/handler/handler.go | 47 ++++++++++++++++++++++++----------- users/handler/handler_test.go | 22 ++++++---------- users/main.go | 11 ++++++-- 4 files changed, 49 insertions(+), 32 deletions(-) diff --git a/go.mod b/go.mod index 4df3474..4f7c473 100644 --- a/go.mod +++ b/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 diff --git a/users/handler/handler.go b/users/handler/handler.go index 644e123..f950085 100644 --- a/users/handler/handler.go +++ b/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 } diff --git a/users/handler/handler_test.go b/users/handler/handler_test.go index 38a183f..e871d5f 100644 --- a/users/handler/handler_test.go +++ b/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) { diff --git a/users/main.go b/users/main.go index 24e8770..90aeb42 100644 --- a/users/main.go +++ b/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 {