fix(restore): report max uid and namespace from offline restore

backup/run.go

@@ -240,7 +240,9 @@ func runRestoreCmd() error {
 			}
 			ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
 			defer cancel()
-			if _, err = zc.AssignIds(ctx, &pb.Num{Val: val, Type: typ}); err != nil {
+			// Bump rather than lease val more ids, as the online restore does: the lease must reach
+			// the backup's max, and a repeated restore must not move it again.
+			if err := worker.BumpLease(ctx, zc, val, typ); err != nil {
 				fmt.Printf("Failed to assign %s %d in Zero: %v\n",
 					pb.NumLeaseType_name[int32(typ)], val, err)
 				return err

worker/lease_test.go

@@ -0,0 +1,59 @@
+//go:build integration
+
+/*
+ * SPDX-FileCopyrightText: © 2017-2026 Istari Digital, Inc.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package worker
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/credentials/insecure"
+
+	"github.com/dgraph-io/dgraph/v25/protos/pb"
+	"github.com/dgraph-io/dgraph/v25/testutil"
+)
+
+// TestBumpLeaseIsIdempotent checks that bumping a lease raises it to the asked value once, and that
+// bumping again to the same or a lower value leaves it where it is. A restore calls BumpLease with
+// the backup's max uid, so restoring the same backup twice must not consume more ids.
+func TestBumpLeaseIsIdempotent(t *testing.T) {
+	conn, err := grpc.NewClient(testutil.GetSockAddrZero(),
+		grpc.WithTransportCredentials(insecure.NewCredentials()))
+	require.NoError(t, err)
+	defer func() { require.NoError(t, conn.Close()) }()
+	zc := pb.NewZeroClient(conn)
+
+	ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
+	defer cancel()
+
+	// nextUid leases a single uid, which tells us where the lease currently stands.
+	nextUid := func() uint64 {
+		res, err := zc.AssignIds(ctx, &pb.Num{Val: 1, Type: pb.Num_UID})
+		require.NoError(t, err)
+		return res.GetStartId()
+	}
+
+	// Anything else in the cluster may lease uids while this runs, so the checks allow some slack
+	// rather than expecting exact values. Leasing val more ids instead of bumping would move the
+	// lease by at least target on every call, far past that slack.
+	const slack = 100
+	target := nextUid() + 1000
+	require.NoError(t, BumpLease(ctx, zc, target, pb.Num_UID))
+	first := nextUid()
+	require.Greater(t, first, target)
+	require.Less(t, first, target+slack)
+
+	// Bumping again to target, or below it, must not move the lease.
+	require.NoError(t, BumpLease(ctx, zc, target, pb.Num_UID))
+	require.NoError(t, BumpLease(ctx, zc, target-500, pb.Num_UID))
+	second := nextUid()
+	require.Greater(t, second, first)
+	require.Less(t, second, first+slack)
+}

worker/offline_restore_test.go

@@ -0,0 +1,115 @@
+/*
+ * SPDX-FileCopyrightText: © 2017-2026 Istari Digital, Inc.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package worker
+
+import (
+	"context"
+	"path/filepath"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"google.golang.org/protobuf/proto"
+
+	"github.com/dgraph-io/badger/v4"
+	"github.com/dgraph-io/badger/v4/options"
+	"github.com/dgraph-io/dgraph/v25/posting"
+	"github.com/dgraph-io/dgraph/v25/protos/pb"
+	"github.com/dgraph-io/dgraph/v25/x"
+)
+
+// writeTestBackup writes a full backup of group 1 to backupDir using the production
+// BackupProcessor. It stores one value posting for every (namespace, uid) pair under the
+// predicate "name" and returns the backup's read timestamp.
+func writeTestBackup(t *testing.T, backupDir string, nsUids map[uint64][]uint64) uint64 {
+	t.Helper()
+	const commitTs, readTs = 5, 10
+
+	db, err := badger.OpenManaged(badger.DefaultOptions(t.TempDir()).WithLogger(nil))
+	require.NoError(t, err)
+	defer func() { require.NoError(t, db.Close()) }()
+
+	val, err := proto.Marshal(&pb.PostingList{Postings: []*pb.Posting{{
+		Uid: 0, Value: []byte("node"), ValType: pb.Posting_STRING, PostingType: pb.Posting_VALUE,
+	}}})
+	require.NoError(t, err)
+
+	var preds []string
+	txn := db.NewTransactionAt(commitTs, true)
+	for ns, uids := range nsUids {
+		attr := x.NamespaceAttr(ns, "name")
+		preds = append(preds, attr)
+		for _, uid := range uids {
+			key := x.DataKey(attr, uid)
+			require.NoError(t, txn.SetEntry(badger.NewEntry(key, val).WithMeta(posting.BitCompletePosting)))
+		}
+	}
+	require.NoError(t, txn.CommitAt(commitTs, nil))
+
+	unixTs := time.Now().UTC().Format("20060102.150405.000")
+	req := &pb.BackupRequest{
+		ReadTs:      readTs,
+		GroupId:     1,
+		UnixTs:      unixTs,
+		Destination: backupDir,
+		Predicates:  preds,
+		ForceFull:   true,
+	}
+	bp := NewBackupProcessor(db, req)
+	defer bp.Close()
+	_, err = bp.WriteBackup(context.Background())
+	require.NoError(t, err)
+
+	m := &Manifest{
+		ManifestBase: ManifestBase{
+			Type:        "full",
+			ReadTs:      readTs,
+			BackupId:    "offline-restore-lease",
+			BackupNum:   1,
+			Version:     x.ManifestVersion,
+			Path:        "dgraph." + unixTs,
+			Compression: "snappy",
+		},
+		Groups: map[uint32][]string{1: preds},
+	}
+	require.NoError(t, bp.CompleteBackup(context.Background(), m))
+	return readTs
+}
+
+// TestRunOfflineRestoreReportsMaxLease checks that the offline restore (the code path behind
+// `dgraph restore`) reports the max UID and namespace present in the backup. `dgraph restore -z`
+// uses these values to bump Zero's UID and namespace leases. If they come back as 0, Zero hands
+// out UIDs that already belong to restored nodes.
+func TestRunOfflineRestoreReportsMaxLease(t *testing.T) {
+	oldBadger, oldTmp := x.WorkerConfig.Badger, x.WorkerConfig.TmpDir
+	// BackupProcessor sizes its per-thread state from this; it must match the DB's stream goroutines.
+	x.WorkerConfig.Badger.NumGoroutines = badger.DefaultOptions("").NumGoroutines
+	x.WorkerConfig.TmpDir = t.TempDir()
+	defer func() { x.WorkerConfig.Badger, x.WorkerConfig.TmpDir = oldBadger, oldTmp }()
+
+	const maxUid, maxNs = uint64(0x2711), uint64(2)
+	nsUids := map[uint64][]uint64{
+		x.RootNamespace: {0x1, 0x2, 0x3, maxUid},
+		maxNs:           {0x10, 0x11},
+	}
+	backupDir := filepath.Join(t.TempDir(), "backups")
+	readTs := writeTestBackup(t, backupDir, nsUids)
+
+	// Control: the map phase does see the max UID and namespace in the backup.
+	mapRes, err := RunMapper(&pb.RestoreRequest{Location: backupDir, GroupId: 1, RestoreTs: 1}, t.TempDir())
+	require.NoError(t, err)
+	require.Equal(t, maxUid, mapRes.maxUid, "mapper max uid")
+	require.Equal(t, maxNs, mapRes.maxNs, "mapper max namespace")
+
+	result := RunOfflineRestore(filepath.Join(t.TempDir(), "p"), backupDir, "", "", nil, options.Snappy, 0)
+	require.NoError(t, result.Err)
+	require.Equal(t, readTs, result.Version)
+	assert.GreaterOrEqual(t, result.MaxLeaseUid, maxUid,
+		"offline restore must report the max uid in the backup so Zero's uid lease can be bumped")
+	assert.GreaterOrEqual(t, result.MaxLeaseNsId, maxNs,
+		"offline restore must report the max namespace in the backup so Zero's ns lease can be bumped")
+}

worker/online_restore.go

@@ -466,23 +466,25 @@ func bumpLease(ctx context.Context, mr *mapResult) error {
 	}
 
 	zc := pb.NewZeroClient(pl.Get())
-	bump := func(val uint64, typ pb.NumLeaseType) error {
-		_, err := zc.AssignIds(ctx, &pb.Num{Val: val, Type: typ, Bump: true})
-		if err != nil && strings.Contains(err.Error(), "Nothing to be leased") {
-			return nil
-		}
-		return err
-	}
-
-	if err := bump(mr.maxUid, pb.Num_UID); err != nil {
+	if err := BumpLease(ctx, zc, mr.maxUid, pb.Num_UID); err != nil {
 		return errors.Wrapf(err, "cannot update max uid lease after restore.")
 	}
-	if err := bump(mr.maxNs, pb.Num_NS_ID); err != nil {
+	if err := BumpLease(ctx, zc, mr.maxNs, pb.Num_NS_ID); err != nil {
 		return errors.Wrapf(err, "cannot update max namespace lease after restore.")
 	}
 	return nil
 }
 
+// BumpLease raises Zero's lease of the given type to at least val. It is a no-op when the lease
+// already reaches val, so a restore can call it any number of times without consuming ids.
+func BumpLease(ctx context.Context, zc pb.ZeroClient, val uint64, typ pb.NumLeaseType) error {
+	_, err := zc.AssignIds(ctx, &pb.Num{Val: val, Type: typ, Bump: true})
+	if err != nil && strings.Contains(err.Error(), "Nothing to be leased") {
+		return nil
+	}
+	return err
+}
+
 // create a config object from the request for use with enc package.
 func getEncConfig(req *pb.RestoreRequest) (*viper.Viper, error) {
 	config := viper.New()
@@ -559,6 +561,7 @@ func RunOfflineRestore(dir, location, backupId, keyFile string, key x.Sensitive,
 		}
 	}
 
+	var maxUid, maxNs uint64
 	for gid := range manifest.Groups {
 		req := &pb.RestoreRequest{
 			Location:          location,
@@ -577,9 +580,12 @@ func RunOfflineRestore(dir, location, backupId, keyFile string, key x.Sensitive,
 			}
 		}()
 
-		if _, err := RunMapper(req, mapDir); err != nil {
+		mapRes, err := RunMapper(req, mapDir)
+		if err != nil {
 			return LoadResult{Err: errors.Wrap(err, "RunRestore failed to map")}
 		}
+		maxUid = x.Max(maxUid, mapRes.maxUid)
+		maxNs = x.Max(maxNs, mapRes.maxNs)
 		pdir := filepath.Join(dir, fmt.Sprintf("p%d", gid))
 		db, err := badger.OpenManaged(badger.DefaultOptions(pdir).
 			WithCompression(ctype).
@@ -609,8 +615,8 @@ func RunOfflineRestore(dir, location, backupId, keyFile string, key x.Sensitive,
 			return LoadResult{Err: errors.Wrap(err, "RunRestore failed to write group id file")}
 		}
 	}
-	// TODO: Fix this return value.
-	return LoadResult{Version: manifest.ValidReadTs()}
+	// The caller bumps Zero's uid and namespace leases past these, as the online restore does.
+	return LoadResult{Version: manifest.ValidReadTs(), MaxLeaseUid: maxUid, MaxLeaseNsId: maxNs}
 }
 
 func buildPredsForDefaultNamespace(restorePreds []string, fromNamespace uint64) []string {