feat(gorm/manager): Added gorm Manager

go.mod

@@ -55,6 +55,7 @@ require (
 	golang.org/x/arch v0.7.0 // indirect
 	golang.org/x/crypto v0.19.0 // indirect
 	golang.org/x/net v0.21.0 // indirect
+	golang.org/x/sync v0.5.0 // indirect
 	golang.org/x/sys v0.17.0 // indirect
 	google.golang.org/genproto/googleapis/rpc v0.0.0-20240213162025-012b6fc9bca9 // indirect
 	google.golang.org/grpc v1.61.1 // indirect

go.sum

@@ -120,6 +120,8 @@ golang.org/x/crypto v0.19.0 h1:ENy+Az/9Y1vSrlrvBSyna3PITt4tiZLf7sgCjZBX7Wo=
 golang.org/x/crypto v0.19.0/go.mod h1:Iy9bg/ha4yyC70EfRS8jz+B6ybOBKMaSxLj6P6oBDfU=
 golang.org/x/net v0.21.0 h1:AQyQV4dYCvJ7vGmJyKki9+PBdyvhkSd8EIx/qb0AYv4=
 golang.org/x/net v0.21.0/go.mod h1:bIjVDfnllIU7BJ2DNgfnXvpSvtn8VRwhlsaeUTyUS44=
+golang.org/x/sync v0.5.0 h1:60k92dhOjHxJkrqnwsfl8KuaHbn/5dl0lUPUklKo3qE=
+golang.org/x/sync v0.5.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk=
 golang.org/x/sys v0.0.0-20220811171246-fbc7d0a398ab/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.5.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
@@ -130,6 +132,8 @@ golang.org/x/term v0.17.0/go.mod h1:lLRBjIVuehSbZlaOtGMbcMncT+aqLLLmKrsjNrUguwk=
 golang.org/x/text v0.14.0 h1:ScX5w1eTa3QqT8oi6+ziP7dTV1S2+ALU0bI+0zXKWiQ=
 golang.org/x/text v0.14.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU=
 golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
+google.golang.org/genproto v0.0.0-20240205150955-31a09d347014 h1:g/4bk7P6TPMkAUbUhquq98xey1slwvuVJPosdBqYJlU=
+google.golang.org/genproto/googleapis/api v0.0.0-20231106174013-bbf56f31fb17 h1:JpwMPBpFN3uKhdaekDpiNlImDdkUAyiJ6ez/uxGaUSo=
 google.golang.org/genproto/googleapis/rpc v0.0.0-20240213162025-012b6fc9bca9 h1:hZB7eLIaYlW9qXRfCq/qDaPdbeY3757uARz5Vvfv+cY=
 google.golang.org/genproto/googleapis/rpc v0.0.0-20240213162025-012b6fc9bca9/go.mod h1:YUWgXUFRPfoYK1IHMuxH5K6nPEXSCzIMljnQ59lLRCk=
 google.golang.org/grpc v1.61.1 h1:kLAiWrZs7YeDM6MumDe7m3y4aM6wacLzM1Y/wiLP9XY=

gorm/manager/manager.go

@@ -0,0 +1,39 @@
+package manager
+
+import "gorm.io/gorm"
+
+type ConnectionFunc func() *gorm.DB
+
+type Manager struct {
+	*gorm.DB
+
+	connections map[string]*gorm.DB
+}
+
+func New(db *gorm.DB) *Manager {
+	return &Manager{
+		DB:          db,
+		connections: make(map[string]*gorm.DB),
+	}
+}
+
+func (m *Manager) Register(name string, db *gorm.DB) {
+	m.connections[name] = db
+}
+
+func (m *Manager) Conn(names ...string) *gorm.DB {
+	var name string
+	if len(names) > 0 {
+		name = names[0]
+	}
+
+	if name == "" {
+		return m.DB
+	}
+
+	if c, ok := m.connections[name]; ok {
+		return c
+	}
+
+	panic("The connection [" + name + "] is not registered.")
+}

gorm/manager/manager_test.go

@@ -0,0 +1,43 @@
+package manager
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"gorm.io/gorm"
+)
+
+var (
+	m                  *Manager
+	db1, db2, db3, db4 *gorm.DB
+)
+
+func init() {
+	m = New(db1)
+	m.Register("db2", db2)
+	m.Register("db3", db3)
+	m.Register("db4", db4)
+}
+
+func TestManager(t *testing.T) {
+	assert.Equal(t, db1, m.DB)
+	assert.Equal(t, db1, m.Conn())
+	assert.Equal(t, db2, m.Conn("db2"))
+	assert.Equal(t, db3, m.Conn("db3"))
+	assert.Equal(t, db4, m.Conn("db4"))
+	assert.Panics(t, func() {
+		m.Conn("db5")
+	})
+}
+
+func BenchmarkManager(b *testing.B) {
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			_ = m.DB
+			_ = m.Conn()
+			_ = m.Conn("db2")
+			_ = m.Conn("db3")
+			_ = m.Conn("db4")
+		}
+	})
+}