feat: limit support in Update/Delete API

tigrisdata-archive · Oct 10, 2022 · c5e3513 · c5e3513
1 parent 3984943
commit c5e3513
Show file tree

Hide file tree

Showing 3 changed files with 102 additions and 1 deletion.
diff --git a/api/proto b/api/proto
diff --git a/server/services/v1/query_runner.go b/server/services/v1/query_runner.go
@@ -493,6 +493,10 @@ func (runner *UpdateQueryRunner) Run(ctx context.Context, tx transaction.Tx, ten
 		runner.queryMetrics.SetWriteType("non-pkey")
 	}
 
+	var limit = int32(0)
+	if runner.req.Options != nil {
+		limit = int32(runner.req.Options.Limit)
+	}
 	modifiedCount := int32(0)
 	var row Row
 	for iterator.Next(&row) {
@@ -515,6 +519,9 @@ func (runner *UpdateQueryRunner) Run(ctx context.Context, tx transaction.Tx, ten
 			return nil, ctx, err
 		}
 		modifiedCount++
+		if limit > 0 && modifiedCount == limit {
+			break
+		}
 	}
 
 	ctx = metrics.UpdateSpanTags(ctx, runner.queryMetrics)
@@ -592,6 +599,10 @@ func (runner *DeleteQueryRunner) Run(ctx context.Context, tx transaction.Tx, ten
 		return nil, ctx, err
 	}
 
+	var limit = int32(0)
+	if runner.req.Options != nil {
+		limit = int32(runner.req.Options.Limit)
+	}
 	modifiedCount := int32(0)
 	var row Row
 	for iterator.Next(&row) {
@@ -605,6 +616,9 @@ func (runner *DeleteQueryRunner) Run(ctx context.Context, tx transaction.Tx, ten
 		}
 
 		modifiedCount++
+		if limit > 0 && modifiedCount == limit {
+			break
+		}
 	}
 
 	ctx = metrics.UpdateSpanTags(ctx, runner.queryMetrics)

diff --git a/test/v1/server/document_test.go b/test/v1/server/document_test.go
@@ -1267,6 +1267,93 @@ func TestUpdate_MultipleRows(t *testing.T) {
 		outDocument)
 }
 
+func TestUpdate_Limit(t *testing.T) {
+	db, coll := setupTests(t)
+	defer cleanupTests(t, db)
+
+	inputDocument := []Doc{
+		{
+			"pkey_int":     110,
+			"int_value":    1000,
+			"string_value": "simple_insert110",
+			"bool_value":   true,
+			"double_value": 1000.000001,
+			"bytes_value":  []byte(`"simple_insert110"`),
+		},
+		{
+			"pkey_int":     120,
+			"int_value":    2000,
+			"string_value": "simple_insert120",
+			"bool_value":   false,
+			"double_value": 2000.22221,
+			"bytes_value":  []byte(`"simple_insert120"`),
+		},
+		{
+			"pkey_int":     130,
+			"int_value":    3000,
+			"string_value": "simple_insert130",
+			"bool_value":   true,
+			"double_value": 3000.999999,
+			"bytes_value":  []byte(`"simple_insert130"`),
+		},
+	}
+
+	// should always succeed with mustNotExists as false
+	insertDocuments(t, db, coll, inputDocument, false).
+		Status(http.StatusOK)
+
+	readFilter := Map{
+		"$or": []Doc{
+			{"pkey_int": 110},
+			{"pkey_int": 120},
+			{"pkey_int": 130},
+		},
+	}
+	readAndValidate(t,
+		db,
+		coll,
+		readFilter,
+		nil,
+		inputDocument)
+
+	filter := Map{
+		"filter": Map{
+			"$or": []Doc{
+				{"pkey_int": 110},
+				{"pkey_int": 120},
+				{"pkey_int": 130},
+			},
+		},
+	}
+	fields := Map{
+		"fields": Map{
+			"$set": Map{
+				"int_value": 12345,
+			},
+		},
+	}
+
+	payload := make(Map)
+	for key, value := range filter {
+		payload[key] = value
+	}
+	for key, value := range fields {
+		payload[key] = value
+	}
+	payload["options"] = Map{
+		"limit": 1,
+	}
+
+	e := expect(t)
+	e.PUT(getDocumentURL(db, coll, "update")).
+		WithJSON(payload).
+		Expect().
+		Status(http.StatusOK).
+		JSON().
+		Object().
+		ValueEqual("modified_count", 1)
+}
+
 func TestUpdate_UsingCollation(t *testing.T) {
 	db, coll := setupTests(t)
 	defer cleanupTests(t, db)