From 8929e622e602ecd4c719c51ce11b5df77b3bd3c9 Mon Sep 17 00:00:00 2001
From: Aniket Pant <aniketp@uber.com>
Date: Sat, 14 Dec 2019 12:41:13 -0500
Subject: [PATCH] Add support for nil partition key for secondary indexes in
 memory connector

Why:
Cassandra supports secondary indexes where the partition key can be
nil. This results in the index not being generated until the value is
set. The test client which uses a memory connector should support the
same behavior.

-

This change addresses the need by:
Add a `nil` check in `partitionKeyBuilder` which is capable of handling
different value types. In these cases, we simply skip over the value
instead of encoding a `nil` value which results in a panic inside the
gob encoder.

-
---
 connectors/memory/memory.go   | 25 +++++++++++++++++++++++++
 testclient/testclient_test.go | 17 +++++++++++++++++
 testentity/testentity.go      | 32 ++++++++++++++++++++------------
 3 files changed, 62 insertions(+), 12 deletions(-)

diff --git a/connectors/memory/memory.go b/connectors/memory/memory.go
index 138afca0..e66700db 100644
--- a/connectors/memory/memory.go
+++ b/connectors/memory/memory.go
@@ -93,6 +93,9 @@ func partitionKeyBuilder(pk *dosa.PrimaryKey, values map[string]dosa.FieldValue)
 	var encodedKey []byte
 	for _, k := range pk.PartitionKeys {
 		if v, ok := values[k]; ok {
+			if isNilInterface(v) {
+				continue
+			}
 			encodedVal, _ := encoder.Encode(v)
 			encodedKey = append(encodedKey, encodedVal...)
 		} else {
@@ -714,6 +717,28 @@ func getStartingPoint(ei *dosa.EntityInfo, token string) (start string, startPar
 	return start, startPartKey, nil
 }
 
+func isNilInterface(v dosa.FieldValue) bool {
+	switch v := v.(type) {
+	case *dosa.UUID:
+		return v == nil
+	case *string:
+		return v == nil
+	case *int64:
+		return v == nil
+	case *int32:
+		return v == nil
+	case *float64:
+		return v == nil
+	case *[]byte:
+		return v == nil
+	case *time.Time:
+		return v == nil
+	case *bool:
+		return v == nil
+	}
+	return false
+}
+
 // CheckSchema is just a stub; there is no schema management for the in memory connector
 // since creating a new one leaves you with no data!
 func (c *Connector) CheckSchema(ctx context.Context, scope, namePrefix string, ed []*dosa.EntityDefinition) (int32, error) {
diff --git a/testclient/testclient_test.go b/testclient/testclient_test.go
index 506eb355..e7b6a769 100644
--- a/testclient/testclient_test.go
+++ b/testclient/testclient_test.go
@@ -47,6 +47,23 @@ func TestTestClient(t *testing.T) {
 	err = client.Upsert(context.Background(), nil, &testEnt)
 	assert.NoError(t, err)
 
+	uuidv := dosa.NewUUID()
+	op := dosa.NewRangeOp(&testentity.TestEntity{}).
+		Eq("UUIDVP", dosa.UUID(uuidv)).
+		Limit(1)
+	results, _, err := client.Range(context.Background(), op)
+	assert.Empty(t, results)
+	assert.NoError(t, err)
+
+	testEnt.UUIDVP = &uuidv
+	err = client.Upsert(context.Background(), nil, &testEnt)
+	assert.NoError(t, err)
+
+	results, _, err = client.Range(context.Background(), op)
+	assert.Len(t, results, 1)
+	assert.Equal(t, &testEnt, results[0].(*testentity.TestEntity))
+	assert.NoError(t, err)
+
 	readEnt := testentity.TestEntity{
 		UUIDKey:  uuid,
 		StrKey:   "key",
diff --git a/testentity/testentity.go b/testentity/testentity.go
index b180b2ff..16e72559 100644
--- a/testentity/testentity.go
+++ b/testentity/testentity.go
@@ -28,18 +28,26 @@ import (
 
 // TestEntity uses common key types and all types in value fields.
 type TestEntity struct {
-	dosa.Entity `dosa:"name=awesome_test_entity, primaryKey=(UUIDKey, StrKey ASC, Int64Key DESC)"`
-	UUIDKey     dosa.UUID `dosa:"name=an_uuid_key"`
-	StrKey      string
-	Int64Key    int64
-	UUIDV       dosa.UUID
-	StrV        string
-	Int64V      int64 `dosa:"name=an_int64_value"`
-	Int32V      int32
-	DoubleV     float64
-	BoolV       bool
-	BlobV       []byte
-	TSV         time.Time
+	dosa.Entity      `dosa:"name=awesome_test_entity, primaryKey=(UUIDKey, StrKey ASC, Int64Key DESC)"`
+	EntityByUUIDVP   dosa.Index `dosa:"key=(UUIDVP)"`
+	EntityByStrVP    dosa.Index `dosa:"key=(StrVP)"`
+	EntityByInt64VP  dosa.Index `dosa:"key=(Int64VP)"`
+	EntityByInt32VP  dosa.Index `dosa:"key=(Int32VP)"`
+	EntityByDoubleVP dosa.Index `dosa:"key=(DoubleVP)"`
+	EntityByBoolVP   dosa.Index `dosa:"key=(BoolVP)"`
+	EntityByTSVP     dosa.Index `dosa:"key=(TSVP)"`
+
+	UUIDKey  dosa.UUID `dosa:"name=an_uuid_key"`
+	StrKey   string
+	Int64Key int64
+	UUIDV    dosa.UUID
+	StrV     string
+	Int64V   int64 `dosa:"name=an_int64_value"`
+	Int32V   int32
+	DoubleV  float64
+	BoolV    bool
+	BlobV    []byte
+	TSV      time.Time
 
 	UUIDVP   *dosa.UUID
 	StrVP    *string