Skip to content

Commit

Permalink
feat(gorm/manager): Added gorm Manager (#126)
Browse files Browse the repository at this point in the history
Signed-off-by: Flc゛ <four_leaf_clover@foxmail.com>
  • Loading branch information
flc1125 authored Mar 4, 2024
1 parent 27680ea commit 061c4ae
Show file tree
Hide file tree
Showing 5 changed files with 87 additions and 0 deletions.
1 change: 1 addition & 0 deletions go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
4 changes: 4 additions & 0 deletions go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -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=
Expand All @@ -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=
Expand Down
Empty file added gorm/manager/README.md
Empty file.
39 changes: 39 additions & 0 deletions gorm/manager/manager.go
Original file line number Diff line number Diff line change
@@ -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.")
}
43 changes: 43 additions & 0 deletions gorm/manager/manager_test.go
Original file line number Diff line number Diff line change
@@ -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")
}
})
}

0 comments on commit 061c4ae

Please sign in to comment.