github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/ccl/storageccl/key_rewriter_test.go (about)

     1  // Copyright 2016 The Cockroach Authors.
     2  //
     3  // Licensed as a CockroachDB Enterprise file under the Cockroach Community
     4  // License (the "License"); you may not use this file except in compliance with
     5  // the License. You may obtain a copy of the License at
     6  //
     7  //     https://github.com/cockroachdb/cockroach/blob/master/licenses/CCL.txt
     8  
     9  package storageccl
    10  
    11  import (
    12  	"bytes"
    13  	"testing"
    14  
    15  	"github.com/cockroachdb/cockroach/pkg/keys"
    16  	"github.com/cockroachdb/cockroach/pkg/roachpb"
    17  	"github.com/cockroachdb/cockroach/pkg/sql/sqlbase"
    18  	"github.com/cockroachdb/cockroach/pkg/util/encoding"
    19  	"github.com/cockroachdb/cockroach/pkg/util/hlc"
    20  	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
    21  	"github.com/cockroachdb/cockroach/pkg/util/protoutil"
    22  )
    23  
    24  func TestPrefixRewriter(t *testing.T) {
    25  	defer leaktest.AfterTest(t)()
    26  
    27  	kr := prefixRewriter{
    28  		{
    29  			OldPrefix: []byte{1, 2, 3},
    30  			NewPrefix: []byte{4, 5, 6},
    31  		},
    32  		{
    33  			OldPrefix: []byte{7, 8, 9},
    34  			NewPrefix: []byte{10},
    35  		},
    36  	}
    37  
    38  	t.Run("match", func(t *testing.T) {
    39  		key := []byte{1, 2, 3, 4}
    40  		newKey, ok := kr.rewriteKey(key)
    41  		if !ok {
    42  			t.Fatalf("expected to match %s but didn't", key)
    43  		}
    44  		if expected := []byte{4, 5, 6, 4}; !bytes.Equal(newKey, expected) {
    45  			t.Fatalf("got %x expected %x", newKey, expected)
    46  		}
    47  	})
    48  
    49  	t.Run("no match", func(t *testing.T) {
    50  		key := []byte{4, 5, 6}
    51  		_, ok := kr.rewriteKey(key)
    52  		if ok {
    53  			t.Fatalf("expected to not match %s but did", key)
    54  		}
    55  	})
    56  }
    57  
    58  func TestKeyRewriter(t *testing.T) {
    59  	defer leaktest.AfterTest(t)()
    60  
    61  	desc := sqlbase.NamespaceTable
    62  	oldID := desc.ID
    63  	newID := desc.ID + 1
    64  	desc.ID = newID
    65  	rekeys := []roachpb.ImportRequest_TableRekey{
    66  		{
    67  			OldID:   uint32(oldID),
    68  			NewDesc: mustMarshalDesc(t, &desc),
    69  		},
    70  	}
    71  
    72  	const notSpan = false
    73  
    74  	kr, err := MakeKeyRewriterFromRekeys(rekeys)
    75  	if err != nil {
    76  		t.Fatal(err)
    77  	}
    78  
    79  	t.Run("normal", func(t *testing.T) {
    80  		key := sqlbase.MakeIndexKeyPrefix(keys.SystemSQLCodec, &sqlbase.NamespaceTable, desc.PrimaryIndex.ID)
    81  		newKey, ok, err := kr.RewriteKey(key, notSpan)
    82  		if err != nil {
    83  			t.Fatal(err)
    84  		}
    85  		if !ok {
    86  			t.Fatal("expected rewrite")
    87  		}
    88  		_, id, err := encoding.DecodeUvarintAscending(newKey)
    89  		if err != nil {
    90  			t.Fatalf("%+v", err)
    91  		}
    92  		if sqlbase.ID(id) != newID {
    93  			t.Fatalf("got %d expected %d", id, newID)
    94  		}
    95  	})
    96  
    97  	t.Run("prefix end", func(t *testing.T) {
    98  		key := roachpb.Key(sqlbase.MakeIndexKeyPrefix(keys.SystemSQLCodec, &sqlbase.NamespaceTable, desc.PrimaryIndex.ID)).PrefixEnd()
    99  		newKey, ok, err := kr.RewriteKey(key, notSpan)
   100  		if err != nil {
   101  			t.Fatal(err)
   102  		}
   103  		if !ok {
   104  			t.Fatalf("expected rewrite")
   105  		}
   106  		_, id, err := encoding.DecodeUvarintAscending(newKey)
   107  		if err != nil {
   108  			t.Fatalf("%+v", err)
   109  		}
   110  		if sqlbase.ID(id) != newID {
   111  			t.Fatalf("got %d expected %d", id, newID)
   112  		}
   113  	})
   114  
   115  	t.Run("multi", func(t *testing.T) {
   116  		desc.ID = oldID + 10
   117  		desc2 := sqlbase.DescriptorTable
   118  		desc2.ID += 10
   119  		newKr, err := MakeKeyRewriterFromRekeys([]roachpb.ImportRequest_TableRekey{
   120  			{OldID: uint32(oldID), NewDesc: mustMarshalDesc(t, &desc)},
   121  			{OldID: uint32(sqlbase.DescriptorTable.ID), NewDesc: mustMarshalDesc(t, &desc2)},
   122  		})
   123  		if err != nil {
   124  			t.Fatal(err)
   125  		}
   126  
   127  		key := sqlbase.MakeIndexKeyPrefix(keys.SystemSQLCodec, &sqlbase.NamespaceTable, desc.PrimaryIndex.ID)
   128  		newKey, ok, err := newKr.RewriteKey(key, notSpan)
   129  		if err != nil {
   130  			t.Fatal(err)
   131  		}
   132  		if !ok {
   133  			t.Fatalf("expected rewrite")
   134  		}
   135  		_, id, err := encoding.DecodeUvarintAscending(newKey)
   136  		if err != nil {
   137  			t.Fatalf("%+v", err)
   138  		}
   139  		if sqlbase.ID(id) != oldID+10 {
   140  			t.Fatalf("got %d expected %d", id, desc.ID+1)
   141  		}
   142  	})
   143  }
   144  
   145  func mustMarshalDesc(t *testing.T, tableDesc *sqlbase.TableDescriptor) []byte {
   146  	desc := sqlbase.WrapDescriptor(tableDesc)
   147  	// Set the timestamp to a non-zero value.
   148  	desc.Table(hlc.Timestamp{WallTime: 1})
   149  	bytes, err := protoutil.Marshal(desc)
   150  	if err != nil {
   151  		t.Fatal(err)
   152  	}
   153  	return bytes
   154  }
   155  
   156  func BenchmarkPrefixRewriter(b *testing.B) {
   157  	kr := prefixRewriter{
   158  		{
   159  			OldPrefix: []byte{1, 2, 3},
   160  			NewPrefix: []byte{4, 5, 6},
   161  		},
   162  		{
   163  			OldPrefix: []byte{7, 8, 9},
   164  			NewPrefix: []byte{10},
   165  		},
   166  	}
   167  
   168  	b.ResetTimer()
   169  	for i := 0; i < b.N; i++ {
   170  		key := []byte{1, 2, 3, 4}
   171  		_, _ = kr.rewriteKey(key)
   172  	}
   173  }