@@ -13,6 +13,7 @@ import (
1313 "gorm.io/driver/postgres"
1414 "gorm.io/gorm"
1515 "gorm.io/hints"
16+ "gorm.io/plugin/dbresolver"
1617)
1718
1819type Order struct {
@@ -37,35 +38,80 @@ func databaseURL() string {
3738 return databaseURL
3839}
3940
41+ func databaseReadURL () string {
42+ databaseURL := os .Getenv ("DATABASE_READ_URL" )
43+ if len (databaseURL ) == 0 {
44+ databaseURL = "postgres://localhost:5432/sharding-read-test?sslmode=disable"
45+ if os .Getenv ("DIALECTOR" ) == "mysql" {
46+ databaseURL = "root@tcp(127.0.0.1:3306)/sharding-read-test?charset=utf8mb4"
47+ }
48+ }
49+ return databaseURL
50+ }
51+
52+ func databaseWriteURL () string {
53+ databaseURL := os .Getenv ("DATABASE_WRITE_URL" )
54+ if len (databaseURL ) == 0 {
55+ databaseURL = "postgres://localhost:5432/sharding-write-test?sslmode=disable"
56+ if os .Getenv ("DIALECTOR" ) == "mysql" {
57+ databaseURL = "root@tcp(127.0.0.1:3306)/sharding-write-test?charset=utf8mb4"
58+ }
59+ }
60+ return databaseURL
61+ }
62+
4063var (
4164 dbConfig = postgres.Config {
4265 DSN : databaseURL (),
4366 PreferSimpleProtocol : true ,
4467 }
45- db * gorm.DB
46-
47- shardingConfig = Config {
48- DoubleWrite : true ,
49- ShardingKey : "user_id" ,
50- NumberOfShards : 4 ,
51- PrimaryKeyGenerator : PKSnowflake ,
68+ dbReadConfig = postgres.Config {
69+ DSN : databaseReadURL (),
70+ PreferSimpleProtocol : true ,
5271 }
72+ dbWriteConfig = postgres.Config {
73+ DSN : databaseWriteURL (),
74+ PreferSimpleProtocol : true ,
75+ }
76+ db , dbRead , dbWrite * gorm.DB
5377
54- middleware = Register (shardingConfig , & Order {})
55- node , _ = snowflake .NewNode (1 )
78+ shardingConfig Config
79+ middleware * Sharding
80+ node , _ = snowflake .NewNode (1 )
5681)
5782
5883func init () {
5984 if os .Getenv ("DIALECTOR" ) == "mysql" {
6085 db , _ = gorm .Open (mysql .Open (databaseURL ()), & gorm.Config {
6186 DisableForeignKeyConstraintWhenMigrating : true ,
6287 })
88+ dbRead , _ = gorm .Open (mysql .Open (databaseReadURL ()), & gorm.Config {
89+ DisableForeignKeyConstraintWhenMigrating : true ,
90+ })
91+ dbWrite , _ = gorm .Open (mysql .Open (databaseWriteURL ()), & gorm.Config {
92+ DisableForeignKeyConstraintWhenMigrating : true ,
93+ })
6394 } else {
6495 db , _ = gorm .Open (postgres .New (dbConfig ), & gorm.Config {
6596 DisableForeignKeyConstraintWhenMigrating : true ,
6697 })
98+ dbRead , _ = gorm .Open (postgres .New (dbReadConfig ), & gorm.Config {
99+ DisableForeignKeyConstraintWhenMigrating : true ,
100+ })
101+ dbWrite , _ = gorm .Open (postgres .New (dbWriteConfig ), & gorm.Config {
102+ DisableForeignKeyConstraintWhenMigrating : true ,
103+ })
104+ }
105+
106+ shardingConfig = Config {
107+ DoubleWrite : true ,
108+ ShardingKey : "user_id" ,
109+ NumberOfShards : 4 ,
110+ PrimaryKeyGenerator : PKSnowflake ,
67111 }
68112
113+ middleware = Register (shardingConfig , & Order {})
114+
69115 fmt .Println ("Clean only tables ..." )
70116 dropTables ()
71117 fmt .Println ("AutoMigrate tables ..." )
@@ -80,6 +126,16 @@ func init() {
80126 user_id bigint,
81127 product text
82128 )` )
129+ dbRead .Exec (`CREATE TABLE ` + table + ` (
130+ id bigint PRIMARY KEY,
131+ user_id bigint,
132+ product text
133+ )` )
134+ dbWrite .Exec (`CREATE TABLE ` + table + ` (
135+ id bigint PRIMARY KEY,
136+ user_id bigint,
137+ product text
138+ )` )
83139 }
84140
85141 db .Use (middleware )
@@ -89,6 +145,8 @@ func dropTables() {
89145 tables := []string {"orders" , "orders_0" , "orders_1" , "orders_2" , "orders_3" , "categories" }
90146 for _ , table := range tables {
91147 db .Exec ("DROP TABLE IF EXISTS " + table )
148+ dbRead .Exec ("DROP TABLE IF EXISTS " + table )
149+ dbWrite .Exec ("DROP TABLE IF EXISTS " + table )
92150 db .Exec (("DROP SEQUENCE IF EXISTS gorm_sharding_" + table + "_id_seq" ))
93151 }
94152}
@@ -264,6 +322,39 @@ func TestPKPGSequence(t *testing.T) {
264322 assert .Equal (t , expected , middleware .LastQuery ())
265323}
266324
325+ func TestReadWriteSplitting (t * testing.T ) {
326+ dbRead .Exec ("INSERT INTO orders_0 (id, product, user_id) VALUES(1, 'iPad', 100)" )
327+ dbWrite .Exec ("INSERT INTO orders_0 (id, product, user_id) VALUES(1, 'iPad', 100)" )
328+
329+ var db * gorm.DB
330+ if os .Getenv ("DIALECTOR" ) == "mysql" {
331+ db , _ = gorm .Open (mysql .Open (databaseURL ()), & gorm.Config {
332+ DisableForeignKeyConstraintWhenMigrating : true ,
333+ })
334+ } else {
335+ db , _ = gorm .Open (postgres .New (dbConfig ), & gorm.Config {
336+ DisableForeignKeyConstraintWhenMigrating : true ,
337+ })
338+ }
339+
340+ db .Use (dbresolver .Register (dbresolver.Config {
341+ Sources : []gorm.Dialector {dbWrite .Dialector },
342+ Replicas : []gorm.Dialector {dbRead .Dialector },
343+ }))
344+ db .Use (middleware )
345+
346+ var order Order
347+ db .Model (& Order {}).Where ("user_id" , 100 ).Find (& order )
348+ assert .Equal (t , "iPad" , order .Product )
349+
350+ db .Model (& Order {}).Where ("user_id" , 100 ).Update ("product" , "iPhone" )
351+ db .Table ("orders_0" ).Where ("user_id" , 100 ).Find (& order )
352+ assert .Equal (t , "iPad" , order .Product )
353+
354+ dbWrite .Table ("orders_0" ).Where ("user_id" , 100 ).Find (& order )
355+ assert .Equal (t , "iPhone" , order .Product )
356+ }
357+
267358func assertQueryResult (t * testing.T , expected string , tx * gorm.DB ) {
268359 t .Helper ()
269360 assert .Equal (t , toDialect (expected ), middleware .LastQuery ())
0 commit comments