Skip to content

Commit

Permalink
add more test cases
Browse files Browse the repository at this point in the history
  • Loading branch information
@3pointer
3pointer committed Jan 22, 2025
1 parent 85bc6a4 commit f000964
Show file tree
Hide file tree
Showing 2 changed files with 160 additions and 14 deletions.
22 changes: 14 additions & 8 deletions br/pkg/restore/utils/rewrite_rule.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -49,9 +49,9 @@ type RewriteRules struct {
// used to record checkpoint data
NewTableID int64

shiftStartTs uint64
startTs uint64
restoredTs uint64
ShiftStartTs uint64
StartTs uint64
RestoredTs uint64
// used to record backup files to pitr.
// note: should NewTableID merged with this?
TableIDRemapHint []TableIDRemap
Expand Down Expand Up @@ -98,23 +98,29 @@ func (r *RewriteRules) Append(other RewriteRules) {
}

func (r *RewriteRules) SetTimeRangeFilter(cfName string) error {
if r.startTs == 0 || r.restoredTs == 0 {
// for some sst files like db restore copy ssts, we don't need to set the time range filter
if r.StartTs == 0 || r.RestoredTs == 0 {
return nil
}

// check if shift start ts is greater than start ts
if r.ShiftStartTs > r.StartTs {
return errors.Errorf("shift start ts %d is greater than start ts %d", r.ShiftStartTs, r.StartTs)
}

var ignoreBeforeTs uint64
switch {
case strings.Contains(cfName, DefaultCFName):
ignoreBeforeTs = r.shiftStartTs
ignoreBeforeTs = r.ShiftStartTs
case strings.Contains(cfName, WriteCFName):
ignoreBeforeTs = r.startTs
ignoreBeforeTs = r.StartTs
default:
return errors.Errorf("unsupported column family type: %s", cfName)
}

for _, rule := range r.Data {
rule.IgnoreBeforeTimestamp = ignoreBeforeTs
rule.IgnoreAfterTimestamp = r.restoredTs
rule.IgnoreAfterTimestamp = r.RestoredTs
}
return nil
}
Expand Down Expand Up @@ -244,7 +250,7 @@ func GetRewriteRuleOfTable(
})
}

return &RewriteRules{Data: dataRules, NewTableID: newTableID, shiftStartTs: shiftStartTs, startTs: startTs, restoredTs: restoredTs, TableIDRemapHint: remaps}
return &RewriteRules{Data: dataRules, NewTableID: newTableID, ShiftStartTs: shiftStartTs, StartTs: startTs, RestoredTs: restoredTs, TableIDRemapHint: remaps}
}

// ValidateFileRewriteRule uses rewrite rules to validate the ranges of a file.
Expand Down
152 changes: 146 additions & 6 deletions br/pkg/restore/utils/rewrite_rule_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ package utils_test

