Merge pull request #571 from go-redis/fix/on-connect

Fix/on connect

redis.go

@@ -21,11 +21,6 @@ func (c *baseClient) String() string {
 	return fmt.Sprintf("Redis<%s db:%d>", c.getAddr(), c.opt.DB)
 }
 
-// Options returns read-only Options that were used to create the client.
-func (c *baseClient) Options() *Options {
-	return c.opt
-}
-
 func (c *baseClient) conn() (*pool.Conn, bool, error) {
 	cn, isNew, err := c.connPool.Get()
 	if err != nil {
@@ -55,13 +50,23 @@ func (c *baseClient) putConn(cn *pool.Conn, err error) bool {
 func (c *baseClient) initConn(cn *pool.Conn) error {
 	cn.Inited = true
 
-	if c.opt.Password == "" && c.opt.DB == 0 && !c.opt.ReadOnly {
+	if c.opt.Password == "" &&
+		c.opt.DB == 0 &&
+		!c.opt.ReadOnly &&
+		c.opt.OnConnect == nil {
 		return nil
 	}
 
-	// Temp client for Auth and Select.
-	client := newClient(c.opt, pool.NewSingleConnPool(cn))
-	_, err := client.Pipelined(func(pipe Pipeliner) error {
+	// Temp client to initialize connection.
+	conn := &Conn{
+		baseClient: baseClient{
+			opt:      c.opt,
+			connPool: pool.NewSingleConnPool(cn),
+		},
+	}
+	conn.setProcessor(conn.Process)
+
+	_, err := conn.Pipelined(func(pipe Pipeliner) error {
 		if c.opt.Password != "" {
 			pipe.Auth(c.opt.Password)
 		}
@@ -76,7 +81,14 @@ func (c *baseClient) initConn(cn *pool.Conn) error {
 
 		return nil
 	})
-	return err
+	if err != nil {
+		return err
+	}
+
+	if c.opt.OnConnect != nil {
+		return c.opt.OnConnect(conn)
+	}
+	return nil
 }
 
 func (c *baseClient) Process(cmd Cmder) error {
@@ -189,7 +201,7 @@ func (c *baseClient) pipelineExecer(p pipelineProcessor) pipelineExecer {
 	}
 }
 
-func (c *baseClient) pipelineProcessCmds(cn *pool.Conn, cmds []Cmder) (retry bool, firstErr error) {
+func (c *baseClient) pipelineProcessCmds(cn *pool.Conn, cmds []Cmder) (bool, error) {
 	cn.SetWriteTimeout(c.opt.WriteTimeout)
 	if err := writeCmd(cn, cmds...); err != nil {
 		setCmdsErr(cmds, err)
@@ -301,7 +313,7 @@ func newClient(opt *Options, pool pool.Pooler) *Client {
 			connPool: pool,
 		},
 	}
-	client.cmdable.process = client.Process
+	client.setProcessor(client.Process)
 	return &client
 }
 
@@ -314,10 +326,15 @@ func NewClient(opt *Options) *Client {
 func (c *Client) copy() *Client {
 	c2 := new(Client)
 	*c2 = *c
-	c2.cmdable.process = c2.Process
+	c2.setProcessor(c2.Process)
 	return c2
 }
 
+// Options returns read-only Options that were used to create the client.
+func (c *Client) Options() *Options {
+	return c.opt
+}
+
 // PoolStats returns connection pool stats.
 func (c *Client) PoolStats() *PoolStats {
 	s := c.connPool.Stats()
@@ -339,8 +356,7 @@ func (c *Client) Pipeline() Pipeliner {
 	pipe := Pipeline{
 		exec: c.pipelineExecer(c.pipelineProcessCmds),
 	}
-	pipe.cmdable.process = pipe.Process
-	pipe.statefulCmdable.process = pipe.Process
+	pipe.setProcessor(pipe.Process)
 	return &pipe
 }
 
@@ -353,8 +369,7 @@ func (c *Client) TxPipeline() Pipeliner {
 	pipe := Pipeline{
 		exec: c.pipelineExecer(c.txPipelineProcessCmds),
 	}
-	pipe.cmdable.process = pipe.Process
-	pipe.statefulCmdable.process = pipe.Process
+	pipe.setProcessor(pipe.Process)
 	return &pipe
 }
 
@@ -384,3 +399,36 @@ func (c *Client) PSubscribe(channels ...string) *PubSub {
 	}
 	return pubsub
 }
+
+//------------------------------------------------------------------------------
+
+// Conn is like Client, but its pool contains single connection.
+type Conn struct {
+	baseClient
+	statefulCmdable
+}
+
+func (c *Conn) Pipelined(fn func(Pipeliner) error) ([]Cmder, error) {
+	return c.Pipeline().pipelined(fn)
+}
+
+func (c *Conn) Pipeline() Pipeliner {
+	pipe := Pipeline{
+		exec: c.pipelineExecer(c.pipelineProcessCmds),
+	}
+	pipe.setProcessor(pipe.Process)
+	return &pipe
+}
+
+func (c *Conn) TxPipelined(fn func(Pipeliner) error) ([]Cmder, error) {
+	return c.TxPipeline().pipelined(fn)
+}
+
+// TxPipeline acts like Pipeline, but wraps queued commands with MULTI/EXEC.
+func (c *Conn) TxPipeline() Pipeliner {
+	pipe := Pipeline{
+		exec: c.pipelineExecer(c.txPipelineProcessCmds),
+	}
+	pipe.setProcessor(pipe.Process)
+	return &pipe
+}