From 96e0f5067386131a71c8f4ef5d2d263779507450 Mon Sep 17 00:00:00 2001
From: Xiang Li <xiangli.cs@gmail.com>
Date: Fri, 20 Jan 2017 09:29:06 -0800
Subject: [PATCH] etcdctlv3: snapshot restore works with lease key

---
 e2e/ctl_v3_snapshot_test.go               | 16 ++++++++++++----
 etcdctl/ctlv3/command/snapshot_command.go |  7 ++++++-
 2 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/e2e/ctl_v3_snapshot_test.go b/e2e/ctl_v3_snapshot_test.go
index 8b5780360..673d84429 100644
--- a/e2e/ctl_v3_snapshot_test.go
+++ b/e2e/ctl_v3_snapshot_test.go
@@ -34,10 +34,18 @@ func TestCtlV3Snapshot(t *testing.T) { testCtl(t, snapshotTest) }
 func snapshotTest(cx ctlCtx) {
 	maintenanceInitKeys(cx)
 
+	leaseID, err := ctlV3LeaseGrant(cx, 100)
+	if err != nil {
+		cx.t.Fatalf("snapshot: ctlV3LeaseGrant error (%v)", err)
+	}
+	if err = ctlV3Put(cx, "withlease", "withlease", leaseID); err != nil {
+		cx.t.Fatalf("snapshot: ctlV3Put error (%v)", err)
+	}
+
 	fpath := "test.snapshot"
 	defer os.RemoveAll(fpath)
 
-	if err := ctlV3SnapshotSave(cx, fpath); err != nil {
+	if err = ctlV3SnapshotSave(cx, fpath); err != nil {
 		cx.t.Fatalf("snapshotTest ctlV3SnapshotSave error (%v)", err)
 	}
 
@@ -45,11 +53,11 @@ func snapshotTest(cx ctlCtx) {
 	if err != nil {
 		cx.t.Fatalf("snapshotTest getSnapshotStatus error (%v)", err)
 	}
-	if st.Revision != 4 {
+	if st.Revision != 5 {
 		cx.t.Fatalf("expected 4, got %d", st.Revision)
 	}
-	if st.TotalKey < 3 {
-		cx.t.Fatalf("expected at least 3, got %d", st.TotalKey)
+	if st.TotalKey < 4 {
+		cx.t.Fatalf("expected at least 4, got %d", st.TotalKey)
 	}
 }
 
diff --git a/etcdctl/ctlv3/command/snapshot_command.go b/etcdctl/ctlv3/command/snapshot_command.go
index 89e421fbb..28f2d8bc6 100644
--- a/etcdctl/ctlv3/command/snapshot_command.go
+++ b/etcdctl/ctlv3/command/snapshot_command.go
@@ -21,6 +21,7 @@ import (
 	"fmt"
 	"hash/crc32"
 	"io"
+	"math"
 	"os"
 	"path"
 	"reflect"
@@ -30,6 +31,7 @@ import (
 	"github.com/coreos/etcd/etcdserver"
 	"github.com/coreos/etcd/etcdserver/etcdserverpb"
 	"github.com/coreos/etcd/etcdserver/membership"
+	"github.com/coreos/etcd/lease"
 	"github.com/coreos/etcd/mvcc"
 	"github.com/coreos/etcd/mvcc/backend"
 	"github.com/coreos/etcd/pkg/fileutil"
@@ -371,7 +373,10 @@ func makeDB(snapdir, dbfile string, commit int) {
 	// update consistentIndex so applies go through on etcdserver despite
 	// having a new raft instance
 	be := backend.NewDefaultBackend(dbpath)
-	s := mvcc.NewStore(be, nil, (*initIndex)(&commit))
+	// a lessor never timeouts leases
+	lessor := lease.NewLessor(be, math.MaxInt64)
+
+	s := mvcc.NewStore(be, lessor, (*initIndex)(&commit))
 	id := s.TxnBegin()
 	btx := be.BatchTx()
 	del := func(k, v []byte) error {