TT-10949 Adding SetIfNotExist to KeyValue (#108)

* Adding SetIfNotExist to KeyValue

* linting

persistent/model/object_id_test.go

@@ -92,7 +92,6 @@ func TestUnmarshalJSON(t *testing.T) {
 
 	var id2 ObjectID
 	err = id2.UnmarshalJSON(idBytes)
-
 	if err != nil {
 		t.Fatal(err)
 	}

temporal/internal/driver/redisv9/keyvalue.go

@@ -195,7 +195,6 @@ func (r *RedisV9) DeleteScanMatch(ctx context.Context, pattern string) (int64, e
 	case *redis.Client:
 		var err error
 		totalDeleted, err = r.deleteScanMatchSingleNode(ctx, client, pattern)
-
 		if err != nil {
 			if errors.Is(err, redis.ErrClosed) {
 				return totalDeleted, temperr.ClosedConnection
@@ -218,7 +217,6 @@ func (r *RedisV9) deleteScanMatchSingleNode(ctx context.Context, client redis.Cm
 
 	var keys []string
 	keys, _, err = client.Scan(ctx, cursor, pattern, 0).Result()
-
 	if err != nil {
 		return int64(deleted), err
 	}
@@ -479,3 +477,16 @@ func fetchKeys(ctx context.Context,
 
 	return keys, cursor, nil
 }
+
+func (r *RedisV9) SetIfNotExist(ctx context.Context, key, value string, expiration time.Duration) (bool, error) {
+	if key == "" {
+		return false, temperr.KeyEmpty
+	}
+
+	res := r.client.SetNX(ctx, key, value, expiration)
+	if res.Err() != nil {
+		return false, res.Err()
+	}
+
+	return res.Val(), nil
+}

temporal/keyvalue/keyvalue_test.go

@@ -1261,3 +1261,96 @@ func TestKeyValue_GetKeysWithOpts(t *testing.T) {
 		}
 	}
 }
+
+func TestKeyValue_SetIfNotExist(t *testing.T) {
+	connectors := testutil.TestConnectors(t)
+	defer testutil.CloseConnectors(t, connectors)
+
+	tcs := []struct {
+		name        string
+		key         string
+		value       string
+		pre         func(kv model.KeyValue)
+		expiration  time.Duration
+		expectedErr error
+		expectedSet bool
+	}{
+		{
+			name:        "set_with_valid_key_and_value",
+			key:         "key1",
+			value:       "value1",
+			expiration:  10 * time.Second,
+			expectedErr: nil,
+			expectedSet: true,
+		},
+		{
+			name:        "set_with_empty_key",
+			key:         "",
+			value:       "value2",
+			expiration:  10 * time.Second,
+			expectedErr: temperr.KeyEmpty,
+		},
+		{
+			name:        "set_with_empty_value",
+			key:         "key3",
+			value:       "",
+			expiration:  10 * time.Second,
+			expectedErr: nil,
+			expectedSet: true,
+		},
+		{
+			name:        "set_with_no_expiration",
+			key:         "key4",
+			value:       "value4",
+			expiration:  10 * time.Second,
+			expectedErr: nil,
+			expectedSet: true,
+		},
+		{
+			name:        "set_already_existing_key",
+			key:         "key5",
+			value:       "value5",
+			expiration:  10 * time.Second,
+			expectedErr: nil,
+			pre: func(kv model.KeyValue) {
+				t.Helper()
+				err := kv.Set(context.Background(), "key5", "value5", 10*time.Second)
+				assert.Nil(t, err)
+			},
+			expectedSet: false,
+		},
+	}
+
+	for _, connector := range connectors {
+		for _, tc := range tcs {
+			t.Run(connector.Type()+"_"+tc.name, func(t *testing.T) {
+				ctx := context.Background()
+
+				kv, err := NewKeyValue(connector)
+				assert.Nil(t, err)
+
+				flusher, err := flusher.NewFlusher(connector)
+				assert.Nil(t, err)
+				defer assert.Nil(t, flusher.FlushAll(ctx))
+
+				if tc.pre != nil {
+					tc.pre(kv)
+				}
+
+				val, err := kv.SetIfNotExist(ctx, tc.key, tc.value, tc.expiration)
+				assert.Equal(t, tc.expectedErr, err)
+				assert.Equal(t, tc.expectedSet, val)
+				if err == nil {
+					actualValue, err := kv.Get(ctx, tc.key)
+					assert.Nil(t, err)
+
+					assert.Equal(t, tc.value, actualValue)
+
+					actualTTL, err := kv.TTL(ctx, tc.key)
+					assert.Nil(t, err)
+					assert.True(t, actualTTL <= int64(tc.expiration.Seconds()))
+				}
+			})
+		}
+	}
+}

temporal/model/types.go

@@ -55,6 +55,9 @@ type KeyValue interface {
 	Get(ctx context.Context, key string) (value string, err error)
 	// Set sets the string value of a key
 	Set(ctx context.Context, key, value string, ttl time.Duration) error
+	// SetIfNotExist sets the string value of a key if the key does not exist.
+	// Returns true if the key was set, false otherwise.
+	SetIfNotExist(ctx context.Context, key, value string, expiration time.Duration) (bool, error)
 	// Delete removes the specified keys
 	Delete(ctx context.Context, key string) error
 	// Increment atomically increments the integer value of a key by one