From c73d043b10ccb3d11319953bbbfb332ad54bc441 Mon Sep 17 00:00:00 2001
From: "Jeffrey J. Persch" <jjpersch@monetate.com>
Date: Sat, 23 Sep 2023 20:28:54 -0400
Subject: [PATCH] Add monetate_crc32 multi-column vindex. (#1)

Compute multi-column vindex compatible with existing application layer sharding.
Supports integer, varchar, and varbinary columns.

Signed-off-by: Jeffrey J. Persch <jjpersch@monetate.com>
---
 go/vt/vtgate/vindexes/monetate_crc32.go      | 96 ++++++++++++++++++++
 go/vt/vtgate/vindexes/monetate_crc32_test.go | 41 +++++++++
 2 files changed, 137 insertions(+)
 create mode 100644 go/vt/vtgate/vindexes/monetate_crc32.go
 create mode 100644 go/vt/vtgate/vindexes/monetate_crc32_test.go

diff --git a/go/vt/vtgate/vindexes/monetate_crc32.go b/go/vt/vtgate/vindexes/monetate_crc32.go
new file mode 100644
index 00000000000..3bdec23a584
--- /dev/null
+++ b/go/vt/vtgate/vindexes/monetate_crc32.go
@@ -0,0 +1,96 @@
+/*
+Copyright 2023 Monetate, Inc.
+*/
+
+package vindexes
+
+import (
+	"bytes"
+	"context"
+	"encoding/binary"
+	"hash/crc32"
+	"strings"
+
+	"vitess.io/vitess/go/sqltypes"
+	"vitess.io/vitess/go/vt/key"
+	vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
+	"vitess.io/vitess/go/vt/vterrors"
+)
+
+var _ MultiColumn = (*MonetateCRC32)(nil)
+
+type MonetateCRC32 struct {
+	name string
+}
+
+// NewMonetateCRC32 creates a new MonetateCRC32.
+func NewMonetateCRC32(name string, m map[string]string) (Vindex, error) {
+	return &MonetateCRC32{name: name}, nil
+}
+
+func (m *MonetateCRC32) String() string {
+	return m.name
+}
+
+func (m *MonetateCRC32) Cost() int {
+	return 1
+}
+
+func (m *MonetateCRC32) IsUnique() bool {
+	return true
+}
+
+func (m *MonetateCRC32) NeedsVCursor() bool {
+	return false
+}
+
+func (m *MonetateCRC32) Map(ctx context.Context, vcursor VCursor, rowsColValues [][]sqltypes.Value) ([]key.Destination, error) {
+	out := make([]key.Destination, 0, len(rowsColValues))
+	for _, colValues := range rowsColValues {
+		ksid, err := ChecksumValues(colValues)
+		if err != nil {
+			out = append(out, key.DestinationNone{})
+			continue
+		}
+		out = append(out, key.DestinationKeyspaceID(ksid))
+	}
+	return out, nil
+}
+
+func (m *MonetateCRC32) Verify(ctx context.Context, vcursor VCursor, rowsColValues [][]sqltypes.Value, ksids [][]byte) ([]bool, error) {
+	out := make([]bool, 0, len(rowsColValues))
+	for idx, colValues := range rowsColValues {
+		ksid, err := ChecksumValues(colValues)
+		if err != nil {
+			return nil, err
+		}
+		out = append(out, bytes.Equal(ksid, ksids[idx]))
+	}
+	return out, nil
+}
+
+func (m *MonetateCRC32) PartialVindex() bool {
+	return false
+}
+
+func ChecksumValues(colValues []sqltypes.Value) ([]byte, error) {
+	// concat string values of columns, separated by slashes
+	var parts []string
+	for _, colVal := range colValues {
+		if !(colVal.IsIntegral() || colVal.IsText() || colVal.IsBinary()) {
+			return nil, vterrors.New(vtrpcpb.Code_INVALID_ARGUMENT, "invalid monetate vindex value")
+		}
+		parts = append(parts, colVal.ToString())
+	}
+	var shardKey = strings.Join(parts, "/")
+	var checksum = crc32.ChecksumIEEE([]byte(shardKey))
+	var vshard = checksum % 1048576 // 20 bit vshard id
+
+	var hashed [4]byte
+	binary.BigEndian.PutUint32(hashed[:], vshard<<12)
+	return hashed[:], nil
+}
+
+func init() {
+	Register("monetate_crc32", NewMonetateCRC32)
+}
diff --git a/go/vt/vtgate/vindexes/monetate_crc32_test.go b/go/vt/vtgate/vindexes/monetate_crc32_test.go
new file mode 100644
index 00000000000..8b3f22d51e6
--- /dev/null
+++ b/go/vt/vtgate/vindexes/monetate_crc32_test.go
@@ -0,0 +1,41 @@
+/*
+Copyright 2023 Monetate, Inc.
+*/
+
+package vindexes
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"vitess.io/vitess/go/sqltypes"
+	"vitess.io/vitess/go/vt/key"
+)
+
+func TestMonetateCRC32Map(t *testing.T) {
+	vindex, err := CreateVindex("monetate_crc32", "monetate", map[string]string{})
+	require.NoError(t, err)
+	mutiCol := vindex.(MultiColumn)
+
+	got, err := mutiCol.Map(context.Background(), nil, [][]sqltypes.Value{{
+		// visitor shard key
+		sqltypes.NewInt64(1), sqltypes.NewInt64(2), sqltypes.NewInt64(3),
+	}, {
+		// customer shard key
+		sqltypes.NewInt64(1), sqltypes.NewVarBinary("customer_id"),
+	}, {
+		// customer shard key, invalid customer id as null
+		sqltypes.NewInt64(1), sqltypes.NULL,
+	}})
+	assert.NoError(t, err)
+
+	want := []key.Destination{
+		key.DestinationKeyspaceID("\x1b\x1c\x60\x00"),
+		key.DestinationKeyspaceID("\xc4\x4b\xe0\x00"),
+		key.DestinationNone{},
+	}
+	assert.Equal(t, want, got)
+}