From c2b32de7174e72a1001f0c7bc37e215a70c1ad2c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Flc=E3=82=9B?= <four_leaf_clover@foxmail.com>
Date: Mon, 4 Mar 2024 18:35:18 +0800
Subject: [PATCH] feat(gorm/manager): Added gorm Manager
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: Flc゛ <four_leaf_clover@foxmail.com>
---
 go.mod                       |  1 +
 go.sum                       |  4 ++++
 gorm/manager/README.md       |  0
 gorm/manager/manager.go      | 39 ++++++++++++++++++++++++++++++++
 gorm/manager/manager_test.go | 43 ++++++++++++++++++++++++++++++++++++
 5 files changed, 87 insertions(+)
 create mode 100644 gorm/manager/README.md
 create mode 100644 gorm/manager/manager.go
 create mode 100644 gorm/manager/manager_test.go

diff --git a/go.mod b/go.mod
index 891fb877..4c793467 100644
--- a/go.mod
+++ b/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
diff --git a/go.sum b/go.sum
index 430a7e3e..ab20ab91 100644
--- a/go.sum
+++ b/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=
diff --git a/gorm/manager/README.md b/gorm/manager/README.md
new file mode 100644
index 00000000..e69de29b
diff --git a/gorm/manager/manager.go b/gorm/manager/manager.go
new file mode 100644
index 00000000..f6cdbdfe
--- /dev/null
+++ b/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.")
+}
diff --git a/gorm/manager/manager_test.go b/gorm/manager/manager_test.go
new file mode 100644
index 00000000..4579629f
--- /dev/null
+++ b/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")
+		}
+	})
+}