import (
"bytes"
"strings"
"testing"

"github.com/pingcap/errors"
Expand Down Expand Up @@ -397,21 +398,52 @@ func TestGetRewriteRulesMap(t *testing.T) {
}

func TestGetRewriteRuleOfTable(t *testing.T) {
// Test basic table prefix rewrite without detailed rules
{
rewriteRules := utils.GetRewriteRuleOfTable(2, 1, 0, map[int64]int64{1: 1, 2: 2}, false)
rewriteRules := utils.GetRewriteRuleOfTable(2, 1, 0, 0, 0, map[int64]int64{1: 1, 2: 2}, false)
require.Equal(t, getNewKeyPrefix(tablecodec.EncodeTablePrefix(2), rewriteRules), tablecodec.EncodeTablePrefix(1))
require.Len(t, rewriteRules.Data, 1) // Only one rule for table prefix
require.Equal(t, rewriteRules.NewTableID, int64(1))
require.Equal(t, rewriteRules.TableIDRemapHint, []utils.TableIDRemap{{Origin: 2, Rewritten: 1}})
}

// Test detailed rules including record and index prefixes
{
rewriteRules := utils.GetRewriteRuleOfTable(2, 1, 0, map[int64]int64{1: 1, 2: 2}, true)
indexIDs := map[int64]int64{1: 1, 2: 2}
rewriteRules := utils.GetRewriteRuleOfTable(2, 1, 0, 0, 0, indexIDs, true)
// Check record prefix
require.Equal(t, getNewKeyPrefix(tablecodec.GenTableRecordPrefix(2), rewriteRules), tablecodec.GenTableRecordPrefix(1))
// Check index prefixes
require.Equal(t, getNewKeyPrefix(tablecodec.EncodeTableIndexPrefix(2, 1), rewriteRules), tablecodec.EncodeTableIndexPrefix(1, 1))
require.Equal(t, getNewKeyPrefix(tablecodec.EncodeTableIndexPrefix(2, 2), rewriteRules), tablecodec.EncodeTableIndexPrefix(1, 2))
// Verify number of rules (1 record + 2 index rules)
require.Len(t, rewriteRules.Data, 3)
}

// Test timestamp fields
{
shiftStartTs := uint64(30)
startTs := uint64(50)
restoredTs := uint64(100)
rewriteRules := utils.GetRewriteRuleOfTable(2, 1, shiftStartTs, startTs, restoredTs, map[int64]int64{1: 1}, true)

// Verify timestamp fields in RewriteRules
require.Equal(t, restoredTs, rewriteRules.RestoredTs)
require.Equal(t, startTs, rewriteRules.StartTs)
require.Equal(t, shiftStartTs, rewriteRules.ShiftStartTs)

// Verify TableIDRemapHint
require.Equal(t, []utils.TableIDRemap{{Origin: 2, Rewritten: 1}}, rewriteRules.TableIDRemapHint)

// Verify NewTableID
require.Equal(t, int64(1), rewriteRules.NewTableID)
}

// Test with empty index IDs
{
rewriteRules := utils.GetRewriteRuleOfTable(2, 1, 100, map[int64]int64{1: 1, 2: 2}, true)
require.Equal(t, rewriteRules.Data[1].IgnoreAfterTimestamp, uint64(100))
require.Equal(t, rewriteRules.Data[2].IgnoreAfterTimestamp, uint64(100))
rewriteRules := utils.GetRewriteRuleOfTable(2, 1, 0, 0, 0, map[int64]int64{}, true)
require.Len(t, rewriteRules.Data, 1) // Only record rule, no index rules
require.Equal(t, getNewKeyPrefix(tablecodec.GenTableRecordPrefix(2), rewriteRules), tablecodec.GenTableRecordPrefix(1))
}
}

Expand All @@ -436,7 +468,7 @@ func rewriteKey(key kv.Key, rule *import_sstpb.RewriteRule) kv.Key {
}

func TestFindMatchedRewriteRule(t *testing.T) {
rewriteRules := utils.GetRewriteRuleOfTable(2, 1, 0, map[int64]int64{1: 10}, true)
rewriteRules := utils.GetRewriteRuleOfTable(2, 1, 0, 0, 0, map[int64]int64{1: 10}, true)
{
applyFile := fakeApplyFile{
StartKey: tablecodec.EncodeRowKeyWithHandle(2, kv.IntHandle(100)),
Expand Down Expand Up @@ -486,3 +518,111 @@ func TestGetRewriteKeyWithDifferentTable(t *testing.T) {
_, _, err = utils.GetRewriteEncodedKeys(applyFile, nil)
require.Error(t, err)
}

func TestSetTimeRangeFilter(t *testing.T) {
testCases := []struct {
name string
rules *utils.RewriteRules
cfName string
expectError bool
}{
{
name: "default cf with valid timestamps",
rules: &utils.RewriteRules{
Data: []*import_sstpb.RewriteRule{
{OldKeyPrefix: []byte("old"), NewKeyPrefix: []byte("new")},
},
ShiftStartTs: 50, // Less than StartTs
StartTs: 100,
RestoredTs: 200,
},
cfName: "default",
expectError: false,
},
{
name: "write cf with valid timestamps",
rules: &utils.RewriteRules{
Data: []*import_sstpb.RewriteRule{
{OldKeyPrefix: []byte("old"), NewKeyPrefix: []byte("new")},
},
ShiftStartTs: 50, // Less than StartTs
StartTs: 100,
RestoredTs: 200,
},
cfName: "write",
expectError: false,
},
{
name: "invalid shift start ts (greater than start ts)",
rules: &utils.RewriteRules{
Data: []*import_sstpb.RewriteRule{
{OldKeyPrefix: []byte("old"), NewKeyPrefix: []byte("new")},
},
ShiftStartTs: 150, // Greater than StartTs
StartTs: 100,
RestoredTs: 200,
},
cfName: "default",
expectError: true,
},
{
name: "invalid cf name",
rules: &utils.RewriteRules{
Data: []*import_sstpb.RewriteRule{
{OldKeyPrefix: []byte("old"), NewKeyPrefix: []byte("new")},
},
ShiftStartTs: 50,
StartTs: 100,
RestoredTs: 200,
},
cfName: "invalid",
expectError: true,
},
{
name: "zero timestamps should skip filter",
rules: &utils.RewriteRules{
Data: []*import_sstpb.RewriteRule{
{OldKeyPrefix: []byte("old"), NewKeyPrefix: []byte("new")},
},
StartTs: 0,
RestoredTs: 0,
ShiftStartTs: 0,
},
cfName: "default",
expectError: false,
},
}

for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
err := tc.rules.SetTimeRangeFilter(tc.cfName)

if tc.expectError {
require.Error(t, err)
return
}
require.NoError(t, err)

if tc.rules.StartTs == 0 || tc.rules.RestoredTs == 0 {
// Should not modify rules when timestamps are zero
for _, rule := range tc.rules.Data {
require.Zero(t, rule.IgnoreBeforeTimestamp)
require.Zero(t, rule.IgnoreAfterTimestamp)
}
return
}

// Verify timestamps are set correctly
for _, rule := range tc.rules.Data {
require.Equal(t, tc.rules.RestoredTs, rule.IgnoreAfterTimestamp)

if strings.Contains(tc.cfName, "default") {
require.Equal(t, tc.rules.ShiftStartTs, rule.IgnoreBeforeTimestamp)
require.Less(t, tc.rules.ShiftStartTs, tc.rules.StartTs)
} else if strings.Contains(tc.cfName, "write") {
require.Equal(t, tc.rules.StartTs, rule.IgnoreBeforeTimestamp)
}
}
})
}
}

0 comments on commit f000964

Please sign in to comment.