DB RPC driver: add master/replica support (#17792)

Automatic Merge

app/plugin_api_tests/test_db_driver/main.go

@@ -32,15 +32,20 @@ func (p *MyPlugin) MessageWillBePosted(_ *plugin.Context, _ *model.Post) (*model
 	store := sqlstore.New(p.API.GetUnsanitizedConfig().SqlSettings, nil)
 	store.GetMaster().Db.Close()
 
-	store.GetMaster().Db = sql.OpenDB(driver.NewConnector(p.Driver))
+	for _, isMaster := range []bool{true, false} {
-	defer store.GetMaster().Db.Close()
+		// We replace the master DB with master and replica both just to make
+		// gorp APIs work.
+		store.GetMaster().Db = sql.OpenDB(driver.NewConnector(p.Driver, isMaster))
 
-	// Testing with a handful of stores
+		// Testing with a handful of stores
-	storetest.TestPostStore(p.t, store, store)
+		storetest.TestPostStore(p.t, store, store)
-	storetest.TestUserStore(p.t, store, store)
+		storetest.TestUserStore(p.t, store, store)
-	storetest.TestTeamStore(p.t, store)
+		storetest.TestTeamStore(p.t, store)
-	storetest.TestChannelStore(p.t, store, store)
+		storetest.TestChannelStore(p.t, store, store)
-	storetest.TestBotStore(p.t, store, store)
+		storetest.TestBotStore(p.t, store, store)
+
+		store.GetMaster().Db.Close()
+	}
 
 	// Use the API to instantiate the driver
 	// And then run the full suite of tests.

app/plugin_db_driver.go

@@ -39,8 +39,12 @@ func NewDriverImpl(s *Server) *DriverImpl {
 	}
 }
 
-func (d *DriverImpl) Conn() (string, error) {
+func (d *DriverImpl) Conn(isMaster bool) (string, error) {
-	conn, err := d.s.sqlStore.GetMaster().Db.Conn(context.Background())
+	dbFunc := d.s.sqlStore.GetMaster
+	if !isMaster {
+		dbFunc = d.s.sqlStore.GetReplica
+	}
+	conn, err := dbFunc().Db.Conn(context.Background())
 	if err != nil {
 		return "", err
 	}

plugin/client.go

@@ -56,6 +56,7 @@ func (p *MattermostPlugin) SetHelpers(helpers Helpers) {
 	p.Helpers = helpers
 }
 
+// SetDriver sets the RPC client implementation to talk with the server.
 func (p *MattermostPlugin) SetDriver(driver Driver) {
 	p.Driver = driver
 }

plugin/db_rpc.go

@@ -42,9 +42,9 @@ type Z_DbBoolReturn struct {
 	A bool
 }
 
-func (db *dbRPCClient) Conn() (string, error) {
+func (db *dbRPCClient) Conn(isMaster bool) (string, error) {
 	ret := &Z_DbStrErrReturn{}
-	err := db.client.Call("Plugin.Conn", struct{}{}, ret)
+	err := db.client.Call("Plugin.Conn", isMaster, ret)
 	if err != nil {
 		log.Printf("error during Plugin.Conn: %v", err)
 	}
@@ -52,8 +52,8 @@ func (db *dbRPCClient) Conn() (string, error) {
 	return ret.A, ret.B
 }
 
-func (db *dbRPCServer) Conn(_ struct{}, ret *Z_DbStrErrReturn) error {
+func (db *dbRPCServer) Conn(isMaster bool, ret *Z_DbStrErrReturn) error {
-	ret.A, ret.B = db.dbImpl.Conn()
+	ret.A, ret.B = db.dbImpl.Conn(isMaster)
 	ret.B = encodableError(ret.B)
 	return nil
 }

plugin/helpers.go

@@ -109,9 +109,13 @@ type ResultContainer struct {
 	RowsAffectedError error
 }
 
+// Driver is a sql driver interface that is used by plugins to perform
+// raw SQL queries without opening DB connections by themselves. This interface
+// is not subject to backward compatibility guarantees and is only meant to be
+// used by plugins built by the Mattermost team.
 type Driver interface {
 	// Connection
-	Conn() (string, error)
+	Conn(isMaster bool) (string, error)
 	ConnPing(connID string) error
 	ConnClose(connID string) error
 	ConnQuery(connID, q string, args []driver.NamedValue) (string, error)         // rows

shared/driver/driver.go

@@ -23,15 +23,19 @@ var (
 // Connector is the DB connector which is used to
 // communicate with the DB API.
 type Connector struct {
-	api plugin.Driver
+	api      plugin.Driver
+	isMaster bool
 }
 
-func NewConnector(api plugin.Driver) *Connector {
+// NewConnector returns a DB connector that can be used to return a sql.DB object.
-	return &Connector{api: api}
+// It takes a plugin.Driver implementation and a boolean flag to indicate whether
+// to connect to a master or replica DB instance.
+func NewConnector(api plugin.Driver, isMaster bool) *Connector {
+	return &Connector{api: api, isMaster: isMaster}
 }
 
 func (c *Connector) Connect(_ context.Context) (driver.Conn, error) {
-	connID, err := c.api.Conn()
+	connID, err := c.api.Conn(c.isMaster)
 	if err != nil {
 		return nil, err
 	}