聊聊gorm的读写分离

type DBResolver struct {
    *gorm.DB
    configs          []Config
    resolvers        map[string]*resolver
    global           *resolver
    prepareStmtStore map[gorm.ConnPool]*gorm.PreparedStmtDB
    compileCallbacks []func(gorm.ConnPool) error
}

func (dr *DBResolver) Name() string {
    return "gorm:db_resolver"
}

func (dr *DBResolver) Initialize(db *gorm.DB) error {
    dr.DB = db
    dr.registerCallbacks(db)
    return dr.compile()
}

func (dr *DBResolver) registerCallbacks(db *gorm.DB) {
    dr.Callback().Create().Before("*").Register("gorm:db_resolver", dr.switchSource)
    dr.Callback().Query().Before("*").Register("gorm:db_resolver", dr.switchReplica)
    dr.Callback().Update().Before("*").Register("gorm:db_resolver", dr.switchSource)
    dr.Callback().Delete().Before("*").Register("gorm:db_resolver", dr.switchSource)
    dr.Callback().Row().Before("*").Register("gorm:db_resolver", dr.switchReplica)
    dr.Callback().Raw().Before("*").Register("gorm:db_resolver", dr.switchGuess)
}

func (dr *DBResolver) switchSource(db *gorm.DB) {
    if !isTransaction(db.Statement.ConnPool) {
        db.Statement.ConnPool = dr.resolve(db.Statement, Write)
    }
}

func (dr *DBResolver) switchReplica(db *gorm.DB) {
    if !isTransaction(db.Statement.ConnPool) {
        if rawSQL := db.Statement.SQL.String(); len(rawSQL) > 0 {
            dr.switchGuess(db)
        } else {
            _, locking := db.Statement.Clauses["FOR"]
            if _, ok := db.Statement.Clauses[writeName]; ok || locking {
                db.Statement.ConnPool = dr.resolve(db.Statement, Write)
            } else {
                db.Statement.ConnPool = dr.resolve(db.Statement, Read)
            }
        }
    }
}

func (dr *DBResolver) switchGuess(db *gorm.DB) {
    if !isTransaction(db.Statement.ConnPool) {
        if _, ok := db.Statement.Clauses[writeName]; ok {
            db.Statement.ConnPool = dr.resolve(db.Statement, Write)
        } else if rawSQL := strings.TrimSpace(db.Statement.SQL.String()); len(rawSQL) > 10 && strings.EqualFold(rawSQL[:6], "select") && !strings.EqualFold(rawSQL[len(rawSQL)-10:], "for update") {
            db.Statement.ConnPool = dr.resolve(db.Statement, Read)
        } else {
            db.Statement.ConnPool = dr.resolve(db.Statement, Write)
        }
    }
}

func (dr *DBResolver) resolve(stmt *gorm.Statement, op Operation) gorm.ConnPool {
    if len(dr.resolvers) > 0 {
        if u, ok := stmt.Clauses[usingName].Expression.(using); ok && u.Use != "" {
            if r, ok := dr.resolvers[u.Use]; ok {
                return r.resolve(stmt, op)
            }
        }

        if stmt.Table != "" {
            if r, ok := dr.resolvers[stmt.Table]; ok {
                return r.resolve(stmt, op)
            }
        }

        if stmt.Schema != nil {
            if r, ok := dr.resolvers[stmt.Schema.Table]; ok {
                return r.resolve(stmt, op)
            }
        }

        if rawSQL := stmt.SQL.String(); rawSQL != "" {
            if r, ok := dr.resolvers[getTableFromRawSQL(rawSQL)]; ok {
                return r.resolve(stmt, op)
            }
        }
    }

    if dr.global != nil {
        return dr.global.resolve(stmt, op)
    }

    return stmt.ConnPool
}

func (dr *DBResolver) compile() error {
    for _, config := range dr.configs {
        if err := dr.compileConfig(config); err != nil {
            return err
        }
    }
    return nil
}

func (dr *DBResolver) compileConfig(config Config) (err error) {
    var (
        connPool = dr.DB.Config.ConnPool
        r        = resolver{
            dbResolver: dr,
            policy:     config.Policy,
        }
    )

    if preparedStmtDB, ok := connPool.(*gorm.PreparedStmtDB); ok {
        connPool = preparedStmtDB.ConnPool
    }

    if len(config.Sources) == 0 {
        r.sources = []gorm.ConnPool{connPool}
    } else if r.sources, err = dr.convertToConnPool(config.Sources); err != nil {
        return err
    }

    if len(config.Replicas) == 0 {
        r.replicas = r.sources
    } else if r.replicas, err = dr.convertToConnPool(config.Replicas); err != nil {
        return err
    }

    if len(config.datas) > 0 {
        for _, data := range config.datas {
            if t, ok := data.(string); ok {
                dr.resolvers[t] = &r
            } else {
                stmt := &gorm.Statement{DB: dr.DB}
                if err := stmt.Parse(data); err == nil {
                    dr.resolvers[stmt.Table] = &r
                } else {
                    return err
                }
            }
        }
    } else if dr.global == nil {
        dr.global = &r
    } else {
        return errors.New("conflicted global resolver")
    }

    for _, fc := range dr.compileCallbacks {
        if err = r.call(fc); err != nil {
            return err
        }
    }

    return nil
}

type resolver struct {
    sources    []gorm.ConnPool
    replicas   []gorm.ConnPool
    policy     Policy
    dbResolver *DBResolver
}

func (r *resolver) resolve(stmt *gorm.Statement, op Operation) (connPool gorm.ConnPool) {
    if op == Read {
        if len(r.replicas) == 1 {
            connPool = r.replicas[0]
        } else {
            connPool = r.policy.Resolve(r.replicas)
        }
    } else if len(r.sources) == 1 {
        connPool = r.sources[0]
    } else {
        connPool = r.policy.Resolve(r.sources)
    }

    if stmt.DB.PrepareStmt {
        if preparedStmt, ok := r.dbResolver.prepareStmtStore[connPool]; ok {
            return &gorm.PreparedStmtDB{
                ConnPool: connPool,
                Mux:      preparedStmt.Mux,
                Stmts:    preparedStmt.Stmts,
            }
        }
    }

    return
}

type Policy interface {
    Resolve([]gorm.ConnPool) gorm.ConnPool
}

type RandomPolicy struct {
}

func (RandomPolicy) Resolve(connPools []gorm.ConnPool) gorm.ConnPool {
    return connPools[rand.Intn(len(connPools))]
}

func dbResolverDemo() {
    db, _ := gorm.Open(mysql.Open("master_dsn"), &gorm.Config{})
    dbResolverCfg := dbresolver.Config{
        Sources:  []gorm.Dialector{mysql.Open("master_dsn")},
        Replicas: []gorm.Dialector{mysql.Open("replica_a_dsn"), mysql.Open("replica_b_dsn")},
        Policy:   dbresolver.RandomPolicy{}}
    readWritePlugin := dbresolver.Register(dbResolverCfg)
    db.Use(readWritePlugin)
}