From 68eaf4083a7316403d8b9517dcc3c9045ddd0dd3 Mon Sep 17 00:00:00 2001
From: Gyu-Ho Lee <gyuhox@gmail.com>
Date: Thu, 12 May 2016 19:25:42 -0700
Subject: [PATCH] clientv3: WithRequireLeader

---
 clientv3/client.go              | 10 ++++++++++
 clientv3/integration/kv_test.go | 26 ++++++++++++++++++++++++++
 2 files changed, 36 insertions(+)

diff --git a/clientv3/client.go b/clientv3/client.go
index 6ea77ca20..0981e35a0 100644
--- a/clientv3/client.go
+++ b/clientv3/client.go
@@ -24,9 +24,12 @@ import (
 	"sync"
 	"time"
 
+	"github.com/coreos/etcd/etcdserver/api/v3rpc/rpctypes"
+
 	"golang.org/x/net/context"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials"
+	"google.golang.org/grpc/metadata"
 )
 
 var (
@@ -160,6 +163,13 @@ func (c *Client) Dial(endpoint string) (*grpc.ClientConn, error) {
 	return conn, nil
 }
 
+// WithRequireLeader requires client requests to only succeed
+// when the cluster has a leader.
+func WithRequireLeader(ctx context.Context) context.Context {
+	md := metadata.Pairs(rpctypes.MetadataRequireLeaderKey, rpctypes.MetadataHasLeader)
+	return metadata.NewContext(ctx, md)
+}
+
 func newClient(cfg *Config) (*Client, error) {
 	if cfg == nil {
 		cfg = &Config{RetryDialer: dialEndpointList}
diff --git a/clientv3/integration/kv_test.go b/clientv3/integration/kv_test.go
index b449e0de0..e7bb6f2c9 100644
--- a/clientv3/integration/kv_test.go
+++ b/clientv3/integration/kv_test.go
@@ -110,6 +110,32 @@ func TestKVPut(t *testing.T) {
 	}
 }
 
+func TestKVPutWithRequireLeader(t *testing.T) {
+	// this test might block for a few seconds, make it parallel to speed up the test.
+	t.Parallel()
+
+	defer testutil.AfterTest(t)
+
+	clus := integration.NewClusterV3(t, &integration.ClusterConfig{Size: 3})
+	defer clus.Terminate(t)
+
+	clus.Members[1].Stop(t)
+	clus.Members[2].Stop(t)
+
+	// wait for election timeout, then member[0] will not have a leader.
+	var (
+		electionTicks = 10
+		tickDuration  = 10 * time.Millisecond
+	)
+	time.Sleep(time.Duration(3*electionTicks) * tickDuration)
+
+	kv := clientv3.NewKV(clus.Client(0))
+	_, err := kv.Put(clientv3.WithRequireLeader(context.Background()), "foo", "bar")
+	if err != rpctypes.ErrNoLeader {
+		t.Fatal(err)
+	}
+}
+
 func TestKVRange(t *testing.T) {
 	defer testutil.AfterTest(t)